fix recursion issue, better macros

examples/bouncing_ball.rs

@@ -4,7 +4,7 @@ use std::time::{Duration, Instant};
 use std::sync::{Arc, Mutex};
 use std::sync::atomic::{AtomicUsize, Ordering};
 use rand::Rng;
-use teleprof::instrument;
+use teleprof::{instrument, instrument_calls};
 
 const WIDTH: usize = 800;
 const HEIGHT: usize = 600;
@@ -39,12 +39,17 @@ fn main() {
     // Name the main thread
     teleprof::set_thread_name("Main");
 
-    println!("Bouncing Ball Demo");
+    println!("Bouncing Ball Demo - Enhanced Instrumentation");
     println!("The ball window should appear alongside the profiler");
     println!("Press Space in profiler window to pause");
     println!("Mouse wheel to zoom, drag to pan in profiler");
     println!("Press Escape in either window to quit");
     println!();
+    println!("Features demonstrated:");
+    println!("- Flame graph rendering (all depth-N spans on row N)");
+    println!("- instrument_calls macro for automatic call instrumentation");
+    println!("- Runtime deduplication preventing double-wrapping");
+    println!();
 
     let mut window = Window::new(
         "Bouncing Ball",
@@ -74,7 +79,10 @@ fn main() {
     }
 }
 
-#[instrument]
+// Using instrument_calls - automatically instruments all function calls in the body
+// This will show update_physics, render, print_status, thread::sleep, Arc::clone, etc.
+// Pause handling is now automatic in SpanGuard!
+#[instrument_calls]
 fn main_frame(
     ball: &Arc<Mutex<Ball>>, 
     framebuffer: &mut [u32], 
@@ -82,16 +90,8 @@ fn main_frame(
     frame_count: &mut u32,
     frame_time: Duration,
 ) {
-    // Check if paused
+    // Update physics - this is also instrumented with #[instrument]
-    if teleprof::PAUSE.try_lock().is_err() {
+    // Runtime dedup will prevent double-wrapping!
-        println!("Paused!");
-        while teleprof::PAUSE.try_lock().is_err() {
-            thread::sleep(Duration::from_millis(100));
-        }
-        println!("Resumed!");
-    }
-
-    // Update physics
     let hit_wall = update_physics(ball, frame_time.as_secs_f32());
     
     // If we hit a wall, spawn a thread to pick a new color
@@ -104,7 +104,7 @@ fn main_frame(
         });
     }
     
-    // Render
+    // Render - also instrumented, so dedup will handle it
     render(ball, framebuffer);
 
     // Update window
@@ -162,7 +162,8 @@ fn pick_new_color(ball: Arc<Mutex<Ball>>) {
     println!("  → New color selected: RGB({}, {}, {})", r, g, b);
 }
 
-#[instrument]
+// Using instrument_calls here too to see the breakdown of rendering
+#[instrument_calls]
 fn render(ball: &Arc<Mutex<Ball>>, framebuffer: &mut [u32]) {
     clear_background(framebuffer);
     draw_ball(ball, framebuffer);

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 {
@@ -56,3 +57,99 @@ 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(),
+    }
+}

teleprof/src/lib.rs

@@ -4,8 +4,8 @@ use std::cell::Cell;
 use std::sync::{Arc, Mutex};
 use std::time::Instant;
 
-// Re-export the macro
+// Re-export the macros
-pub use teleprof_macros::instrument;
+pub use teleprof_macros::{instrument, instrument_calls};
 
 // ============================================================================
 // Public API
@@ -33,17 +33,52 @@ pub fn set_thread_name(name: impl Into<String>) {
 }
 
 /// Create a profiling span
+#[allow(dead_code)]
 pub struct SpanGuard {
     span_id: u64,
+    previous_parent: Option<u64>,
+    previous_parent_name: Option<&'static str>,
+    name: &'static str,
+    is_duplicate: bool,
 }
 
 impl SpanGuard {
     pub fn new(name: &'static str) -> Self {
-        let span_id = next_span_id();
+        // Handle pause mechanism BEFORE creating the span
+        // This way the pause time isn't counted in the span
+        if let Err(_guard) = PAUSE.try_lock() {
+            // We're paused - block until unpaused
+            // The _guard will be dropped when we acquire the lock, releasing it immediately
+            let _pause_lock = PAUSE.lock().unwrap();
+            // Now we're unpaused, continue with span creation
+        }
+        
         let thread_id = std::thread::current().id();
         let parent_id = PARENT_SPAN.with(|p| p.get());
+        let parent_name = PARENT_SPAN_NAME.with(|p| p.get());
+        
+        // Check if this is a duplicate span (parent has same name)
+        let is_duplicate = parent_name == Some(name);
+        
+        if is_duplicate {
+            // Don't create a new span, just pass through
+            return Self {
+                span_id: 0, // Dummy value, won't be used
+                previous_parent: parent_id,
+                previous_parent_name: parent_name,
+                name,
+                is_duplicate: true,
+            };
+        }
+        
+        let span_id = next_span_id();
+        
+        // Store the previous parent so we can restore it when this span ends
+        let previous_parent = parent_id;
+        let previous_parent_name = parent_name;
         
         PARENT_SPAN.with(|p| p.set(Some(span_id)));
+        PARENT_SPAN_NAME.with(|p| p.set(Some(name)));
         
         EVENT_SENDER.send(Event::SpanStart {
             span_id,
@@ -53,19 +88,31 @@ impl SpanGuard {
             timestamp: Instant::now(),
         }).ok();
         
-        Self { span_id }
+        Self { 
+            span_id, 
+            previous_parent,
+            previous_parent_name,
+            name,
+            is_duplicate: false,
+        }
     }
 }
 
 impl Drop for SpanGuard {
     fn drop(&mut self) {
+        if self.is_duplicate {
+            // Don't send end event or modify thread-local state
+            return;
+        }
+        
         EVENT_SENDER.send(Event::SpanEnd {
             span_id: self.span_id,
             timestamp: Instant::now(),
         }).ok();
         
-        // Restore parent
+        // Restore the previous parent (grandparent of this span)
-        PARENT_SPAN.with(|p| p.set(None));
+        PARENT_SPAN.with(|p| p.set(self.previous_parent));
+        PARENT_SPAN_NAME.with(|p| p.set(self.previous_parent_name));
     }
 }
 
@@ -102,6 +149,7 @@ pub(crate) enum Event {
 
 thread_local! {
     static PARENT_SPAN: Cell<Option<u64>> = Cell::new(None);
+    static PARENT_SPAN_NAME: Cell<Option<&'static str>> = Cell::new(None);
 }
 
 static EVENT_SENDER: Lazy<Sender<Event>> = Lazy::new(|| {
@@ -719,62 +767,84 @@ mod window {
         view: &ViewState,
         font: &fontdue::Font,
     ) {
-        // Build tree structure
+        // Build tree structure and compute depths
-        let mut roots = Vec::new();
         let mut children: HashMap<u64, Vec<&CompletedSpan>> = HashMap::new();
+        let mut depths: HashMap<u64, usize> = HashMap::new();
+        let mut roots = Vec::new();
 
         for span in spans {
             if let Some(parent) = span.parent_id {
                 children.entry(parent).or_default().push(span);
             } else {
-                roots.push(*span);
+                roots.push(span.span_id);
+                depths.insert(span.span_id, 0);
             }
         }
 
-        // Sort roots by start time
+        // Compute depths for all spans using BFS
-        roots.sort_by_key(|s| s.start);
+        let mut queue: Vec<u64> = roots.clone();
+        while let Some(span_id) = queue.pop() {
+            let current_depth = depths[&span_id];
+            if let Some(child_list) = children.get(&span_id) {
+                for child in child_list {
+                    depths.insert(child.span_id, current_depth + 1);
+                    queue.push(child.span_id);
+                }
+            }
+        }
 
-        // Render each root and its children recursively
+        // Group spans by depth level
+        let mut spans_by_depth: HashMap<usize, Vec<&CompletedSpan>> = HashMap::new();
+        for span in spans {
+            if let Some(&depth) = depths.get(&span.span_id) {
+                spans_by_depth.entry(depth).or_default().push(*span);
+            }
+        }
+
+        // Find max depth
+        let max_depth = spans_by_depth.keys().max().copied().unwrap_or(0);
+
+        // Render each depth level as a row (flame graph style)
         let row_height = 24;
-        let mut y_offset = view.icicle_scroll_y as usize;
+        let scroll_offset = view.icicle_scroll_y as i32;
+        
+        for depth in 0..=max_depth {
+            let y = (depth * row_height) as i32 - scroll_offset;
             
-        for root in roots {
             // Only render if in view
-            if y_offset < height {
+            if y + row_height as i32 > 0 && y < height as i32 {
-                y_offset = render_icicle_span(
+                if let Some(depth_spans) = spans_by_depth.get(&depth) {
+                    for span in depth_spans {
+                        render_icicle_span_at_depth(
                             framebuffer,
                             width,
                             height,
-                    root,
+                            span,
-                    &children,
                             earliest,
-                    y_offset,
+                            y.max(0) as usize,
                             row_height,
                             view,
                             font,
-                    0, // depth
                         );
-            } else {
+                    }
-                break;
+                }
             }
         }
     }
 
-    fn render_icicle_span(
+    fn render_icicle_span_at_depth(
         framebuffer: &mut [u32],
         width: usize,
         height: usize,
         span: &CompletedSpan,
-        children: &HashMap<u64, Vec<&CompletedSpan>>,
         earliest: Instant,
         y: usize,
         row_height: usize,
         view: &ViewState,
         font: &fontdue::Font,
-        depth: usize,
+    ) {
-    ) -> usize {
         if y + row_height > height {
-            return y;
+            return;
         }
 
         let start_time = (span.start - earliest).as_secs_f64();
@@ -799,31 +869,6 @@ mod window {
                 draw_text(framebuffer, width, x2.max(0) as usize + 4, y + 4, span.name, font, 14.0, TEXT_COLOR);
             }
         }
-        
-        // Render children at next depth level
-        let mut child_y = y + row_height;
-        if let Some(child_spans) = children.get(&span.span_id) {
-            for child in child_spans {
-                if child_y >= height {
-                    break;
-                }
-                child_y = render_icicle_span(
-                    framebuffer,
-                    width,
-                    height,
-                    child,
-                    children,
-                    earliest,
-                    child_y,
-                    row_height,
-                    view,
-                    font,
-                    depth + 1,
-                );
-            }
-        }
-        
-        child_y
     }
 
     fn render_timeline(