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

Posted 2023-02-19

【中文标题】我们可以在过程宏属性中获取调用者的源代码位置吗？

【英文标题】：Can we get the source code location of the caller in a procedural macro attribute?

【发布时间】：2020-06-26 17:57:56

【问题描述】：

我需要获取每个方法的调用者的源位置。我正在尝试创建一个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

【问题讨论】：

另见How can I access a function's calling location each time it's called?

【参考方案1】：

提供现成的解决方案（请参阅@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

【讨论】：

谢谢维克多。我实际上应该问不同的问题。我没有找到修改令牌流中的函数的方法，这是我从您的示例中得到的。再次感谢。

我的荣幸。您是否需要更多关于如何修改功能的解释？如果需要，您还可以提出其他问题

会的。我从你的例子中得到了足够的信息。再次感谢维克多

【参考方案2】：

TL;DR

这是一个程序宏，它使用 syn 和 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"

深入讲解

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

如何编写一个函数来跟踪调用者的位置？ 见How can I access a function's calling location each time it's called?

简答：要获取调用函数的位置，用#[track_caller] 标记它并在其主体中使用std::panic::Location::caller。

如何编写创建此类函数的过程宏？

初步尝试

我们想要一个过程宏

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

例如，它会像这样转换一个函数：

fn foo() 
    // body of foo

进入

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

下面，我展示了一个程序宏，它可以准确地执行该转换——尽管，正如您将在以后的版本中看到的那样，您可能想要一些不同的东西。要尝试此代码，就像之前在 TL;DR 部分中一样，将其放入自己的 crate 并将其依赖项添加到 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 fns 组成：

它不打印调用者位置，而是打印调用我们的宏 (#[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 fns 的问题，例如，如果我们希望我们的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

我们可以添加一个特殊情况，将这种转换应用于我们的宏。但是，该转换将函数的公共 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] 函数bar、Location::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)

【讨论】：

这太棒了。谢谢。