call-trace-tls 0.4.0

Provides a way to annotate a function to inject tracing code before and after its body.
Documentation
#![warn(missing_docs)]

/*!

# Example (simple)
```

```

Output:

```text

```

# Example (custom callback)

```

```

Output

```text

*/

pub use call_trace::trace_with;
pub use call_trace::Trace;
pub use call_trace::CallContext;
use std::cell::RefCell;
use std::rc::Rc;

/// A callback. It gets called by `#[trace]`.
pub type Callback = Rc<dyn Fn(&mut Context, Event)>;

/// The thread-local callback context for `#[trace]`.
pub struct Context {
    callback: Option<Callback>,
    stack: Vec<CallContext>,
}
impl Context {
    fn new() -> Self {
        let mut r = Context {
            callback: None,
            stack: Vec::new(),
        };
        r.register_callback(|_| Self::default_callback);
        r
    }

    /// The default callback registered by `Self::new()`.
    pub fn default_callback(ctx: &mut Context, event: Event) {
        match event {
            Event::Call => {
                eprintln!("[{}:{}] => {}()", ctx.top().file, ctx.top().line, ctx.top().fn_name);
            }
            Event::Return => {
                eprintln!("[{}:{}] <= {}()", ctx.top().file, ctx.top().line, ctx.top().fn_name);
            }
        }

    }

    /// Returns the current call stack.
    pub fn stack(&self) -> &[CallContext] {
        &self.stack
    }

    /// Returns the top of the current call stack.
    pub fn top(&self) -> &CallContext {
        self.stack.last().expect("something went wrong with the callstack")
    }

    /// Registers a callback. A previously registered callback is
    /// passed to the outer closure.
    pub fn register_callback<F, G>(&mut self, f: F)
    where F: FnOnce(Option<Callback>) -> G,
          G: Fn(&mut Context, Event) + 'static
    {
        self.callback = Some(Rc::new(f(self.callback.take())));
    }

    /// Unregisters a callback.
    pub fn unregister_callback(&mut self) -> Option<Callback> {
        self.callback.take()
    }
}

#[derive(Copy, Clone, Eq, PartialEq, Ord, PartialOrd, Hash, Debug)]
/// Indicates why the callback got called by `#[trace]`
pub enum Event {
    /// The function just got called.
    Call,

    /// The function is about to return.
    Return,
}

thread_local! {
    static CONTEXT: RefCell<Context> = RefCell::new(Context::new());
}

/// Access the thread-local context of `#[trace]`
pub fn thread_access_with<T, F: FnOnce(&mut Context) -> T>(f: F) -> T {
    CONTEXT.with(move |ctx| f(&mut ctx.borrow_mut()))
}

/// Registers a callback in the thread-local context. This is a shortcut
/// for accessing the `Context` via `thread_access_with()`.
pub fn thread_register_callback<F, G>(f: F)
    where F: FnOnce(Option<Callback>) -> G,
          G: Fn(&mut Context, Event) + 'static
{
    thread_access_with(move |ctx| {
        ctx.register_callback(f)
    })
}

/// Unregisters a callback in the thread-local context. This is a shortcut
/// for accessing the `Context` via `thread_access_with()`.
pub fn thread_unregister_callback() -> Option<Callback> {
    thread_access_with(move |ctx| {
        ctx.unregister_callback()
    })
}

fn on_event(cctx: &CallContext, event: Event) {
    CONTEXT.with(|ctx| {
        let mut ctx = ctx.borrow_mut();
        if let Event::Call = event {
            ctx.stack.push(cctx.clone());
        }
        let ctx = &mut *ctx;
        if let Some(cb) = &ctx.callback {
            let cb: Callback = cb.clone();
            cb(ctx, event);
        }
        if let Event::Return = event {
            ctx.stack.pop().expect("something went wrong with the call stack");
        }
    })
}

/// Delegates to a trace kept in TLS.
pub struct Tls;

impl Trace for Tls {
    fn on_pre(&mut self, ctx: &CallContext) {
        on_event(&ctx, Event::Call);
    }
    fn on_post(&mut self, ctx: &CallContext) {
        on_event(&ctx, Event::Return);
    }
}