feat: Add crate argument to derive(ShrinkToFit) (#62)

crates/shrink-to-fit-macro/Cargo.toml

@@ -4,7 +4,7 @@ edition     = { workspace = true }
 license     = { workspace = true }
 name        = "shrink-to-fit-macro"
 repository  = { workspace = true }
-version     = "0.2.2"
+version     = "0.2.3"
 
 [lib]
 proc-macro = true

crates/shrink-to-fit-macro/src/lib.rs

@@ -1,16 +1,17 @@
-use quote::quote;
-use syn::{spanned::Spanned, Ident};
+use quote::{quote, ToTokens};
+use syn::{spanned::Spanned, Attribute, Expr, Ident, Lit, Meta};
 
-#[proc_macro_derive(ShrinkToFit)]
+#[proc_macro_derive(ShrinkToFit, attributes(shrink_to_fit))]
 pub fn derive_shrink_to_fit(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
     let input: syn::DeriveInput = syn::parse_macro_input!(input as syn::DeriveInput);
+    let type_attr: TypeAttr = TypeAttr::parse(&input.attrs);
 
     let name = &input.ident;
     let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
 
     let body_impl = match &input.data {
         syn::Data::Struct(s) => {
-            let (field_bindings, body_code) = expand_fields(&s.fields);
+            let (field_bindings, body_code) = expand_fields(&type_attr, &s.fields);
 
             quote!(
                 match self {
@@ -27,7 +28,7 @@ pub fn derive_shrink_to_fit(input: proc_macro::TokenStream) -> proc_macro::Token
             for v in e.variants.iter() {
                 let variant_name = &v.ident;
 
-                let (field_bindings, body_code) = expand_fields(&v.fields);
+                let (field_bindings, body_code) = expand_fields(&type_attr, &v.fields);
 
                 arms.extend(quote!(
                     Self::#variant_name { #field_bindings } => {
@@ -58,8 +59,48 @@ pub fn derive_shrink_to_fit(input: proc_macro::TokenStream) -> proc_macro::Token
     .into()
 }
 
+#[derive(Default)]
+struct TypeAttr {
+    crate_name: Option<syn::Path>,
+}
+impl TypeAttr {
+    fn parse(attrs: &[Attribute]) -> TypeAttr {
+        let mut data_attr = TypeAttr::default();
+
+        for attr in attrs {
+            if attr.path().is_ident("shrink_to_fit") {
+                if let Meta::List(meta) = &attr.meta {
+                    let tokens = meta.tokens.clone();
+                    let kv = syn::parse2::<syn::MetaNameValue>(tokens).unwrap();
+
+                    if kv.path.is_ident("crate") {
+                        if let Expr::Lit(syn::ExprLit {
+                            lit: Lit::Str(s), ..
+                        }) = &kv.value
+                        {
+                            let path = syn::parse_str::<syn::Path>(&s.value()).unwrap();
+                            data_attr.crate_name = Some(path);
+                        }
+                    }
+                }
+            }
+        }
+
+        data_attr
+    }
+}
+
 /// Returns `(field_bindings, body_code)`
-fn expand_fields(fields: &syn::Fields) -> (proc_macro2::TokenStream, proc_macro2::TokenStream) {
+fn expand_fields(
+    type_attr: &TypeAttr,
+    fields: &syn::Fields,
+) -> (proc_macro2::TokenStream, proc_macro2::TokenStream) {
+    let crate_name = type_attr
+        .crate_name
+        .as_ref()
+        .map(|q| q.to_token_stream())
+        .unwrap_or_else(|| quote!(shrink_to_fit));
+
     let mut field_bindings = proc_macro2::TokenStream::new();
     let mut body_impl = proc_macro2::TokenStream::new();
 
@@ -73,7 +114,7 @@ fn expand_fields(fields: &syn::Fields) -> (proc_macro2::TokenStream, proc_macro2
                 ));
 
                 body_impl.extend(quote!(
-                    shrink_to_fit::helpers::ShrinkToFitDerefSpecialization::new(#field_name).shrink_to_fit();
+                    #crate_name::helpers::ShrinkToFitDerefSpecialization::new(#field_name).shrink_to_fit();
                 ));
             }
         }
@@ -83,7 +124,7 @@ fn expand_fields(fields: &syn::Fields) -> (proc_macro2::TokenStream, proc_macro2
                 let field_name = Ident::new(&format!("_{}", i), field.span());
 
                 body_impl.extend(quote!(
-                    shrink_to_fit::helpers::ShrinkToFitDerefSpecialization::new(#field_name).shrink_to_fit()
+                    #crate_name::helpers::ShrinkToFitDerefSpecialization::new(#field_name).shrink_to_fit()
                 ));
 
                 let index = syn::Index::from(i);

crates/shrink-to-fit/Cargo.toml

@@ -4,7 +4,7 @@ edition     = { workspace = true }
 license     = { workspace = true }
 name        = "shrink-to-fit"
 repository  = { workspace = true }
-version     = "0.2.6"
+version     = "0.2.7"
 
 [package.metadata.docs.rs]
 all-features = true
@@ -14,7 +14,7 @@ rustdoc-args = ["--cfg", "docsrs"]
 hashbrown           = { workspace = true, optional = true, default-features = false }
 indexmap            = { workspace = true, optional = true, default-features = false }
 serde_json          = { workspace = true, optional = true, default-features = false }
-shrink-to-fit-macro = { version = "0.2.2", path = "../shrink-to-fit-macro", optional = true }
+shrink-to-fit-macro = { version = "0.2.3", path = "../shrink-to-fit-macro", optional = true }
 
 
 [features]

crates/shrink-to-fit/tests/derive.rs

@@ -65,3 +65,30 @@ fn test_nightly_specialization() {
     assert_eq!(buf[0].a.capacity(), 1);
     assert_eq!(buf[0].b.capacity(), 0);
 }
+
+#[deny(unused)]
+mod helpers {
+    pub use shrink_to_fit;
+}
+
+#[derive(Debug, ShrinkToFit)]
+#[shrink_to_fit(crate = "crate::helpers::shrink_to_fit")]
+struct ArgCheck {
+    a: String,
+    b: String,
+}
+
+#[test]
+fn test_arg_check() {
+    let mut s = ArgCheck {
+        a: String::with_capacity(100),
+        b: String::with_capacity(100),
+    };
+
+    s.a.push('a');
+
+    s.shrink_to_fit();
+
+    assert_eq!(s.a.capacity(), 1);
+    assert_eq!(s.b.capacity(), 0);
+}