如何在Rust中定位提取被调用函数/方法逻辑以实现自动微分宏

核心思路：避免直接提取函数体，用属性宏预生成导数 #

直接提取外部函数体的方案要么依赖不稳定API，要么受限于编译单元，最简洁低侵入的方式是通过自定义属性宏为目标函数预生成导数版本，让自动微分宏直接调用预生成的导数逻辑。

1. 实现#[differentiable]属性宏 #

这个属性宏会为目标函数生成对应的导数函数，复用你已有的微分逻辑处理函数体：

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

// 假设你的自动微分宏名为`diff`，用于处理表达式生成导数
#[proc_macro_attribute]
pub fn differentiable(_attr: TokenStream, item: TokenStream) -> TokenStream {
    let input = parse_macro_input!(item as ItemFn);
    let fn_name = &input.sig.ident;
    let deriv_fn_name = syn::Ident::new(&format!("{}_deriv", fn_name), fn_name.span());
    let fn_block = &input.block;
    let fn_args = &input.sig.inputs;

    // 用你的微分宏处理原函数体，生成导数逻辑
    let deriv_block = quote! {
        #diff! {
            |#fn_args| #fn_block
        }
    };

    // 输出原函数 + 自动生成的导数函数
    let output = quote! {
        #input

        pub fn #deriv_fn_name#fn_args -> f64 {
            #deriv_block
        }
    };

    output.into()
}

2. 给外部函数（以sigmoid为例）加属性 #

同一crate内的目标函数只需加#[differentiable]，自动生成导数函数：

#[differentiable]
fn sigmoid(x: f64) -> f64 {
    1.0 / (1.0 + (-x).exp())
}

此时会自动生成sigmoid_deriv(x)函数，你的自动微分宏在遇到sigmoid(x)调用时，直接替换为调用sigmoid_deriv(x)即可完成求导。

3. 处理跨crate的外部函数 #

如果目标函数来自第三方库，无法直接修改源代码，只需写一个包装函数并加属性：

use third_party_lib::sigmoid;

#[differentiable]
fn sigmoid_wrap(x: f64) -> f64 {
    sigmoid(x)
}

之后在微分逻辑中调用sigmoid_wrap即可，属性宏会自动为包装函数生成导数。

（可选）直接提取同一crate内函数体的方案 #

如果一定要提取函数体（不推荐，依赖不稳定API），可以利用syn解析当前文件的AST，配合过程宏的Span获取源码路径：

use proc_macro::TokenStream;
use syn::{parse_file, File, ItemFn};
use std::fs;

#[proc_macro]
pub fn diff(expr: TokenStream) -> TokenStream {
    // 获取调用宏的当前文件路径（需启用不稳定特性#![feature(proc_macro_span)]）
    let call_site_span = proc_macro::Span::call_site();
    let source_path = call_site_span.source_file().path();
    let source_content = fs::read_to_string(source_path).unwrap();
    let ast = parse_file(&source_content).unwrap();

    // 遍历AST查找目标函数（比如sigmoid）
    let target_fn = ast.items.iter()
        .find_map(|item| {
            if let ItemFn(fn_item) = item {
                if fn_item.sig.ident == "sigmoid" {
                    Some(fn_item.clone())
                } else {
                    None
                }
            } else {
                None
            }
        }).expect("Target function not found in current file");

    // 用你的微分逻辑处理提取到的函数体
    let deriv_logic = quote! { /* 你的微分处理代码 */ };

    deriv_logic.into()
}

注意：这个方案只能处理同一文件内的函数，且依赖Rust的不稳定特性，跨crate函数无法通过此方式提取函数体。

内容的提问来源于stack exchange，提问作者Joe McCain III