feat: allow to add attribute to specific soa type

Author: curryyzheng

soa-derive-internal/src/input.rs

@@ -1,6 +1,10 @@
+use std::convert::TryInto;
+
 use proc_macro2::{Span, TokenStream};
-use syn::{Data, DeriveInput, Ident, Field, Visibility, Meta, MetaNameValue, Lit};
 use quote::quote;
+use syn::{
+    Data, DeriveInput, Field, Ident, Lit, Meta, MetaList, MetaNameValue, NestedMeta, Visibility,
+};
 
 /// Representing the struct we are deriving
 pub struct Input {
@@ -11,29 +15,73 @@ pub struct Input {
     /// The list of fields in the struct
     pub fields: Vec<Field>,
     /// The struct overall visibility
-    pub visibility: Visibility
+    pub visibility: Visibility,
+
+    pub vec_attrs: Vec<Meta>,
+    pub slice_attrs: Vec<Meta>,
+    pub ref_attrs: Vec<Meta>,
+    pub ptr_attrs: Vec<Meta>,
 }
 
 impl Input {
     pub fn new(input: DeriveInput) -> Input {
         let fields = match input.data {
-            Data::Struct(s) => {
-                s.fields.iter().cloned().collect::<Vec<_>>()
-            }
+            Data::Struct(s) => s.fields.iter().cloned().collect::<Vec<_>>(),
             _ => panic!("#[derive(StructOfArray)] only supports structs."),
         };
 
         let mut derives: Vec<Ident> = vec![];
+        let mut vec_attrs = Vec::new();
+        let mut slice_attrs = Vec::new();
+        let mut ref_attrs = Vec::new();
+        let mut ptr_attrs = Vec::new();
         for attr in input.attrs {
             if let Ok(meta) = attr.parse_meta() {
                 if meta.path().is_ident("soa_derive") {
                     match meta {
-                        Meta::NameValue(MetaNameValue{lit: Lit::Str(string), ..}) => {
+                        Meta::NameValue(MetaNameValue {
+                            lit: Lit::Str(string),
+                            ..
+                        }) => {
                             for value in string.value().split(',') {
                                 derives.push(Ident::new(value.trim(), Span::call_site()));
                             }
                         }
-                        _ => panic!("expected #[soa_derive = \"Traits, To, Derive\"], got #[{}]", quote!(#meta))
+                        _ => panic!(
+                            "expected #[soa_derive = \"Traits, To, Derive\"], got #[{}]",
+                            quote!(#meta)
+                        ),
+                    }
+                } else if meta.path().is_ident("soa_attr") {
+                    match meta.clone() {
+                        Meta::List(MetaList { nested, .. }) => {
+                            let [soa_type, attr]: [NestedMeta; 2] =
+                                nested.into_iter().collect::<Vec<_>>().try_into().unwrap_or_else(|_| panic!("expected #[soa_derive(\"Types, To, Add, Attribute\", \"Attributes\", got #[{}])]", quote!(#meta)));
+                            let attr = match attr {
+                                NestedMeta::Meta(meta) => meta,
+                                NestedMeta::Lit(_) => {
+                                    panic!("expected a attribute, got {}", quote!(attr))
+                                }
+                            };
+                            match soa_type {
+                                NestedMeta::Meta(Meta::Path(path)) => match path.get_ident() {
+                                    Some(ident) => match ident.to_string().as_str() {
+                                        "Vec" => vec_attrs.push(attr),
+                                        "Slice" => slice_attrs.push(attr),
+                                        "Ref" => ref_attrs.push(attr),
+                                        "ptr" => ptr_attrs.push(attr),
+                                        _ => panic!("expected a soa type, got {}", ident),
+                                    },
+                                    None => {
+                                        panic!("expected a soa type, got {}", quote!(#path))
+                                    }
+                                },
+                                _ => {
+                                    panic!("expected a soa type, got {}", quote!(#soa_type))
+                                }
+                            }
+                        }
+                        _ => panic!("expected #[soa_attr(...), got #[{}]]", quote!(#meta)),
                     }
                 }
             }
@@ -43,7 +91,11 @@ impl Input {
             name: input.ident,
             derives: derives,
             fields: fields,
-            visibility: input.vis
+            visibility: input.vis,
+            vec_attrs,
+            slice_attrs,
+            ref_attrs,
+            ptr_attrs,
         }
     }
 
@@ -64,12 +116,14 @@ impl Input {
         if self.derives.is_empty() {
             TokenStream::new()
         } else {
-            let derives = &self.derives.iter()
-                                       .cloned()
-                                       .filter(|name| name != "Clone")
-                                       .filter(|name| name != "Deserialize")
-                                       .filter(|name| name != "Serialize")
-                                       .collect::<Vec<_>>();
+            let derives = &self
+                .derives
+                .iter()
+                .cloned()
+                .filter(|name| name != "Clone")
+                .filter(|name| name != "Deserialize")
+                .filter(|name| name != "Serialize")
+                .collect::<Vec<_>>();
             quote!(
                 #[derive(
                     #(#derives,)*

soa-derive-internal/src/lib.rs

@@ -17,7 +17,7 @@ mod refs;
 mod slice;
 mod vec;
 
-#[proc_macro_derive(StructOfArray, attributes(soa_derive))]
+#[proc_macro_derive(StructOfArray, attributes(soa_derive, soa_attr))]
 pub fn soa_derive(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let ast = syn::parse(input).unwrap();
     let input = input::Input::new(ast);

soa-derive-internal/src/ptr.rs

@@ -7,6 +7,7 @@ pub fn derive(input: &Input) -> TokenStream {
     let name = &input.name;
     let visibility = &input.visibility;
     let other_derive = &input.derive_with_exceptions();
+    let attrs = &input.ptr_attrs;
     let vec_name = &input.vec_name();
     let ptr_name = &input.ptr_name();
     let ptr_mut_name = &input.ptr_mut_name();
@@ -42,6 +43,7 @@ pub fn derive(input: &Input) -> TokenStream {
         #[doc = #doc_url]
         /// with struct of array layout.
         #other_derive
+        #(#[#attrs])*
         #[derive(Copy, Clone)]
         #visibility struct #ptr_name {
             #(

soa-derive-internal/src/refs.rs

@@ -7,6 +7,7 @@ pub fn derive(input: &Input) -> TokenStream {
     let name = &input.name;
     let visibility = &input.visibility;
     let other_derive = &input.derive_with_exceptions();
+    let attrs = &input.ref_attrs;
     let vec_name = &input.vec_name();
     let ref_name = &input.ref_name();
     let ref_mut_name = &input.ref_mut_name();
@@ -38,6 +39,7 @@ pub fn derive(input: &Input) -> TokenStream {
         #[doc = #doc_url]
         /// with struct of array layout.
         #other_derive
+        #(#[#attrs])*
         #[derive(Copy, Clone)]
         #visibility struct #ref_name<'a> {
             #(
@@ -50,6 +52,7 @@ pub fn derive(input: &Input) -> TokenStream {
         #[doc = #doc_url]
         /// with struct of array layout.
         #other_derive
+        #(#[#attrs])*
         #visibility struct #ref_mut_name<'a> {
             #(
                 #[doc = #fields_mut_doc]

soa-derive-internal/src/slice.rs

@@ -9,6 +9,7 @@ pub fn derive(input: &Input) -> TokenStream {
     let other_derive = &input.derive_with_exceptions();
     let visibility = &input.visibility;
     let slice_name = &input.slice_name();
+    let attrs = &input.slice_attrs;
     let vec_name = &input.vec_name();
     let ref_name = &input.ref_name();
     let ptr_name = &input.ptr_name();
@@ -48,6 +49,7 @@ pub fn derive(input: &Input) -> TokenStream {
         #[allow(dead_code)]
         #[derive(Copy, Clone)]
         #other_derive
+        #(#[#attrs])*
         #visibility struct #slice_name<'a> {
             #(
                 #[doc = #fields_doc]
@@ -239,6 +241,7 @@ 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 attrs = &input.slice_attrs;
     let ref_mut_name = &input.ref_mut_name();
     let ptr_name = &input.ptr_name();
     let ptr_mut_name = &input.ptr_mut_name();
@@ -280,6 +283,7 @@ pub fn derive_mut(input: &Input) -> TokenStream {
         /// .
         #[allow(dead_code)]
         #other_derive
+        #(#[#attrs])*
         #visibility struct #slice_mut_name<'a> {
             #(
                 #[doc = #fields_doc]

soa-derive-internal/src/vec.rs

@@ -9,6 +9,7 @@ pub fn derive(input: &Input) -> TokenStream {
     let name = &input.name;
     let vec_name_str = format!("Vec<{}>", name);
     let other_derive = &input.derive();
+    let attrs = &input.vec_attrs;
     let visibility = &input.visibility;
     let vec_name = &input.vec_name();
     let slice_name = &input.slice_name();
@@ -42,6 +43,7 @@ pub fn derive(input: &Input) -> TokenStream {
         /// ` with Struct of Array (SoA) layout
         #[allow(dead_code)]
         #other_derive
+        #(#[#attrs])*
         #visibility struct #vec_name {
             #(
                 #[doc = #fields_doc]

tests/soa_attr.rs

@@ -0,0 +1,27 @@
+use soa_derive::StructOfArray;
+
+#[derive(Debug, Clone, PartialEq, StructOfArray)]
+#[soa_attr(Vec, cfg_attr(test, derive(PartialEq, Debug)))]
+pub struct Particle {
+    pub name: String,
+    pub mass: f64,
+}
+
+impl Particle {
+    pub fn new(name: String, mass: f64) -> Self {
+        Particle { name, mass }
+    }
+}
+
+#[test]
+fn eq_test() {
+    let particles0 = ParticleVec {
+        name: vec![String::from("foo"), String::from("bar")],
+        mass: vec![1.0, 2.0],
+    };
+    let particles1 = ParticleVec {
+        name: vec![String::from("foo"), String::from("bar")],
+        mass: vec![1.0, 2.0],
+    };
+    assert_eq!(particles0, particles1);
+}