call_trace_tls/
lib.rs

1#![warn(missing_docs)]
2
3/*!
4
5# Example (simple)
6```
7
8```
9
10Output:
11
12```text
13
14```
15
16# Example (custom callback)
17
18```
19
20```
21
22Output
23
24```text
25
26*/
27
28pub use call_trace::trace_with;
29pub use call_trace::Trace;
30pub use call_trace::CallContext;
31use std::cell::RefCell;
32use std::rc::Rc;
33
34/// A callback. It gets called by `#[trace]`.
35pub type Callback = Rc<dyn Fn(&mut Context, Event)>;
36
37/// The thread-local callback context for `#[trace]`.
38pub struct Context {
39    callback: Option<Callback>,
40    stack: Vec<CallContext>,
41}
42impl Context {
43    fn new() -> Self {
44        let mut r = Context {
45            callback: None,
46            stack: Vec::new(),
47        };
48        r.register_callback(|_| Self::default_callback);
49        r
50    }
51
52    /// The default callback registered by `Self::new()`.
53    pub fn default_callback(ctx: &mut Context, event: Event) {
54        match event {
55            Event::Call => {
56                eprintln!("[{}:{}] => {}()", ctx.top().file, ctx.top().line, ctx.top().fn_name);
57            }
58            Event::Return => {
59                eprintln!("[{}:{}] <= {}()", ctx.top().file, ctx.top().line, ctx.top().fn_name);
60            }
61        }
62
63    }
64
65    /// Returns the current call stack.
66    pub fn stack(&self) -> &[CallContext] {
67        &self.stack
68    }
69
70    /// Returns the top of the current call stack.
71    pub fn top(&self) -> &CallContext {
72        self.stack.last().expect("something went wrong with the callstack")
73    }
74
75    /// Registers a callback. A previously registered callback is
76    /// passed to the outer closure.
77    pub fn register_callback<F, G>(&mut self, f: F)
78    where F: FnOnce(Option<Callback>) -> G,
79          G: Fn(&mut Context, Event) + 'static
80    {
81        self.callback = Some(Rc::new(f(self.callback.take())));
82    }
83
84    /// Unregisters a callback.
85    pub fn unregister_callback(&mut self) -> Option<Callback> {
86        self.callback.take()
87    }
88}
89
90#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
91/// Indicates why the callback got called by `#[trace]`
92pub enum Event {
93    /// The function just got called.
94    Call,
95
96    /// The function is about to return.
97    Return,
98}
99
100thread_local! {
101    static CONTEXT: RefCell<Context> = RefCell::new(Context::new());
102}
103
104/// Access the thread-local context of `#[trace]`
105pub fn thread_access_with<T, F: FnOnce(&mut Context) -> T>(f: F) -> T {
106    CONTEXT.with(move |ctx| f(&mut ctx.borrow_mut()))
107}
108
109/// Registers a callback in the thread-local context. This is a shortcut
110/// for accessing the `Context` via `thread_access_with()`.
111pub fn thread_register_callback<F, G>(f: F)
112    where F: FnOnce(Option<Callback>) -> G,
113          G: Fn(&mut Context, Event) + 'static
114{
115    thread_access_with(move |ctx| {
116        ctx.register_callback(f)
117    })
118}
119
120/// Unregisters a callback in the thread-local context. This is a shortcut
121/// for accessing the `Context` via `thread_access_with()`.
122pub fn thread_unregister_callback() -> Option<Callback> {
123    thread_access_with(move |ctx| {
124        ctx.unregister_callback()
125    })
126}
127
128fn on_event(cctx: &CallContext, event: Event) {
129    CONTEXT.with(|ctx| {
130        let mut ctx = ctx.borrow_mut();
131        if let Event::Call = event {
132            ctx.stack.push(cctx.clone());
133        }
134        let ctx = &mut *ctx;
135        if let Some(cb) = &ctx.callback {
136            let cb: Callback = cb.clone();
137            cb(ctx, event);
138        }
139        if let Event::Return = event {
140            ctx.stack.pop().expect("something went wrong with the call stack");
141        }
142    })
143}
144
145/// Delegates to a trace kept in TLS.
146pub struct Tls;
147
148impl Trace for Tls {
149    fn on_pre(&mut self, ctx: &CallContext) {
150        on_event(&ctx, Event::Call);
151    }
152    fn on_post(&mut self, ctx: &CallContext) {
153        on_event(&ctx, Event::Return);
154    }
155}
156