fix recursion issue, better macros

teleprof-macros/Cargo.toml

@@ -7,6 +7,6 @@ edition = "2021"
 proc-macro = true
 
 [dependencies]
-syn = { version = "2.0", features = ["full"] }
+syn = { version = "2.0", features = ["full", "visit-mut"] }
 quote = "1.0"
 proc-macro2 = "1.0"

teleprof-macros/src/lib.rs

@@ -1,6 +1,7 @@
 use proc_macro::TokenStream;
 use quote::quote;
-use syn::{parse_macro_input, ItemFn, Lit};
+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 {
@@ -55,4 +56,100 @@ pub fn instrument(args: TokenStream, input: TokenStream) -> TokenStream {
     };
 
     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);
+
+    TokenStream::from(quote! { #input_fn })
+}
+
+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(),
+    }
 }