如何编写用于函数转换的Rust过程宏？

没问题，我来一步步教你实现这个过程宏！这类修改函数签名和内部逻辑的过程宏，核心是解析Rust的语法树，再重新生成修改后的代码，咱们从搭建环境开始：

第一步：搭建proc-macro项目结构 #

Rust的过程宏必须放在单独的proc-macro类型 crate 里，所以先创建一个工作区结构：

# 创建主项目（用来测试宏）
cargo new my_project
cd my_project
# 创建proc-macro子 crate
cargo new --lib my_proc_macros

然后修改my_proc_macros/Cargo.toml，声明它是proc-macro并添加依赖：

[package]
name = "my_proc_macros"
version = "0.1.0"
edition = "2021"

[lib]
proc-macro = true

[dependencies]
# 用于解析Rust语法树
syn = { version = "2.0", features = ["full", "visit-mut"] }
# 用于生成Rust代码
quote = "1.0"
# 辅助处理TokenStream
proc-macro2 = "1.0"

第二步：实现宏的核心逻辑 #

我们的宏需要完成三个核心操作：

1. 解析输入的函数，提取它的返回类型、参数和函数体
2. 修改函数签名：把返回类型改成()，新增一个*mut T类型的参数（T是原返回类型）
3. 遍历函数体，将所有return 值;替换为给指针赋值后return，把函数最后一行的表达式替换为给指针赋值

这里用syn::visit_mut来递归遍历语法树，确保所有嵌套的return语句都被处理：

use proc_macro::TokenStream;
use quote::quote;
use syn::{visit_mut::VisitMut, ItemFn, ReturnType, Type, Error};

// 语法树访问器，用来修改函数里的return语句和最后表达式
struct ReturnTransformer {
    res_ident: syn::Ident,
}

impl VisitMut for ReturnTransformer {
    // 处理所有return语句
    fn visit_expr_return_mut(&mut self, i: &mut syn::ExprReturn) {
        if let Some(expr) = &mut i.expr {
            let res = &self.res_ident;
            // 把`return value;`替换成`unsafe { *res = value; return; }`
            *i = syn::parse2(quote! {
                unsafe { *#res = #expr; }
                return;
            }).expect("Failed to parse transformed return statement");
        }
        // 继续递归处理子节点（保持完整性）
        syn::visit_mut::visit_expr_return_mut(self, i);
    }

    // 处理函数块的最后一个表达式
    fn visit_block_mut(&mut self, i: &mut syn::Block) {
        if let Some(last_stmt) = i.stmts.last_mut() {
            // 如果最后一行是表达式（不是return或let等语句），就替换为赋值给res
            if let syn::Stmt::Expr(expr) = last_stmt {
                let res = &self.res_ident;
                *last_stmt = syn::parse2(quote! {
                    unsafe { *#res = #expr; }
                }).expect("Failed to parse transformed final expression");
            }
        }
        // 递归处理块内的其他语句
        syn::visit_mut::visit_block_mut(self, i);
    }
}

#[proc_macro_attribute]
pub fn r(_attr: TokenStream, item: TokenStream) -> TokenStream {
    // 解析输入的函数，如果解析失败返回编译错误
    let mut item_fn: ItemFn = match syn::parse(item) {
        Ok(fn_item) => fn_item,
        Err(err) => return err.to_compile_error().into(),
    };

    // 提取原函数的返回类型，必须是有具体返回值的函数
    let return_type = match &item_fn.sig.output {
        ReturnType::Type(_, ty) => ty.as_ref(),
        ReturnType::Default => {
            let err = Error::new_spanned(&item_fn.sig, "`r` macro can only be applied to functions that return a non-unit value");
            return err.to_compile_error().into();
        }
    };

    // 创建新的参数：res: *mut T
    let res_ident = syn::Ident::new("res", proc_macro2::Span::call_site());
    let res_param = syn::parse2(quote! {
        #res_ident: *mut #return_type
    }).expect("Failed to create res parameter");

    // 修改函数签名：移除原返回类型，添加res参数
    item_fn.sig.output = ReturnType::Default;
    item_fn.sig.inputs.push(syn::FnArg::Typed(res_param));

    // 转换函数体里的所有return和最后表达式
    let mut transformer = ReturnTransformer { res_ident };
    transformer.visit_block_mut(&mut item_fn.block);

    // 生成最终的代码并返回
    let output = quote! {
        #item_fn
    };
    output.into()
}

第三步：测试宏的效果 #

在主项目的Cargo.toml里添加proc-macro依赖：

[dependencies]
my_proc_macros = { path = "./my_proc_macros" }

然后编写测试代码：

use my_proc_macros::r;

#[r]
fn bar(a: usize) -> usize {
    if a % 2 == 0 {
        return 0;
    }
    1
}

fn main() {
    let mut res = 0;
    // 调用转换后的函数，传入结果指针
    bar(2, &mut res as *mut usize);
    println!("bar(2) = {}", res); // 输出 0

    bar(3, &mut res as *mut usize);
    println!("bar(3) = {}", res); // 输出 1
}

一些关键注意事项 #

• 这个宏会自动处理所有嵌套的return语句，不管函数体有多复杂
• 因为操作原始指针是unsafe的，生成的代码里已经包裹了unsafe块，但调用时你需要自己保证指针的有效性（比如不能是空指针，指向的内存未被释放）
• 我添加了错误处理，如果把宏用在返回()的函数上，会给出清晰的编译错误

备注：内容来源于stack exchange，提问作者Eugene Usachev