Chad
Chad

Reputation: 1530

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()
}

Upvotes: 4

Views: 3542

Answers (1)

Chad
Chad

Reputation: 1530

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

Upvotes: 6

Related Questions