1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156

#![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);
    }
}