Add access methods to the generated SoA vectors

soa-derive-internal/src/lib.rs

@@ -9,6 +9,7 @@ extern crate proc_macro;
 use proc_macro2::TokenStream;
 use quote::TokenStreamExt;
 
+mod index;
 mod input;
 mod iter;
 mod ptr;
@@ -27,6 +28,7 @@ pub fn soa_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     generated.append_all(ptr::derive(&input));
     generated.append_all(slice::derive(&input));
     generated.append_all(slice::derive_mut(&input));
+    generated.append_all(index::derive(&input));
     generated.append_all(iter::derive(&input));
     generated.append_all(derive_trait(&input));
     generated.into()

soa-derive-internal/src/slice.rs

@@ -153,22 +153,47 @@ pub fn derive(input: &Input) -> TokenStream {
             /// Similar to [`
             #[doc = #slice_name_str]
             /// ::get()`](https://doc.rust-lang.org/std/primitive.slice.html#method.get).
-            pub fn get(&self, i: usize) -> Option<#ref_name> {
-                if self.is_empty() || i >= self.len() {
-                    None
-                } else {
-                    Some(#ref_name {
-                        #(#fields_names_1: self.#fields_names_2.get(i).unwrap(),)*
-                    })
-                }
+            pub fn get<'b, I>(&'b self, index: I) -> Option<I::RefOutput>
+            where
+                I: ::soa_derive::SoAIndex<#slice_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_name<'b> = self.reborrow();
+                index.get(slice)
             }
 
             /// Similar to [`
             #[doc = #slice_name_str]
             /// ::get_unchecked()`](https://doc.rust-lang.org/std/primitive.slice.html#method.get_unchecked).
-            pub unsafe fn get_unchecked(&self, i: usize) -> #ref_name {
-                #ref_name {
-                    #(#fields_names_1: self.#fields_names_2.get_unchecked(i),)*
+            pub unsafe fn get_unchecked<'b, I>(&'b self, index: I) -> I::RefOutput
+            where
+                I: ::soa_derive::SoAIndex<#slice_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_name<'b> = self.reborrow();
+                index.get_unchecked(slice)
+            }
+
+            /// Similar to [`std::ops::Index` trait](https://doc.rust-lang.org/std/ops/trait.Index.html) on
+            #[doc = #slice_name_str]
+            /// .
+            /// This is required because we cannot implement that trait.
+            pub fn index<'b, I>(&'b self, index: I) -> I::RefOutput
+            where
+                I: ::soa_derive::SoAIndex<#slice_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_name<'b> = self.reborrow();
+                index.index(slice)
+            }
+
+            /// Reborrows the slices in a more narrower lifetime
+            pub fn reborrow<'b>(&'b self) -> #slice_name<'b>
+            where
+                'a: 'b
+            {
+                #slice_name {
+                    #(#fields_names_1: &self.#fields_names_2,)*
                 }
             }
 
@@ -215,7 +240,6 @@ pub fn derive_mut(input: &Input) -> TokenStream {
     let slice_name = &input.slice_name();
     let slice_mut_name = &input.slice_mut_name();
     let vec_name = &input.vec_name();
-    let ref_name = &input.ref_name();
     let ref_mut_name = &input.ref_mut_name();
     let ptr_name = &input.ptr_name();
     let ptr_mut_name = &input.ptr_mut_name();
@@ -382,44 +406,95 @@ pub fn derive_mut(input: &Input) -> TokenStream {
             /// Similar to [`
             #[doc = #slice_name_str]
             /// ::get()`](https://doc.rust-lang.org/std/primitive.slice.html#method.get).
-            pub fn get(&self, i: usize) -> Option<#ref_name> {
-                if self.is_empty() || i >= self.len() {
-                    None
-                } else {
-                    Some(#ref_name {
-                        #(#fields_names_1: self.#fields_names_2.get(i).unwrap(),)*
-                    })
-                }
+            pub fn get<'b, I>(&'b self, index: I) -> Option<I::RefOutput>
+            where
+                I: ::soa_derive::SoAIndex<#slice_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_name<'b> = self.as_slice();
+                index.get(slice)
             }
 
             /// Similar to [`
             #[doc = #slice_name_str]
             /// ::get_unchecked()`](https://doc.rust-lang.org/std/primitive.slice.html#method.get_unchecked).
-            pub unsafe fn get_unchecked(&self, i: usize) -> #ref_name {
-                #ref_name {
-                    #(#fields_names_1: self.#fields_names_2.get_unchecked(i),)*
-                }
+            pub unsafe fn get_unchecked<'b, I>(&'b self, index: I) -> I::RefOutput
+            where
+                I: ::soa_derive::SoAIndex<#slice_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_name<'b> = self.as_slice();
+                index.get_unchecked(slice)
+            }
+
+
+            /// Similar to [`std::ops::Index` trait](https://doc.rust-lang.org/std/ops/trait.Index.html) on
+            #[doc = #slice_name_str]
+            /// .
+            /// This is required because we cannot implement that trait.
+            pub fn index<'b, I>(&'b self, index: I) -> I::RefOutput
+            where
+                I: ::soa_derive::SoAIndex<#slice_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_name<'b> = self.as_slice();
+                index.index(slice)
             }
 
             /// Similar to [`
             #[doc = #slice_name_str]
             /// ::get_mut()`](https://doc.rust-lang.org/std/primitive.slice.html#method.get_mut).
-            pub fn get_mut(&mut self, i: usize) -> Option<#ref_mut_name> {
-                if self.is_empty() || i >= self.len() {
-                    None
-                } else {
-                    Some(#ref_mut_name {
-                        #(#fields_names_1: self.#fields_names_2.get_mut(i).unwrap(),)*
-                    })
-                }
+            pub fn get_mut<'b, I>(&'b mut self, index: I) -> Option<I::MutOutput>
+            where
+                I: ::soa_derive::SoAIndexMut<#slice_mut_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_mut_name<'b> = self.reborrow();
+                index.get_mut(slice)
             }
 
             /// Similar to [`
             #[doc = #slice_name_str]
             /// ::get_unchecked_mut()`](https://doc.rust-lang.org/std/primitive.slice.html#method.get_unchecked_mut).
-            pub unsafe fn get_unchecked_mut(&mut self, i: usize) -> #ref_mut_name {
-                #ref_mut_name {
-                    #(#fields_names_1: self.#fields_names_2.get_unchecked_mut(i),)*
+            pub unsafe fn get_unchecked_mut<'b, I>(&'b mut self, index: I) -> I::MutOutput
+            where
+                I: ::soa_derive::SoAIndexMut<#slice_mut_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_mut_name<'b> = self.reborrow();
+                index.get_unchecked_mut(slice)
+            }
+
+            /// Similar to [`std::ops::IndexMut` trait](https://doc.rust-lang.org/std/ops/trait.IndexMut.html) on
+            #[doc = #slice_name_str]
+            /// .
+            /// This is required because we cannot implement that trait.
+            pub fn index_mut<'b, I>(&'b mut self, index: I) -> I::MutOutput
+            where
+                I: ::soa_derive::SoAIndexMut<#slice_mut_name<'b>>,
+                'a: 'b
+            {
+                let slice: #slice_mut_name<'b> = self.reborrow();
+                index.index_mut(slice)
+            }
+
+            /// Returns a non-mutable slice from this mutable slice.
+            pub fn as_slice<'b>(&'b self) -> #slice_name<'b>
+            where
+                'a: 'b
+            {
+                #slice_name {
+                    #(#fields_names_1: &self.#fields_names_2,)*
+                }
+            }
+
+            /// Reborrows the slices in a more narrower lifetime
+            pub fn reborrow<'b>(&'b mut self) -> #slice_mut_name<'b>
+            where
+                'a: 'b
+            {
+                #slice_mut_name {
+                    #(#fields_names_1: &mut *self.#fields_names_2,)*
                 }
             }
 

soa-derive-internal/src/vec.rs

@@ -259,6 +259,66 @@ pub fn derive(input: &Input) -> TokenStream {
                 }
             }
 
+            /// Similar to [`
+            #[doc = #vec_name_str]
+            /// ::get<I>()`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.get).
+            pub fn get<'a, I>(&'a self, index: I) -> Option<I::RefOutput>
+            where
+                I: ::soa_derive::SoAIndex<&'a #vec_name>
+            {
+                index.get(self)
+            }
+
+            /// Similar to [`
+            #[doc = #vec_name_str]
+            /// ::get_unchecked<I>()`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.get_unchecked).
+            pub unsafe fn get_unchecked<'a, I>(&'a self, index: I) -> I::RefOutput
+            where
+                I: ::soa_derive::SoAIndex<&'a #vec_name>
+            {
+                index.get_unchecked(self)
+            }
+
+            /// Similar to [`
+            #[doc = #vec_name_str]
+            /// ::index<I>()`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.index).
+            pub fn index<'a, I>(&'a self, index: I) -> I::RefOutput
+            where
+                I: ::soa_derive::SoAIndex<&'a #vec_name>
+            {
+                index.index(self)
+            }
+
+            /// Similar to [`
+            #[doc = #vec_name_str]
+            /// ::get_mut<I>()`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.get_mut).
+            pub fn get_mut<'a, I>(&'a mut self, index: I) -> Option<I::MutOutput>
+            where
+                I: ::soa_derive::SoAIndexMut<&'a mut #vec_name>
+            {
+                index.get_mut(self)
+            }
+
+            /// Similar to [`
+            #[doc = #vec_name_str]
+            /// ::get_unchecked_mut<I>()`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.get_unchecked_mut).
+            pub unsafe fn get_unchecked_mut<'a, I>(&'a mut self, index: I) -> I::MutOutput
+            where
+                I: ::soa_derive::SoAIndexMut<&'a mut #vec_name>
+            {
+                index.get_unchecked_mut(self)
+            }
+
+            /// Similar to [`
+            #[doc = #vec_name_str]
+            /// ::index_mut<I>()`](https://doc.rust-lang.org/std/vec/struct.Vec.html#method.index_mut).
+            pub fn index_mut<'a, I>(&'a mut self, index: I) -> I::MutOutput
+            where
+                I: ::soa_derive::SoAIndexMut<&'a mut #vec_name>
+            {
+                index.index_mut(self)
+            }
+
             /// Similar to [`
             #[doc = #vec_name_str]
             /// ::as_ptr()`](https://doc.rust-lang.org/std/struct.Vec.html#method.as_ptr).

src/lib.rs

@@ -168,6 +168,53 @@ pub trait StructOfArray {
     type Type;
 }
 
+
+mod private_soa_indexs {
+    // From [`std::slice::SliceIndex`](https://doc.rust-lang.org/std/slice/trait.SliceIndex.html) code.
+    // Limits the types that may implement the SoA index traits.
+    // It's also helpful to have the exaustive list of all accepted types.
+
+    use ::std::ops;
+
+    pub trait Sealed {}
+
+    impl Sealed for usize {}                        // [a]
+    impl Sealed for ops::Range<usize> {}            // [a..b]
+    impl Sealed for ops::RangeTo<usize> {}          // [..b]
+    impl Sealed for ops::RangeFrom<usize> {}        // [a..]
+    impl Sealed for ops::RangeFull {}               // [..]
+    impl Sealed for ops::RangeInclusive<usize> {}   // [a..=b]
+    impl Sealed for ops::RangeToInclusive<usize> {} // [..=b]
+}
+
+/// Helper trait used for indexing operations.
+/// Inspired by [`std::slice::SliceIndex`](https://doc.rust-lang.org/std/slice/trait.SliceIndex.html).
+pub trait SoAIndex<T>: private_soa_indexs::Sealed {
+    /// The output for the non-mutable functions
+    type RefOutput;
+
+    /// Returns the reference output in this location if in bounds. None otherwise.
+    fn get(self, soa: T) -> Option<Self::RefOutput>;
+    /// Returns the reference output in this location withotu performing any bounds check.
+    unsafe fn get_unchecked(self, soa: T) -> Self::RefOutput;
+    /// Returns the reference output in this location. Panics if it is not in bounds.
+    fn index(self, soa: T) -> Self::RefOutput;
+}
+
+/// Helper trait used for indexing operations returning mutable.
+/// Inspired by [`std::slice::SliceIndex`](https://doc.rust-lang.org/std/slice/trait.SliceIndex.html).
+pub trait SoAIndexMut<T>: private_soa_indexs::Sealed {
+    /// The output for the mutable functions
+    type MutOutput;
+
+    /// Returns the mutable reference output in this location if in bounds. None otherwise.
+    fn get_mut(self, soa: T) -> Option<Self::MutOutput>;
+    /// Returns the mutable reference output in this location withotu performing any bounds check.
+    unsafe fn get_unchecked_mut(self, soa: T) -> Self::MutOutput;
+    /// Returns the mutable reference output in this location. Panics if it is not in bounds.
+    fn index_mut(self, soa: T) -> Self::MutOutput;
+}
+
 /// Create an iterator over multiple fields in a Struct of array style vector.
 ///
 /// This macro takes two main arguments: the array/slice container, and a list