teleprof/teleprof-macros/src/lib.rs

use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, ItemFn, Lit, Expr, Stmt, ExprCall, ExprMethodCall};
use syn::visit_mut::{self, VisitMut};

#[proc_macro_attribute]
pub fn instrument(args: TokenStream, input: TokenStream) -> TokenStream {
    let input_fn = parse_macro_input!(input as ItemFn);

    // Parse custom name from attributes
    let mut custom_name: Option<String> = None;
    
    let args_parser = syn::meta::parser(|meta| {
        if meta.path.is_ident("name") {
            let value: Lit = meta.value()?.parse()?;
            if let Lit::Str(lit_str) = value {
                custom_name = Some(lit_str.value());
            }
            Ok(())
        } else {
            Err(meta.error("unsupported attribute"))
        }
    });
    
    let _ = parse_macro_input!(args with args_parser);

    let fn_name = &input_fn.sig.ident;
    let span_name = custom_name.unwrap_or_else(|| fn_name.to_string());
    
    let fn_vis = &input_fn.vis;
    let fn_sig = &input_fn.sig;
    let fn_block = &input_fn.block;
    let fn_attrs = &input_fn.attrs;

    // Check if function is async
    let is_async = fn_sig.asyncness.is_some();

    let instrumented = if is_async {
        // For async functions, we need to instrument the future
        quote! {
            #(#fn_attrs)*
            #fn_vis #fn_sig {
                let _guard = ::teleprof::SpanGuard::new(#span_name);
                async move #fn_block
            }
        }
    } else {
        // For sync functions, straightforward instrumentation
        quote! {
            #(#fn_attrs)*
            #fn_vis #fn_sig {
                let _guard = ::teleprof::SpanGuard::new(#span_name);
                #fn_block
            }
        }
    };

    TokenStream::from(instrumented)
}

#[proc_macro_attribute]
pub fn instrument_calls(_args: TokenStream, input: TokenStream) -> TokenStream {
    let mut input_fn = parse_macro_input!(input as ItemFn);

    // Transform the function body to wrap all function calls
    let mut visitor = CallInstrumenter;
    visitor.visit_block_mut(&mut input_fn.block);

    // Add an outer span for the function itself
    let fn_name = input_fn.sig.ident.to_string();
    let fn_vis = &input_fn.vis;
    let fn_sig = &input_fn.sig;
    let fn_block = &input_fn.block;
    let fn_attrs = &input_fn.attrs;

    let instrumented = quote! {
        #(#fn_attrs)*
        #fn_vis #fn_sig {
            let _guard = ::teleprof::SpanGuard::new(#fn_name);
            #fn_block
        }
    };

    TokenStream::from(instrumented)
}

struct CallInstrumenter;

impl VisitMut for CallInstrumenter {
    fn visit_expr_mut(&mut self, expr: &mut Expr) {
        // First, recurse into children
        visit_mut::visit_expr_mut(self, expr);

        // Then wrap this expression if it's a function call
        match expr {
            Expr::Call(call) => {
                *expr = wrap_call(call);
            }
            Expr::MethodCall(method_call) => {
                *expr = wrap_method_call(method_call);
            }
            _ => {}
        }
    }

    fn visit_stmt_mut(&mut self, stmt: &mut Stmt) {
        // Handle statements that contain expressions
        match stmt {
            Stmt::Expr(expr, _) => {
                self.visit_expr_mut(expr);
            }
            Stmt::Local(local) => {
                if let Some(init) = &mut local.init {
                    self.visit_expr_mut(&mut init.expr);
                }
            }
            _ => {
                // Use default visitor for other statement types
                visit_mut::visit_stmt_mut(self, stmt);
            }
        }
    }
}

fn wrap_call(call: &ExprCall) -> Expr {
    // Extract function name from the call expression
    let func_name = extract_function_name(&call.func);
    
    let func = &call.func;
    let args = &call.args;
    let attrs = &call.attrs;
    
    syn::parse_quote! {
        {
            #(#attrs)*
            let _guard = ::teleprof::SpanGuard::new(#func_name);
            (#func)(#args)
        }
    }
}

fn wrap_method_call(method_call: &ExprMethodCall) -> Expr {
    let method_name = method_call.method.to_string();
    
    let receiver = &method_call.receiver;
    let method = &method_call.method;
    let args = &method_call.args;
    let turbofish = &method_call.turbofish;
    let attrs = &method_call.attrs;
    
    syn::parse_quote! {
        {
            #(#attrs)*
            let _guard = ::teleprof::SpanGuard::new(#method_name);
            (#receiver).#method #turbofish(#args)
        }
    }
}

fn extract_function_name(func: &Expr) -> String {
    match func {
        Expr::Path(path) => {
            // Get the last segment of the path
            path.path.segments.last()
                .map(|seg| seg.ident.to_string())
                .unwrap_or_else(|| "unknown".to_string())
        }
        _ => "unknown".to_string(),
    }
}