stack_debug/
lib.rs

1use proc_macro::TokenStream;
2use quote::quote;
3use syn::{ItemFn, parse_macro_input};
4
5#[proc_macro_attribute]
6pub fn instrument(attr: TokenStream, item: TokenStream) -> TokenStream {
7    let input = parse_macro_input!(item as ItemFn);
8
9    // Ensure no attributes are passed to `instrument` for now.
10    // If attributes were needed, they would be parsed from `attr`.
11    if !attr.is_empty() {
12        return syn::Error::new_spanned(
13            proc_macro2::TokenStream::from(attr),
14            "#[stack_debug::instrument] does not take any arguments",
15        )
16        .to_compile_error()
17        .into();
18    }
19
20    let block = input.block;
21    let sig = input.sig;
22    let vis = input.vis;
23    let attrs = input.attrs;
24
25    #[cfg(feature = "tracing")]
26    let expanded = quote! {
27        #(#attrs)*
28        #[inline(never)]
29        #[tracing::instrument(skip_all)]
30        #vis #sig {
31            let rbp: usize;
32            let rsp: usize;
33
34            unsafe {
35                // These instructions are for x86_64.
36                // Other architectures like AArch64 may have different
37                // conventions or registers (e.g., `fp`).
38                #[cfg(target_arch = "x86_64")]
39                std::arch::asm!(
40                    "mov {}, rbp", // Get the base pointer
41                    "mov {}, rsp", // Get the stack pointer
42                    out(reg) rbp,
43                    out(reg) rsp,
44                );
45                // Add cfgs for other architectures if needed.
46            }
47
48            // rbp should be greater than rsp. If it's not, frame pointers
49            // were likely omitted, and the result is meaningless.
50            #[cfg(debug_assertions)]
51            let frame_size = rbp - rsp - 1520; // subtract overhead of instrumentation
52            #[cfg(not(debug_assertions))]
53            let frame_size = rbp - rsp - 224; // subtract overhead of instrumentation
54            tracing::info!("stack frame size: {frame_size}");
55
56            #block
57        }
58    };
59
60    #[cfg(not(feature = "tracing"))]
61    let expanded = {
62        let function_name = sig.ident.to_string();
63        quote! {
64            #(#attrs)*
65            #[inline(never)]
66            #vis #sig {
67                let rbp: usize;
68                let rsp: usize;
69
70                unsafe {
71                    // These instructions are for x86_64.
72                    // Other architectures like AArch64 may have different
73                    // conventions or registers (e.g., `fp`).
74                    #[cfg(target_arch = "x86_64")]
75                    std::arch::asm!(
76                        "mov {}, rbp", // Get the base pointer
77                        "mov {}, rsp", // Get the stack pointer
78                        out(reg) rbp,
79                        out(reg) rsp,
80                    );
81                    // Add cfgs for other architectures if needed.
82                }
83
84                // rbp should be greater than rsp. If it's not, frame pointers
85                // were likely omitted, and the result is meaningless.
86                #[cfg(debug_assertions)]
87                let frame_size = rbp - rsp - 128; // subtract overhead of instrumentation
88                #[cfg(not(debug_assertions))]
89                let frame_size = rbp - rsp - 80; // subtract overhead of instrumentation
90                println!("{}::{}(): stack frame size: {frame_size}", module_path!(), #function_name);
91
92                #block
93            }
94        }
95    };
96
97    expanded.into()
98}
stack_debug/lib.rs

stack_debug/
lib.rs
