arb_rpc/
stylus_tracer.rs

1//! Stylus host-I/O tracer: records each WASM host function call made
2//! during Stylus program execution so `debug_traceTransaction` can
3//! surface them alongside EVM events.
4
5use std::{
6    collections::HashMap,
7    sync::{Arc, Mutex, OnceLock},
8};
9
10use alloy_primitives::{Address, Bytes, B256};
11use serde::{Deserialize, Serialize};
12
13/// One host-I/O record captured during Stylus execution.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15#[serde(rename_all = "camelCase")]
16pub struct HostioTraceInfo {
17    /// Host function name (e.g., `storage_load_bytes32`, `contract_call`).
18    pub name: String,
19    /// Arguments passed to the host function.
20    pub args: Bytes,
21    /// Outputs returned from the host function.
22    pub outs: Bytes,
23    /// Ink (gas) counter at entry.
24    pub start_ink: u64,
25    /// Ink counter at exit.
26    pub end_ink: u64,
27    /// Target address for CALL/CREATE family host functions.
28    #[serde(default, skip_serializing_if = "Option::is_none")]
29    pub address: Option<Address>,
30    /// Nested host-I/O records for sub-call frames.
31    #[serde(default, skip_serializing_if = "Vec::is_empty")]
32    pub steps: Vec<HostioTraceInfo>,
33}
34
35/// Shared recording buffer — Stylus runtime pushes; debug handler drains.
36#[derive(Debug, Default, Clone)]
37pub struct StylusTraceBuffer {
38    inner: Arc<Mutex<Vec<HostioTraceInfo>>>,
39}
40
41impl StylusTraceBuffer {
42    pub fn new() -> Self {
43        Self::default()
44    }
45
46    /// Append a single host-I/O record.
47    pub fn push(&self, record: HostioTraceInfo) {
48        if let Ok(mut g) = self.inner.lock() {
49            g.push(record);
50        }
51    }
52
53    /// Drain + return the collected records.
54    pub fn drain(&self) -> Vec<HostioTraceInfo> {
55        self.inner
56            .lock()
57            .map(|mut g| std::mem::take(&mut *g))
58            .unwrap_or_default()
59    }
60
61    /// Clear the buffer.
62    pub fn clear(&self) {
63        if let Ok(mut g) = self.inner.lock() {
64            g.clear();
65        }
66    }
67
68    /// Number of records currently buffered.
69    pub fn len(&self) -> usize {
70        self.inner.lock().map(|g| g.len()).unwrap_or(0)
71    }
72
73    /// Whether the buffer has any records.
74    pub fn is_empty(&self) -> bool {
75        self.len() == 0
76    }
77}
78
79/// Top-level tracer output attached to a `debug_traceTransaction`
80/// result when the transaction invoked a Stylus contract.
81#[derive(Debug, Clone, Default, Serialize, Deserialize)]
82#[serde(rename_all = "camelCase")]
83pub struct StylusTraceOutput {
84    /// Flat list of host-I/O records captured during tx execution.
85    pub hostio_records: Vec<HostioTraceInfo>,
86}
87
88impl From<Vec<HostioTraceInfo>> for StylusTraceOutput {
89    fn from(hostio_records: Vec<HostioTraceInfo>) -> Self {
90        Self { hostio_records }
91    }
92}
93
94/// Global cache of `tx_hash -> HostioTraceInfo[]` populated by the
95/// block producer when a tx touches a Stylus program with tracing
96/// enabled, and drained by `arb_traceStylusHostio`. Size-bounded LRU
97/// semantics (oldest entries evicted when the cap is reached).
98fn trace_cache() -> &'static Mutex<HashMap<B256, Vec<HostioTraceInfo>>> {
99    static CACHE: OnceLock<Mutex<HashMap<B256, Vec<HostioTraceInfo>>>> = OnceLock::new();
100    CACHE.get_or_init(|| Mutex::new(HashMap::new()))
101}
102
103/// Bound on cached tx-hash entries. A conservative default: most
104/// debug_trace workflows only care about recent txs.
105const TRACE_CACHE_MAX_ENTRIES: usize = 1024;
106
107/// Store the Stylus host-I/O records for a tx-hash. Called by the
108/// block producer after executing a tx with the trace buffer active.
109pub fn cache_trace(tx_hash: B256, records: Vec<HostioTraceInfo>) {
110    if records.is_empty() {
111        return;
112    }
113    if let Ok(mut m) = trace_cache().lock() {
114        if m.len() >= TRACE_CACHE_MAX_ENTRIES {
115            // Simple eviction: drop a random existing entry to stay
116            // bounded. (HashMap iteration order is nondeterministic.)
117            if let Some(key) = m.keys().next().cloned() {
118                m.remove(&key);
119            }
120        }
121        m.insert(tx_hash, records);
122    }
123}
124
125/// Retrieve + remove the cached trace for a tx-hash. Matches Nitro's
126/// one-shot retrieval semantics: the buffer is drained on first read.
127pub fn take_cached_trace(tx_hash: B256) -> Vec<HostioTraceInfo> {
128    trace_cache()
129        .lock()
130        .ok()
131        .and_then(|mut m| m.remove(&tx_hash))
132        .unwrap_or_default()
133}
134
135/// Run `f` with a fresh Stylus host-I/O trace buffer installed on the
136/// current thread. The buffer is drained and returned after `f`
137/// completes so callers can attach the records to a debug response.
138///
139/// This is the integration seam between `debug_traceTransaction` and
140/// the Stylus runtime: the debug handler wraps its tx execution in a
141/// `with_trace_buffer` call and then surfaces the records alongside
142/// the standard EVM trace.
143pub fn with_trace_buffer<F, T>(f: F) -> (T, Vec<HostioTraceInfo>)
144where
145    F: FnOnce() -> T,
146{
147    use std::sync::{Arc, Mutex};
148
149    let buf = Arc::new(Mutex::new(Vec::<arb_stylus::trace::HostioRecord>::new()));
150    arb_stylus::trace::enable(buf.clone());
151    let result = f();
152    arb_stylus::trace::disable();
153
154    let raw = buf.lock().map(|g| g.clone()).unwrap_or_default();
155    let records = raw.into_iter().map(translate_record).collect();
156    (result, records)
157}
158
159/// Recursively translate an `arb_stylus` host record (which can carry
160/// nested sub-call records under `steps`) into the RPC trace shape.
161fn translate_record(r: arb_stylus::trace::HostioRecord) -> HostioTraceInfo {
162    HostioTraceInfo {
163        name: r.name.to_string(),
164        args: r.args,
165        outs: r.outs,
166        start_ink: r.start_ink,
167        end_ink: r.end_ink,
168        address: r.address,
169        steps: r.steps.into_iter().map(translate_record).collect(),
170    }
171}
172
173#[cfg(test)]
174mod tests {
175    use super::*;
176
177    fn mk(name: &str, start_ink: u64, end_ink: u64) -> HostioTraceInfo {
178        HostioTraceInfo {
179            name: name.to_string(),
180            args: Bytes::new(),
181            outs: Bytes::new(),
182            start_ink,
183            end_ink,
184            address: None,
185            steps: Vec::new(),
186        }
187    }
188
189    #[test]
190    fn buffer_default_empty() {
191        let b = StylusTraceBuffer::new();
192        assert!(b.is_empty());
193        assert_eq!(b.len(), 0);
194    }
195
196    #[test]
197    fn buffer_push_and_drain() {
198        let b = StylusTraceBuffer::new();
199        b.push(mk("storage_load_bytes32", 100, 50));
200        b.push(mk("contract_call", 50, 10));
201        assert_eq!(b.len(), 2);
202        let drained = b.drain();
203        assert_eq!(drained.len(), 2);
204        assert_eq!(drained[0].name, "storage_load_bytes32");
205        assert!(b.is_empty());
206    }
207
208    #[test]
209    fn buffer_clear() {
210        let b = StylusTraceBuffer::new();
211        b.push(mk("emit_log", 200, 150));
212        b.clear();
213        assert!(b.is_empty());
214    }
215
216    #[test]
217    fn buffer_clone_shares_inner() {
218        let b1 = StylusTraceBuffer::new();
219        let b2 = b1.clone();
220        b1.push(mk("getCaller", 10, 9));
221        assert_eq!(b2.len(), 1);
222    }
223
224    #[test]
225    fn hostio_serde_roundtrips() {
226        let r = HostioTraceInfo {
227            name: "contract_call".to_string(),
228            args: Bytes::from(vec![0xDE, 0xAD]),
229            outs: Bytes::from(vec![0xBE, 0xEF]),
230            start_ink: 1_000,
231            end_ink: 500,
232            address: Some(Address::repeat_byte(0xAB)),
233            steps: Vec::new(),
234        };
235        let json = serde_json::to_string(&r).unwrap();
236        let back: HostioTraceInfo = serde_json::from_str(&json).unwrap();
237        assert_eq!(back.name, r.name);
238        assert_eq!(back.start_ink, r.start_ink);
239        assert_eq!(back.address, r.address);
240    }
241
242    #[test]
243    fn nested_steps_supported() {
244        let mut parent = mk("contract_call", 1_000, 400);
245        parent.steps.push(mk("storage_load_bytes32", 900, 800));
246        parent.steps.push(mk("emit_log", 800, 600));
247        assert_eq!(parent.steps.len(), 2);
248    }
249}
arb_rpc/stylus_tracer.rs

arb_rpc/
stylus_tracer.rs
