我们可以在过程宏属性中获取调用者的源代码位置吗?

Can we get the source code location of the caller in a procedural macro attribute?

我有要求获取每个方法的调用者的源位置。我正在尝试创建一个 proc_macro_attribute 来捕获位置并打印它。

#[proc_macro_attribute]
pub fn get_location(attr: TokenStream, item: TokenStream) -> TokenStream {
    // Get and print file!(), line!() of source
    // Should print line no. 11
    item
}
#[get_location]
fn add(x: u32, y: u32) -> u32 {
    x + y
}

fn main() {
    add(1, 5); // Line No. 11
}

现成可用的解决方案(请参阅@timotree 的评论)。如果你想自己做,有更多的灵活性或学习,你可以编写一个过程宏来解析回溯(从被调用的函数内部获得)并打印你需要的信息。这是 lib.rs:

中的程序宏
extern crate proc_macro;
use proc_macro::{TokenStream, TokenTree};

#[proc_macro_attribute]
pub fn get_location(_attr: TokenStream, item: TokenStream) -> TokenStream {

    // prefix code to be added to the function's body
    let mut prefix: TokenStream = "
        // find earliest symbol in source file using backtrace
        let ps = Backtrace::new().frames().iter()
            .flat_map(BacktraceFrame::symbols)
            .skip_while(|s| s.filename()
                .map(|p|!p.ends_with(file!())).unwrap_or(true))
            .nth(1 as usize).unwrap();

        println!(\"Called from {:?} at line {:?}\",
            ps.filename().unwrap(), ps.lineno().unwrap());
    ".parse().unwrap(); // parse string into TokenStream

    item.into_iter().map(|tt| { // edit input TokenStream
        match tt { 
            TokenTree::Group(ref g) // match the function's body
                if g.delimiter() == proc_macro::Delimiter::Brace => { 

                    prefix.extend(g.stream()); // add parsed string

                    TokenTree::Group(proc_macro::Group::new(
                        proc_macro::Delimiter::Brace, prefix.clone()))
            },
            other => other, // else just forward TokenTree
        }
    }).collect()
} 

回溯被解析以找到源文件中最早的符号(使用 file!() 检索,另一个宏)。我们需要添加到函数中的代码是在一个字符串中定义的,然后将其解析为 TokenStream 并添加到函数主体的开头。我们可以在末尾添加这个逻辑,但是返回一个没有分号的值将不再有效。然后,您可以在 main.rs 中使用程序宏,如下所示:

extern crate backtrace;
use backtrace::{Backtrace, BacktraceFrame};
use mylib::get_location;

#[get_location]
fn add(x: u32, y: u32) -> u32 { x + y }

fn main() { 
    add(1, 41);
    add(41, 1);
}

输出为:

> Called from "src/main.rs" at line 10
> Called from "src/main.rs" at line 11

不要忘记通过将这两行添加到 Cargo.toml:

来指定您的 lib crate 正在提供程序宏
[lib]
proc-macro = true

TL;DR

这是一个程序宏,它使用 syn and quote 来执行您所描述的操作:

// print_caller_location/src/lib.rs

use proc_macro::TokenStream;
use quote::quote;
use syn::spanned::Spanned;

// Create a procedural attribute macro
//
// Notably, this must be placed alone in its own crate
#[proc_macro_attribute]
pub fn print_caller_location(_attr: TokenStream, item: TokenStream) -> TokenStream {
    // Parse the passed item as a function
    let func = syn::parse_macro_input!(item as syn::ItemFn);

    // Break the function down into its parts
    let syn::ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = func;

    // Ensure that it isn't an `async fn`
    if let Some(async_token) = sig.asyncness {
        // Error out if so
        let error = syn::Error::new(
            async_token.span(),
            "async functions do not support caller tracking functionality
    help: consider returning `impl Future` instead",
        );

        return TokenStream::from(error.to_compile_error());
    }

    // Wrap body in a closure only if function doesn't already have #[track_caller]
    let block = if attrs.iter().any(|attr| attr.path.is_ident("track_caller")) {
        quote! { #block }
    } else {
        quote! {
            (move || #block)()
        }
    };

    // Extract function name for prettier output
    let name = format!("{}", sig.ident);

    // Generate the output, adding `#[track_caller]` as well as a `println!`
    let output = quote! {
        #[track_caller]
        #(#attrs)*
        #vis #sig {
            println!(
                "entering `fn {}`: called from `{}`",
                #name,
                ::core::panic::Location::caller()
            );
            #block
        }
    };

    // Convert the output from a `proc_macro2::TokenStream` to a `proc_macro::TokenStream`
    TokenStream::from(output)
}

确保将其放入包装箱并将这些行添加到其 Cargo.toml:

# print_caller_location/Cargo.toml

[lib]
proc-macro = true

[dependencies]
syn = {version = "1.0.16", features = ["full"]}
quote = "1.0.3"
proc-macro2 = "1.0.9"

深入解释

宏只能扩展为可以手写的代码。知道这一点,我在这里看到两个问题:

  1. 如何编写一个函数来跟踪其调用者的位置?
    • Short answer: to obtain the location in which your function gets called, mark it with #[track_caller] and use std::panic::Location::caller in its body.

  2. 如何编写创建此类函数的过程宏?

初次尝试

我们想要一个程序宏

  • 取一个函数,
  • 标记为#[track_caller],
  • 并添加一行打印 Location::caller.

例如,它会像这样转换一个函数:

fn foo() {
    // body of foo
}

进入

#[track_caller]
fn foo() {
    println!("{}", std::panic::Location::caller());
    // body of foo
}

下面,我展示了一个完全执行该转换的过程宏 — 尽管正如您将在以后的版本中看到的那样,您可能想要一些不同的东西。要尝试此代码,就像之前在 TL;DR 部分中一样,将其放入自己的板条箱并将其依赖项添加到 Cargo.toml.

// print_caller_location/src/lib.rs

use proc_macro::TokenStream;
use quote::quote;

// Create a procedural attribute macro
//
// Notably, this must be placed alone in its own crate
#[proc_macro_attribute]
pub fn print_caller_location(_attr: TokenStream, item: TokenStream) -> TokenStream {
    // Parse the passed item as a function
    let func = syn::parse_macro_input!(item as syn::ItemFn);

    // Break the function down into its parts
    let syn::ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = func;

    // Extract function name for prettier output
    let name = format!("{}", sig.ident);

    // Generate the output, adding `#[track_caller]` as well as a `println!`
    let output = quote! {
        #[track_caller]
        #(#attrs)*
        #vis #sig {
            println!(
                "entering `fn {}`: called from `{}`",
                #name,
                ::core::panic::Location::caller()
            );
            #block
        }
    };

    // Convert the output from a `proc_macro2::TokenStream` to a `proc_macro::TokenStream`
    TokenStream::from(output)
}

用法示例:

// example1/src/main.rs

#![feature(track_caller)]

#[print_caller_location::print_caller_location]
fn add(x: u32, y: u32) -> u32 {
    x + y
}

fn main() {
    add(1, 5); // entering `fn add`: called from `example1/src/main.rs:11:5`
    add(1, 5); // entering `fn add`: called from `example1/src/main.rs:12:5`
}

不幸的是,我们无法摆脱那个简单的版本。该版本至少有两个问题:

  • 如何与async fn组合:

    • 它不打印调用者位置,而是打印我们的宏 (#[print_caller_location]) 被调用的位置。例如:

    // example2/src/main.rs
    
    #![feature(track_caller)]
    
    #[print_caller_location::print_caller_location]
    async fn foo() {}
    
    fn main() {
        let future = foo();
        // ^ oops! prints nothing
        futures::executor::block_on(future);
        // ^ oops! prints "entering `fn foo`: called from `example2/src/main.rs:5:1`"
        let future = foo();
        // ^ oops! prints nothing
        futures::executor::block_on(future);
        // ^ oops! prints "entering `fn foo`: called from `example2/src/main.rs:5:1`"
    }
    
  • 它如何与自身的其他调用一起工作,或者一般来说,#[track_caller]:

    • 带有 #[print_caller_location] 的嵌套函数将打印根调用者的位置,而不是给定函数的直接调用者。例如:

    // example3/src/main.rs
    
    #![feature(track_caller)]
    
    #[print_caller_location::print_caller_location]
    fn add(x: u32, y: u32) -> u32 {
        x + y
    }
    
    #[print_caller_location::print_caller_location]
    fn add_outer(x: u32, y: u32) -> u32 {
        add(x, y)
        // ^ we would expect "entering `fn add`: called from `example3/src/main.rs:12:5`"
    }
    
    fn main() {
        add(1, 5);
        // ^ "entering `fn add`: called from `example3/src/main.rs:17:5`"
        add(1, 5);
        // ^ "entering `fn add`: called from `example3/src/main.rs:19:5`"
        add_outer(1, 5);
        // ^ "entering `fn add_outer`: called from `example3/src/main.rs:21:5`"
        // ^ oops! "entering `fn add`: called from `example3/src/main.rs:21:5`"
        //
        // In reality, `add` was called on line 12, from within the body of `add_outer`
        add_outer(1, 5);
        // ^ "entering `fn add_outer`: called from `example3/src/main.rs:26:5`"
        // oops! ^ entering `fn add`: called from `example3/src/main.rs:26:5`
        //
        // In reality, `add` was called on line 12, from within the body of `add_outer`
    }
    

寻址 async fns

可以使用 -> impl Future 来解决 async fn 的问题,例如,如果我们希望我们的 async fn 反例能够正常工作,我们可以改写:

// example4/src/main.rs

#![feature(track_caller)]

use std::future::Future;

#[print_caller_location::print_caller_location]
fn foo() -> impl Future<Output = ()> {
    async move {
        // body of foo
    }
}

fn main() {
    let future = foo();
    // ^ prints "entering `fn foo`: called from `example4/src/main.rs:15:18`"
    futures::executor::block_on(future);
    // ^ prints nothing
    let future = foo();
    // ^ prints "entering `fn foo`: called from `example4/src/main.rs:19:18`"
    futures::executor::block_on(future);
    // ^ prints nothing
}

我们可以添加一个特例,将此转换应用于我们的宏。但是,除了影响返回的未来可能具有的自动特征之外,该转换还将函数的 public API 从 async fn foo() 更改为 fn foo() -> impl Future<Output = ()>

因此,我建议我们允许用户根据需要使用该解决方法,如果我们的宏用于 async fn,则只需发出错误。我们可以通过将这些行添加到我们的宏代码来做到这一点:

// Ensure that it isn't an `async fn`
if let Some(async_token) = sig.asyncness {
    // Error out if so
    let error = syn::Error::new(
        async_token.span(),
        "async functions do not support caller tracking functionality
    help: consider returning `impl Future` instead",
    );

    return TokenStream::from(error.to_compile_error());
}

修复 #[print_caller_location] 函数的嵌套行为

问题行为最小化为以下事实:当 #[track_caller] 函数 foo 直接调用另一个 #[track_caller] 函数 barLocation::caller 将使他们都可以访问 foo 的来电者。换句话说,Location::caller 在嵌套 #[track_caller] 函数的情况下允许访问根调用者:

#![feature(track_caller)]

fn main() {
    foo(); // prints `src/main.rs:4:5` instead of the line number in `foo`
}

#[track_caller]
fn foo() {
   bar();
}

#[track_caller]
fn bar() {
    println!("{}", std::panic::Location::caller());
}

playground link

为了解决这个问题,我们需要断开 #[track_caller] 调用链。我们可以通过在闭包中隐藏对 bar 的嵌套调用来打破链条:

#![feature(track_caller)]

fn main() {
    foo();
}

#[track_caller]
fn foo() {
    (move || {
        bar(); // prints `src/main.rs:10:9`
    })()
}

#[track_caller]
fn bar() {
    println!("{}", std::panic::Location::caller());
}

playground link

既然我们知道如何打破 #[track_caller] 函数链,我们就可以解决这个问题了。我们只需要确保如果用户实际上故意用 #[track_caller] 标记他们的功能,我们将避免插入闭包和破坏链。

我们可以将这些行添加到我们的解决方案中:

// Wrap body in a closure only if function doesn't already have #[track_caller]
let block = if attrs.iter().any(|attr| attr.path.is_ident("track_caller")) {
    quote! { #block }
} else {
    quote! {
        (move || #block)()
    }
};

最终解决方案

经过这两项更改后,我们得到了以下代码:

// print_caller_location/src/lib.rs

use proc_macro::TokenStream;
use quote::quote;
use syn::spanned::Spanned;

// Create a procedural attribute macro
//
// Notably, this must be placed alone in its own crate
#[proc_macro_attribute]
pub fn print_caller_location(_attr: TokenStream, item: TokenStream) -> TokenStream {
    // Parse the passed item as a function
    let func = syn::parse_macro_input!(item as syn::ItemFn);

    // Break the function down into its parts
    let syn::ItemFn {
        attrs,
        vis,
        sig,
        block,
    } = func;

    // Ensure that it isn't an `async fn`
    if let Some(async_token) = sig.asyncness {
        // Error out if so
        let error = syn::Error::new(
            async_token.span(),
            "async functions do not support caller tracking functionality
    help: consider returning `impl Future` instead",
        );

        return TokenStream::from(error.to_compile_error());
    }

    // Wrap body in a closure only if function doesn't already have #[track_caller]
    let block = if attrs.iter().any(|attr| attr.path.is_ident("track_caller")) {
        quote! { #block }
    } else {
        quote! {
            (move || #block)()
        }
    };

    // Extract function name for prettier output
    let name = format!("{}", sig.ident);

    // Generate the output, adding `#[track_caller]` as well as a `println!`
    let output = quote! {
        #[track_caller]
        #(#attrs)*
        #vis #sig {
            println!(
                "entering `fn {}`: called from `{}`",
                #name,
                ::core::panic::Location::caller()
            );
            #block
        }
    };

    // Convert the output from a `proc_macro2::TokenStream` to a `proc_macro::TokenStream`
    TokenStream::from(output)
}