「深挖Rust」Rust macro 实例 - 5

「这是我参与2022首次更文挑战的第 21 天，活动详情查看：2022首次更文挑战」。

首先，解析属性宏所依附的代码。 syn 提供了一个内置的Rust函数语法的解析器。ItemFn 会对函数进行解析，如果语法无效会抛出一个错误。

#[proc_macro_attribute]
pub fn trace_vars(_metadata: TokenStream, input: TokenStream) -> TokenStream{
		// parsing rust function to easy to use struct
    let input_fn = parse_macro_input!(input as ItemFn);
    TokenStream::from(quote!{fn dummy(){}})
}

现在我们把输入解析了，然后我们看看元数据。对于元数据，没有内置的解析器，所以我们必须用syn 的解析模块自己写一个。

#[trace_vars(a,c,b)] // we need to parse a "," seperated list of tokens
// code

struct Args{
    vars:HashSet<Ident>
}

impl Parse for Args{
    fn parse(input: ParseStream) -> Result<Self> {
        // parses a,b,c, or a,b,c where a,b and c are Indent
        let vars = Punctuated::<Ident, Token![,]>::parse_terminated(input)?;
        Ok(Args {
            vars: vars.into_iter().collect(),
        })
    }
}

为了能和syn结合使用，我们需要实现 syn 提供的 Parse Trait。 Punctuated 是用来创建一个用 , 分隔的 Indent 集合。

一旦我们实现了 Parse Trait ，我们就可以使用 parse_macro_input macro 来解析元数据。

#[proc_macro_attribute]
pub fn trace_vars(metadata: TokenStream, input: TokenStream) -> TokenStream 
		let input_fn = parse_macro_input!(input as ItemFn);
		// using newly created struct Args
		let args= parse_macro_input!(metadata as Args);
		TokenStream::from(quote!{fn dummy(){}})
}

现在我们将修改 input_fn，以便在变量改变值时添加 println!，我们需要过滤有赋值操作的代码行，并在该行后插入一个print语句。

impl Args {
    fn should_print_expr(&self, e: &Expr) -> bool {
        match *e {
            Expr::Path(ref e) => {
                // variable shouldn't start wiht ::
                if e.path.leading_colon.is_some() {
                    false
                // should be a single variable like `x=8` not n::x=0
                } else if e.path.segments.len() != 1 {
                    false
                } else {
                    // get the first part
                    let first = e.path.segments.first().unwrap();
                    // check if the variable name is in the Args.vars hashset
                    self.vars.contains(&first.ident) && first.arguments.is_empty()
                }
            }
            _ => false,
        }
    }

    // used for checking if to print let i=0 etc or not
    fn should_print_pat(&self, p: &Pat) -> bool {
        match p {
            // check if variable name is present in set
            Pat::Ident(ref p) => self.vars.contains(&p.ident),
            _ => false,
        }
    }

    // manipulate tree to insert print statement
    fn assign_and_print(&mut self, left: Expr, op: &dyn ToTokens, right: Expr) -> Expr {
        // recurive call on right of the assigment statement
        let right = fold::fold_expr(self, right);
        // returning manipulated sub-tree
        parse_quote!({
            #left #op #right;
            println!(concat!(stringify!(#left), " = {:?}"), #left);
        })
    }

    // manipulating let statement
    fn let_and_print(&mut self, local: Local) -> Stmt {
        let Local { pat, init, .. } = local;
        let init = self.fold_expr(*init.unwrap().1);
        // get the variable name of assigned variable
        let ident = match pat {
            Pat::Ident(ref p) => &p.ident,
            _ => unreachable!(),
        };
        // new sub tree
        parse_quote! {
            let #pat = {
                #[allow(unused_mut)]
                let #pat = #init;
                println!(concat!(stringify!(#ident), " = {:?}"), #ident);
                #ident
            };
        }
    }
}

在上面的例子中，quote macro 用于模板化和编写产生的Rust代码。# 是用来注入变量的值。