简体   繁体   中英

Procedural attribute macro to inject code at the beginning of a function

I'm trying to build a simple attribute which will inject a let binding at the beginning of a function, so the result would be:

#[foo]
fn bar(){
    // start injected code
    let x = 0;
    // end injected code

    ...
}

I've gotten to this:

use proc_macro::TokenStream;
use quote::{quote, ToTokens};
use syn::{parse_macro_input, ItemFn, Item, Stmt};



#[proc_macro_attribute]
pub fn hello(attr: TokenStream, stream: TokenStream) -> TokenStream {
    let input = parse_macro_input!(stream as ItemFn);

    let block = input.block;
    let res = quote! {let x = 1;};

    // doesn't work, can't convert TokenStream into Stmt
    block.stmts.insert(0, res.into());
    TokenStream::from(input.into_token_stream())
}

However, I'm having trouble handling these items. For instance, the block is of type Vec<Stmt> , now a statement is made of pieces ( Local , Item Expr , Semi ) and when trying to work on these pieces, I just get lost. I feel like I'm missing something about how to handle these pieces, but looking at the provided example trace-vars is not helping, and a lot of the ressources online are outdated.

I've also tried a very silly approach of creating an ItemFn using quote, parsing it, and getting the Stmt from it, but then I get another error due to TokenStream actually being from two different crates, proc_macro and proc_macro2 :

#[proc_macro_attribute]
pub fn hello(attr: TokenStream, stream: TokenStream) -> TokenStream {
    let input = parse_macro_input!(stream as ItemFn);


    let res = quote! {fn () { let x = 1; } };
    let res = parse_macro_input!(res as ItemFn);

    let block = input.block;
    block.stmts.insert(0, res.block.stmts[0].clone());
    TokenStream::from(input.into_token_stream())
}
error[E0308]: mismatched types
  --> bsub-procedural/src/lib.rs:13:15
   |
13 |     let res = parse_macro_input!(res as ItemFn);
   |               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ expected struct `proc_macro::TokenStream`, found struct `proc_macro2::TokenStream`

Running on

[dependencies]
syn = { version = "1.0.67", features=["full", "fold", "printing"] }
quote = "1.0.9"

The "silly approach" is actually quite a valid solution in some situation. To parse a proc_macro2::TokenStream , use syn::parse2 instead of the macro.

The parsing sure seems wasted, since the result is immediately converted into a tokenstream again. But for small things it's fine. The alternative is to convert all parts of the function to tokens individually, eg:

let ItemFn { attrs, vis, sig, block } = input;
let stmts = &block.stmts;
quote! {
    #(#attrs)* #vis #sig {
        let x = 1; // <- your extra code
        #(#stmts)*
    }
}

The technical post webpages of this site follow the CC BY-SA 4.0 protocol. If you need to reprint, please indicate the site URL or the original address.Any question please contact:yoyou2525@163.com.

 
粤ICP备18138465号  © 2020-2024 STACKOOM.COM