arb_stylus/
pages.rs

1use alloy_primitives::Address;
2use std::{
3    cell::{Cell, RefCell},
4    collections::HashMap,
5};
6
7thread_local! {
8    /// Currently open WASM memory pages across all active Stylus calls.
9    static STYLUS_PAGES_OPEN: Cell<u16> = const { Cell::new(0) };
10    /// High-water mark of pages ever open during this transaction.
11    static STYLUS_PAGES_EVER: Cell<u16> = const { Cell::new(0) };
12    /// Per-address count of open Stylus execution contexts (for reentrancy).
13    static STYLUS_PROGRAM_COUNTS: RefCell<HashMap<Address, u32>> = RefCell::new(HashMap::new());
14}
15
16/// Reset Stylus page counters at transaction start.
17pub fn reset_stylus_pages() {
18    STYLUS_PAGES_OPEN.with(|v| v.set(0));
19    STYLUS_PAGES_EVER.with(|v| v.set(0));
20    STYLUS_PROGRAM_COUNTS.with(|v| v.borrow_mut().clear());
21}
22
23/// Get current (open, ever) page counts.
24pub fn get_stylus_pages() -> (u16, u16) {
25    let open = STYLUS_PAGES_OPEN.with(|v| v.get());
26    let ever = STYLUS_PAGES_EVER.with(|v| v.get());
27    (open, ever)
28}
29
30/// Add pages for a new Stylus call. Returns previous (open, ever).
31pub fn add_stylus_pages(footprint: u16) -> (u16, u16) {
32    let open = STYLUS_PAGES_OPEN.with(|v| v.get());
33    let ever = STYLUS_PAGES_EVER.with(|v| v.get());
34    let new_open = open.saturating_add(footprint);
35    STYLUS_PAGES_OPEN.with(|v| v.set(new_open));
36    STYLUS_PAGES_EVER.with(|v| v.set(ever.max(new_open)));
37    (open, ever)
38}
39
40/// Restore page count after Stylus call returns.
41pub fn set_stylus_pages_open(open: u16) {
42    STYLUS_PAGES_OPEN.with(|v| v.set(open));
43}
44
45/// Push a Stylus program address onto the reentrancy tracker.
46/// Returns true if this is a reentrant call (address was already active).
47pub fn push_stylus_program(addr: Address) -> bool {
48    STYLUS_PROGRAM_COUNTS.with(|v| {
49        let mut map = v.borrow_mut();
50        let count = map.entry(addr).or_insert(0);
51        *count += 1;
52        *count > 1
53    })
54}
55
56/// Get the current Stylus program count for an address (no mutation).
57pub fn get_stylus_program_count(addr: Address) -> u32 {
58    STYLUS_PROGRAM_COUNTS.with(|v| v.borrow().get(&addr).copied().unwrap_or(0))
59}
60
61/// Pop a Stylus program address from the reentrancy tracker.
62pub fn pop_stylus_program(addr: Address) {
63    STYLUS_PROGRAM_COUNTS.with(|v| {
64        let mut map = v.borrow_mut();
65        if let Some(count) = map.get_mut(&addr) {
66            *count = count.saturating_sub(1);
67            if *count == 0 {
68                map.remove(&addr);
69            }
70        }
71    });
72}
arb_stylus/pages.rs

arb_stylus/
pages.rs
