arb_stylus/
trace.rs

1//! Host-I/O trace buffer for `debug_traceTransaction` on Stylus programs.
2//!
3//! Host functions call [`record`] to push a trace entry; a tracing
4//! driver installs a buffer via [`enable`] before running the program
5//! and [`take`]s the records afterwards.
6
7use std::{
8    cell::RefCell,
9    sync::{Arc, Mutex},
10};
11
12use alloy_primitives::{Address, Bytes};
13
14/// Single recorded host-I/O call.
15#[derive(Debug, Clone)]
16pub struct HostioRecord {
17    pub name: &'static str,
18    pub args: Bytes,
19    pub outs: Bytes,
20    pub start_ink: u64,
21    pub end_ink: u64,
22    pub address: Option<Address>,
23    /// Sub-frame records for CALL/CREATE family. Empty for leaf hostios.
24    pub steps: Vec<HostioRecord>,
25}
26
27thread_local! {
28    static ACTIVE: RefCell<Option<Arc<Mutex<Vec<HostioRecord>>>>> = const { RefCell::new(None) };
29    /// Stack of sub-call frames. While non-empty, recording goes into the
30    /// top frame instead of the active buffer; the parent CALL/CREATE
31    /// hostio attaches the popped frame as its `steps` field.
32    static FRAMES: RefCell<Vec<Vec<HostioRecord>>> = const { RefCell::new(Vec::new()) };
33}
34
35/// Push a fresh sub-call frame. Subsequent [`record`] calls (until the
36/// matching [`exit_subcall`]) accumulate inside this frame.
37pub fn enter_subcall() {
38    FRAMES.with(|f| f.borrow_mut().push(Vec::new()));
39}
40
41/// Pop the top sub-call frame and return its accumulated records. The
42/// parent CALL/CREATE hostio attaches this list as its `steps`.
43pub fn exit_subcall() -> Vec<HostioRecord> {
44    FRAMES.with(|f| f.borrow_mut().pop().unwrap_or_default())
45}
46
47/// Install a buffer for the current thread. Subsequent [`record`]
48/// calls append to it until [`disable`] is called.
49pub fn enable(buf: Arc<Mutex<Vec<HostioRecord>>>) {
50    ACTIVE.with(|slot| *slot.borrow_mut() = Some(buf));
51}
52
53/// Clear the active buffer for the current thread.
54pub fn disable() {
55    ACTIVE.with(|slot| *slot.borrow_mut() = None);
56}
57
58/// Take and clear the active buffer's contents.
59pub fn take() -> Vec<HostioRecord> {
60    ACTIVE
61        .with(|slot| {
62            slot.borrow()
63                .as_ref()
64                .and_then(|b| b.lock().ok().map(|mut v| std::mem::take(&mut *v)))
65        })
66        .unwrap_or_default()
67}
68
69/// Push one record into the active buffer (or the open sub-frame, if
70/// any). A no-op when tracing is disabled — zero cost on the hot path.
71pub fn record(
72    name: &'static str,
73    args: Bytes,
74    outs: Bytes,
75    start_ink: u64,
76    end_ink: u64,
77    address: Option<Address>,
78) {
79    record_with_steps(name, args, outs, start_ink, end_ink, address, Vec::new());
80}
81
82/// Like [`record`] but with pre-collected sub-frame records attached
83/// as `steps` (used by CALL/CREATE family hostios after popping their
84/// own sub-frame).
85pub fn record_with_steps(
86    name: &'static str,
87    args: Bytes,
88    outs: Bytes,
89    start_ink: u64,
90    end_ink: u64,
91    address: Option<Address>,
92    steps: Vec<HostioRecord>,
93) {
94    let rec = HostioRecord {
95        name,
96        args,
97        outs,
98        start_ink,
99        end_ink,
100        address,
101        steps,
102    };
103    let leftover = FRAMES.with(|f| {
104        let mut frames = f.borrow_mut();
105        if let Some(top) = frames.last_mut() {
106            top.push(rec);
107            None
108        } else {
109            Some(rec)
110        }
111    });
112    if let Some(rec) = leftover {
113        ACTIVE.with(|slot| {
114            if let Some(buf) = slot.borrow().as_ref() {
115                if let Ok(mut v) = buf.lock() {
116                    v.push(rec);
117                }
118            }
119        });
120    }
121}
122
123/// Whether tracing is active on the current thread — cheap check the
124/// host functions can use to avoid building args/outs when disabled.
125pub fn is_active() -> bool {
126    ACTIVE.with(|slot| slot.borrow().is_some())
127}
128
129/// Convenience wrapper for host functions that want to record the
130/// call name with optional args + outs and no ink delta (e.g., leaf
131/// host functions that never block or touch state).
132#[inline]
133pub fn record_leaf(name: &'static str, args: Bytes, outs: Bytes) {
134    if is_active() {
135        record(name, args, outs, 0, 0, None);
136    }
137}
138
139/// Record a host-function call with an ink delta captured by the
140/// caller. Used where args/outs aren't meaningful but ink cost is.
141#[inline]
142pub fn record_ink(name: &'static str, start_ink: u64, end_ink: u64) {
143    if is_active() {
144        record(name, Bytes::new(), Bytes::new(), start_ink, end_ink, None);
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn record_when_disabled_is_noop() {
154        disable();
155        assert!(!is_active());
156        record("x", Bytes::new(), Bytes::new(), 100, 50, None);
157        assert!(take().is_empty());
158    }
159
160    #[test]
161    fn enable_captures_records() {
162        let buf = Arc::new(Mutex::new(Vec::new()));
163        enable(buf.clone());
164        assert!(is_active());
165        record(
166            "storage_load_bytes32",
167            Bytes::from(vec![1]),
168            Bytes::from(vec![2]),
169            100,
170            90,
171            None,
172        );
173        record(
174            "contract_call",
175            Bytes::new(),
176            Bytes::new(),
177            90,
178            40,
179            Some(Address::repeat_byte(0xAA)),
180        );
181        let records = take();
182        assert_eq!(records.len(), 2);
183        assert_eq!(records[0].name, "storage_load_bytes32");
184        assert_eq!(records[1].address, Some(Address::repeat_byte(0xAA)));
185        disable();
186    }
187
188    #[test]
189    fn subcall_frame_nests_records() {
190        let buf = Arc::new(Mutex::new(Vec::new()));
191        enable(buf.clone());
192
193        // Top-level record before sub-call.
194        record(
195            "storage_load_bytes32",
196            Bytes::new(),
197            Bytes::new(),
198            100,
199            90,
200            None,
201        );
202
203        // Sub-call: enter, record two inner hostios, exit, record parent.
204        enter_subcall();
205        record(
206            "storage_load_bytes32",
207            Bytes::new(),
208            Bytes::new(),
209            80,
210            70,
211            None,
212        );
213        record("emit_log", Bytes::new(), Bytes::new(), 70, 60, None);
214        let steps = exit_subcall();
215        assert_eq!(steps.len(), 2);
216        record_with_steps(
217            "call_contract",
218            Bytes::new(),
219            Bytes::new(),
220            85,
221            55,
222            Some(Address::repeat_byte(0xCC)),
223            steps,
224        );
225
226        let records = take();
227        assert_eq!(records.len(), 2);
228        assert_eq!(records[0].name, "storage_load_bytes32");
229        assert_eq!(records[0].steps.len(), 0);
230        assert_eq!(records[1].name, "call_contract");
231        assert_eq!(records[1].steps.len(), 2);
232        assert_eq!(records[1].steps[0].name, "storage_load_bytes32");
233        assert_eq!(records[1].steps[1].name, "emit_log");
234        disable();
235    }
236
237    #[test]
238    fn nested_subcalls_compose() {
239        let buf = Arc::new(Mutex::new(Vec::new()));
240        enable(buf.clone());
241        enter_subcall();
242        record("a", Bytes::new(), Bytes::new(), 0, 0, None);
243        enter_subcall();
244        record("b", Bytes::new(), Bytes::new(), 0, 0, None);
245        record("c", Bytes::new(), Bytes::new(), 0, 0, None);
246        let inner = exit_subcall();
247        assert_eq!(inner.len(), 2);
248        record_with_steps("inner_call", Bytes::new(), Bytes::new(), 0, 0, None, inner);
249        record("d", Bytes::new(), Bytes::new(), 0, 0, None);
250        let outer = exit_subcall();
251        assert_eq!(outer.len(), 3);
252        assert_eq!(outer[1].name, "inner_call");
253        assert_eq!(outer[1].steps.len(), 2);
254        record_with_steps("outer_call", Bytes::new(), Bytes::new(), 0, 0, None, outer);
255        let records = take();
256        assert_eq!(records.len(), 1);
257        assert_eq!(records[0].name, "outer_call");
258        assert_eq!(records[0].steps.len(), 3);
259        disable();
260    }
261
262    #[test]
263    fn take_clears_buffer() {
264        let buf = Arc::new(Mutex::new(Vec::new()));
265        enable(buf);
266        record("foo", Bytes::new(), Bytes::new(), 10, 5, None);
267        assert_eq!(take().len(), 1);
268        assert_eq!(take().len(), 0);
269        disable();
270    }
271}