Search code examples
rustmacrosrust-proc-macros

use Rust proc macros to generate dynamically named struct instance methods


I'm trying to write a procedural macro that generates methods for doubling all fields that are f64. I have it working for a single field with ./src/main.rs

use attr_macro::DoubleF64;

#[derive(DoubleF64)]
struct MyStruct {
    my_string: String,
    my_number: f64,
    my_other_number: f64,
}


fn main() {
    let mystruct = MyStruct {
        my_string: "some str".to_string(),
        my_number: 2.0,
        my_other_number: 2.0,
    };
    println!("my_number * 2: {}", mystruct.double_my_number());
}

and ./proc_macro/src/lib.rs:

extern crate proc_macro;
use proc_macro::TokenStream;
use quote::{format_ident, quote};
use syn::{parse_macro_input, DeriveInput, FieldsNamed};

#[proc_macro_derive(DoubleF64)]
pub fn double_f64(input: TokenStream) -> TokenStream {
    let DeriveInput { ident, data, .. } = parse_macro_input!(input);

    let (func_name, fident) = if let syn::Data::Struct(s) = data {
        if let syn::Fields::Named(FieldsNamed { named, .. }) = s.fields {
            let f = named[1].ident.clone().unwrap();
            (format_ident!("double_{}", f), f)
        } else {
            (format_ident!(""), format_ident!(""))
        }
    } else {
        (format_ident!(""), format_ident!(""))
    };

    let output = quote! {
        impl #ident {
            // func_str.parse.unwrap();
            // fn double_f64(&self) -> f64 {
            //     self.my_number * 2.
            // }
            fn #func_name(&self) -> f64 { self.#fident * 2. }
        }
    };

    output.into()
}

but I'm struggling to figure out how to construct a loop that generates a valid TokenStream to extend this to all the fields. Here's what I've tried:

#[proc_macro_derive(DoubleF64)]
pub fn double_f64(input: TokenStream) -> TokenStream {
    let DeriveInput { ident, data, .. } = parse_macro_input!(input);

    let mut func_stream_vec: Vec<TokenStream> = Vec::new();

    if let syn::Data::Struct(s) = data {
        if let syn::Fields::Named(FieldsNamed { named, .. }) = s.fields {
            let fields = named.iter().map(|f| &f.ident);
            let ftypes = named.iter().map(|f| &f.ty);

            for (field, ftype) in fields.into_iter().zip(ftypes) {
                if stringify!(#ftype) == "f64" {
                    let fname = format_ident!("double_{}", field.clone().unwrap());
                    func_stream_vec
                        .push(quote! { fn #fname(&self) -> f64 { self.#field * 2.0 } }.into());
                }
            }
        }
    };

    let output = quote! {
        impl #ident {
            #(#func_stream_vec)*
        }
    };

    output.into()
}

Solution

  • src/main.rs:

    // much of this code is bowrrowed from https://blog.logrocket.com/procedural-macros-in-rust/
    
    use proc_macro::DoubleF64;
    
    #[derive(DoubleF64)]
    struct MyStruct {
        my_string: String,
        my_number: f64,
        my_other_number: f64,
    }
    
    fn main() {
        let mystruct = MyStruct {
            my_string: "some str".to_string(),
            my_number: 2.0,
            my_other_number: 17.0,
        };
        println!("my_number * 2: {}", mystruct.double_my_number());
        println!("my_other_number * 2: {}", mystruct.double_my_other_number());
    }
    

    proc_macro/src/lib.rs:

    extern crate proc_macro2;
    use proc_macro2::TokenStream as TokenStream2;
    extern crate proc_macro;
    use proc_macro::TokenStream;
    use quote::{format_ident, quote, ToTokens};
    use syn::{parse_macro_input, DeriveInput, FieldsNamed, Type};
    
    extern crate quote;
    extern crate syn;
    
    #[proc_macro_derive(DoubleF64)]
    pub fn double_f64(input: TokenStream) -> TokenStream {
        let DeriveInput { ident, data, .. } = parse_macro_input!(input);
    
        let mut func_stream = TokenStream2::default();
    
        if let syn::Data::Struct(s) = data {
            if let syn::Fields::Named(FieldsNamed { named, .. }) = s.fields {
                let fields = named.iter().map(|f| &f.ident);
                let ftypes = named.iter().map(|f| &f.ty);
    
                for (field, ftype) in fields.into_iter().zip(ftypes.into_iter()) {
                    match ftype {
                        Type::Path(type_path)
                            if type_path.clone().into_token_stream().to_string() == "f64" =>
                        {
                            let fname = format_ident!("double_{}", field.clone().unwrap());
                            func_stream.extend::<TokenStream2>(
                                quote! { fn #fname(&self) -> f64 { self.#field * 2.0 } },
                            );
                        }
                        _ => {}
                    };
                }
            }
        };
    
        let output = quote! {
            impl #ident {
                #func_stream
            }
        };
    
        output.into()
    }
    

    generates:

        Finished dev [unoptimized + debuginfo] target(s) in 0.00s
         Running `target/debug/use_attr_macro`
    my_number * 2: 4
    my_other_number * 2: 34
    

    see https://github.com/calbaker/rust_proc_macro_play/tree/8afb5e088d6db81e98a2aa3f31f7831dc1e3746e