arbos/programs/
mod.rs

1pub mod api;
2pub mod data_pricer;
3pub mod memory;
4pub mod params;
5pub mod types;
6
7use alloy_primitives::B256;
8use arb_chainspec::arbos_version::ARBOS_VERSION_STYLUS_FIXES;
9use arb_primitives::multigas::{MultiGas, ResourceKind};
10use revm::Database;
11
12use arb_storage::Storage;
13
14pub use self::types::{
15    evm_memory_cost, to_word_size, ActivationResult, EvmData, ProgParams, RequestType, UserOutcome,
16};
17use self::{
18    data_pricer::{init_data_pricer, open_data_pricer, DataPricer, ARBITRUM_START_TIME},
19    memory::MemoryModel,
20    params::{init_stylus_params, StylusParams},
21};
22use crate::address_set::{open_address_set, AddressSet};
23
24const PARAMS_KEY: &[u8] = &[0];
25const PROGRAM_DATA_KEY: &[u8] = &[1];
26const MODULE_HASHES_KEY: &[u8] = &[2];
27const DATA_PRICER_KEY: &[u8] = &[3];
28const CACHE_MANAGERS_KEY: &[u8] = &[4];
29
30/// Per-program metadata stored in state.
31#[derive(Debug, Clone, Copy)]
32pub struct Program {
33    pub version: u16,
34    pub init_cost: u16,
35    pub cached_cost: u16,
36    pub footprint: u16,
37    pub asm_estimate_kb: u32, // uint24 in Go
38    pub activated_at: u32,    // uint24 hours since Arbitrum began
39    pub age_seconds: u64,     // not stored in state
40    pub cached: bool,
41}
42
43impl Program {
44    /// Decode a program from a 32-byte storage word.
45    pub fn from_storage(data: B256, time: u64) -> Self {
46        let b = data.as_slice();
47        let version = u16::from_be_bytes([b[0], b[1]]);
48        let init_cost = u16::from_be_bytes([b[2], b[3]]);
49        let cached_cost = u16::from_be_bytes([b[4], b[5]]);
50        let footprint = u16::from_be_bytes([b[6], b[7]]);
51        let activated_at = (b[8] as u32) << 16 | (b[9] as u32) << 8 | b[10] as u32;
52        let asm_estimate_kb = (b[11] as u32) << 16 | (b[12] as u32) << 8 | b[13] as u32;
53        let cached = b[14] != 0;
54        let age_seconds = hours_to_age(time, activated_at);
55        Program {
56            version,
57            init_cost,
58            cached_cost,
59            footprint,
60            asm_estimate_kb,
61            activated_at,
62            age_seconds,
63            cached,
64        }
65    }
66
67    /// Encode the program to a 32-byte storage word.
68    pub fn to_storage(&self) -> B256 {
69        let mut data = [0u8; 32];
70        data[0..2].copy_from_slice(&self.version.to_be_bytes());
71        data[2..4].copy_from_slice(&self.init_cost.to_be_bytes());
72        data[4..6].copy_from_slice(&self.cached_cost.to_be_bytes());
73        data[6..8].copy_from_slice(&self.footprint.to_be_bytes());
74        // activated_at: uint24
75        data[8] = (self.activated_at >> 16) as u8;
76        data[9] = (self.activated_at >> 8) as u8;
77        data[10] = self.activated_at as u8;
78        // asm_estimate_kb: uint24
79        data[11] = (self.asm_estimate_kb >> 16) as u8;
80        data[12] = (self.asm_estimate_kb >> 8) as u8;
81        data[13] = self.asm_estimate_kb as u8;
82        data[14] = self.cached as u8;
83        B256::from(data)
84    }
85
86    /// Estimated ASM size in bytes.
87    pub fn asm_size(&self) -> u32 {
88        self.asm_estimate_kb.saturating_mul(1024)
89    }
90
91    /// Gas cost for program initialization.
92    pub fn init_gas(&self, params: &StylusParams) -> u64 {
93        let base = (params.min_init_gas as u64).saturating_mul(params::MIN_INIT_GAS_UNITS);
94        let dyno = (self.init_cost as u64)
95            .saturating_mul((params.init_cost_scalar as u64) * params::COST_SCALAR_PERCENT);
96        base.saturating_add(div_ceil(dyno, 100))
97    }
98
99    /// Gas cost for cached program initialization.
100    pub fn cached_gas(&self, params: &StylusParams) -> u64 {
101        let base = (params.min_cached_init_gas as u64).saturating_mul(params::MIN_CACHED_GAS_UNITS);
102        let dyno = (self.cached_cost as u64)
103            .saturating_mul((params.cached_cost_scalar as u64) * params::COST_SCALAR_PERCENT);
104        base.saturating_add(div_ceil(dyno, 100))
105    }
106}
107
108/// Stylus programs state.
109pub struct Programs<D> {
110    pub arbos_version: u64,
111    pub backing_storage: Storage<D>,
112    programs: Storage<D>,
113    module_hashes: Storage<D>,
114    pub data_pricer: DataPricer<D>,
115    pub cache_managers: AddressSet<D>,
116}
117
118impl<D: Database> Programs<D> {
119    pub fn initialize(arbos_version: u64, sto: &Storage<D>) {
120        let params_sto = sto.open_sub_storage(PARAMS_KEY);
121        init_stylus_params(arbos_version, &params_sto);
122        let data_pricer_sto = sto.open_sub_storage(DATA_PRICER_KEY);
123        init_data_pricer(&data_pricer_sto);
124    }
125
126    pub fn open(arbos_version: u64, sto: Storage<D>) -> Self {
127        let data_pricer_sto = sto.open_sub_storage(DATA_PRICER_KEY);
128        let data_pricer = open_data_pricer(&data_pricer_sto);
129        let programs = sto.open_sub_storage(PROGRAM_DATA_KEY);
130        let module_hashes = sto.open_sub_storage(MODULE_HASHES_KEY);
131        let cache_managers_sto = sto.open_sub_storage(CACHE_MANAGERS_KEY);
132        let cache_managers = open_address_set(cache_managers_sto);
133        Self {
134            arbos_version,
135            backing_storage: sto,
136            programs,
137            module_hashes,
138            data_pricer,
139            cache_managers,
140        }
141    }
142
143    /// Load the current Stylus parameters.
144    pub fn params(&self) -> Result<StylusParams, ()> {
145        let sto = self.backing_storage.open_sub_storage(PARAMS_KEY);
146        StylusParams::load(self.arbos_version, &sto)
147    }
148
149    /// Retrieve a program entry (may be expired or unactivated).
150    pub fn get_program(&self, code_hash: B256, time: u64) -> Result<Program, ()> {
151        let data = self.programs.get(code_hash)?;
152        Ok(Program::from_storage(data, time))
153    }
154
155    /// Retrieve and validate an active program.
156    pub fn get_active_program(
157        &self,
158        code_hash: B256,
159        time: u64,
160        params: &StylusParams,
161    ) -> Result<Program, ()> {
162        let program = self.get_program(code_hash, time)?;
163        if program.version == 0 {
164            return Err(());
165        }
166        if program.version != params.version {
167            return Err(());
168        }
169        if program.age_seconds > days_to_seconds(params.expiry_days) {
170            return Err(());
171        }
172        Ok(program)
173    }
174
175    /// Store a program entry.
176    pub fn set_program(&self, code_hash: B256, program: Program) -> Result<(), ()> {
177        self.programs.set(code_hash, program.to_storage())
178    }
179
180    /// Check if a program exists and its status.
181    pub fn program_exists(
182        &self,
183        code_hash: B256,
184        time: u64,
185        params: &StylusParams,
186    ) -> Result<(u16, bool, bool), ()> {
187        let program = self.get_program(code_hash, time)?;
188        let expired = program.activated_at == 0
189            || hours_to_age(time, program.activated_at) > days_to_seconds(params.expiry_days);
190        Ok((program.version, expired, program.cached))
191    }
192
193    /// Get the module hash for a code hash.
194    pub fn get_module_hash(&self, code_hash: B256) -> Result<B256, ()> {
195        self.module_hashes.get(code_hash)
196    }
197
198    /// Set the module hash for a code hash.
199    pub fn set_module_hash(&self, code_hash: B256, module_hash: B256) -> Result<(), ()> {
200        self.module_hashes.set(code_hash, module_hash)
201    }
202
203    /// Build runtime parameters for a program invocation.
204    pub fn prog_params(&self, version: u16, debug_mode: bool, params: &StylusParams) -> ProgParams {
205        ProgParams {
206            version,
207            max_depth: params.max_stack_depth,
208            ink_price: params.ink_price,
209            debug_mode,
210        }
211    }
212
213    /// Activate a Stylus program. Records metadata and charges data fees.
214    ///
215    /// Returns `(version, code_hash, module_hash, data_fee)` on success.
216    pub fn activate_program(
217        &self,
218        code_hash: B256,
219        wasm: &[u8],
220        time: u64,
221        page_limit: u16,
222        debug: bool,
223        activate_fn: impl FnOnce(&[u8], u16, u64, u16, bool) -> Result<ActivationResult, String>,
224    ) -> Result<(u16, B256, alloy_primitives::U256), String> {
225        let params = self.params().map_err(|_| "failed to load params")?;
226        let stylus_version = params.version;
227
228        let (current_version, expired, cached) = self
229            .program_exists(code_hash, time, &params)
230            .map_err(|_| "failed to read program")?;
231
232        if current_version == stylus_version && !expired {
233            return Err("program up to date".into());
234        }
235
236        let info = activate_fn(wasm, stylus_version, self.arbos_version, page_limit, debug)?;
237
238        // If previously cached, remove old module.
239        if cached {
240            // Old module eviction would happen at the runtime layer.
241        }
242
243        self.set_module_hash(code_hash, info.module_hash)
244            .map_err(|_| "failed to set module hash")?;
245
246        let estimate_kb = div_ceil(info.asm_estimate as u64, 1024) as u32;
247
248        let data_fee = self
249            .data_pricer
250            .update_model(info.asm_estimate, time)
251            .map_err(|_| "failed to update data pricer")?;
252
253        let program = Program {
254            version: stylus_version,
255            init_cost: info.init_gas,
256            cached_cost: info.cached_init_gas,
257            footprint: info.footprint,
258            asm_estimate_kb: estimate_kb.min(0xFF_FFFF), // uint24 max
259            activated_at: hours_since_arbitrum(time),
260            age_seconds: 0,
261            cached,
262        };
263
264        self.set_program(code_hash, program)
265            .map_err(|_| "failed to set program")?;
266
267        Ok((stylus_version, info.module_hash, data_fee))
268    }
269
270    /// Compute gas costs for calling a Stylus program.
271    ///
272    /// Returns `(call_gas_cost, memory_model)`.
273    pub fn call_gas_cost(
274        &self,
275        code_hash: B256,
276        time: u64,
277        pages_open: u16,
278        recent_cache_hit: bool,
279    ) -> Result<(u64, Program, MemoryModel), ()> {
280        let params = self.params()?;
281        let program = self.get_active_program(code_hash, time, &params)?;
282        let model = MemoryModel::new(params.free_pages, params.page_gas);
283
284        let mut cost = model.gas_cost(program.footprint, pages_open, pages_open);
285
286        let cached = program.cached || recent_cache_hit;
287        if cached || program.version > 1 {
288            cost = cost.saturating_add(program.cached_gas(&params));
289        }
290        if !cached {
291            cost = cost.saturating_add(program.init_gas(&params));
292        }
293
294        Ok((cost, program, model))
295    }
296
297    /// Extend a program's expiry by resetting its activation time.
298    ///
299    /// Returns the data fee charged.
300    pub fn program_keepalive(
301        &self,
302        code_hash: B256,
303        time: u64,
304    ) -> Result<alloy_primitives::U256, String> {
305        let params = self.params().map_err(|_| "failed to load params")?;
306        let mut program = self
307            .get_active_program(code_hash, time, &params)
308            .map_err(|_| "program not active")?;
309
310        if program.age_seconds < days_to_seconds(params.keepalive_days) {
311            return Err("keepalive too soon".into());
312        }
313        if program.version != params.version {
314            return Err("program needs upgrade".into());
315        }
316
317        let data_fee = self
318            .data_pricer
319            .update_model(program.asm_size(), time)
320            .map_err(|_| "failed to update data pricer")?;
321
322        program.activated_at = hours_since_arbitrum(time);
323        self.set_program(code_hash, program)
324            .map_err(|_| "failed to set program")?;
325
326        Ok(data_fee)
327    }
328
329    /// Execute a Stylus program call.
330    ///
331    /// Computes gas costs, resolves program metadata, and delegates to the
332    /// provided `call_fn` which handles the actual WASM runtime execution.
333    /// After execution, enforces return data cost parity with the EVM and
334    /// attributes residual gas to WasmComputation.
335    ///
336    /// The `call_fn` receives `(program, prog_params, evm_data, calldata, gas)`
337    /// and returns `(output_bytes, gas_left)` or an error.
338    pub fn call_program<F>(
339        &self,
340        code_hash: B256,
341        time: u64,
342        pages_open: u16,
343        calldata: &[u8],
344        evm_data: EvmData,
345        _reentrant: bool,
346        debug_mode: bool,
347        gas: u64,
348        used_multi_gas: &mut MultiGas,
349        recent_cache_hit: bool,
350        call_fn: F,
351    ) -> Result<(Vec<u8>, u64), String>
352    where
353        F: FnOnce(Program, ProgParams, EvmData, &[u8], u64) -> Result<(Vec<u8>, u64), String>,
354    {
355        let (call_cost, program, _model) = self
356            .call_gas_cost(code_hash, time, pages_open, recent_cache_hit)
357            .map_err(|_| "failed to compute call cost")?;
358
359        if gas < call_cost {
360            return Err("insufficient gas for Stylus call".into());
361        }
362
363        let gas_for_program = gas - call_cost;
364        let params = self.prog_params(
365            program.version,
366            debug_mode,
367            &self.params().map_err(|_| "failed to load params")?,
368        );
369
370        let starting_gas = gas;
371        let (output, gas_left) = call_fn(program, params, evm_data, calldata, gas_for_program)?;
372
373        // Ensure return data costs at least as much as it would in the EVM.
374        let gas_left = if !output.is_empty() && self.arbos_version >= ARBOS_VERSION_STYLUS_FIXES {
375            let evm_cost = evm_memory_cost(output.len() as u64);
376            if starting_gas < evm_cost {
377                // Burn all remaining gas.
378                attribute_wasm_computation(used_multi_gas, starting_gas, 0);
379                return Err("out of gas".into());
380            }
381            let max_gas_to_return = starting_gas - evm_cost;
382            gas_left.min(max_gas_to_return)
383        } else {
384            gas_left
385        };
386
387        attribute_wasm_computation(used_multi_gas, starting_gas, gas_left);
388
389        Ok((output, gas_left))
390    }
391
392    /// Update the cached status of a program.
393    pub fn set_program_cached(
394        &self,
395        code_hash: B256,
396        cache: bool,
397        time: u64,
398    ) -> Result<(), String> {
399        let params = self.params().map_err(|_| "failed to load params")?;
400        let mut program = self
401            .get_program(code_hash, time)
402            .map_err(|_| "failed to read program")?;
403
404        let expired = program.age_seconds > days_to_seconds(params.expiry_days);
405
406        if program.version != params.version && cache {
407            return Err("program needs upgrade".into());
408        }
409        if expired && cache {
410            return Err("program expired".into());
411        }
412        if program.cached == cache {
413            return Ok(());
414        }
415
416        program.cached = cache;
417        self.set_program(code_hash, program)
418            .map_err(|_| "failed to set program")?;
419
420        Ok(())
421    }
422}
423
424/// Information returned from program activation.
425#[derive(Debug, Clone)]
426pub struct ActivationInfo {
427    pub module_hash: B256,
428    pub init_gas: u16,
429    pub cached_init_gas: u16,
430    pub asm_estimate: u32,
431    pub footprint: u16,
432}
433
434/// Attribute residual WASM gas consumption to the WasmComputation resource kind.
435///
436/// After a Stylus program executes, the total gas consumed may exceed what was
437/// individually tracked through MultiGas accounting. This function assigns the
438/// residual (unaccounted) gas to `ResourceKind::WasmComputation`.
439pub fn attribute_wasm_computation(used_multi_gas: &mut MultiGas, starting_gas: u64, gas_left: u64) {
440    let used_gas = starting_gas.saturating_sub(gas_left);
441    let accounted_gas = used_multi_gas.single_gas();
442
443    let residual = if accounted_gas > used_gas {
444        tracing::trace!(
445            used_gas,
446            accounted_gas,
447            "negative WASM computation residual"
448        );
449        0
450    } else {
451        used_gas - accounted_gas
452    };
453
454    let (updated, overflow) =
455        used_multi_gas.safe_increment(ResourceKind::WasmComputation, residual);
456    if overflow {
457        tracing::trace!(residual, "WASM computation gas overflow");
458    }
459    *used_multi_gas = updated;
460}
461
462/// Hours since Arbitrum began, rounded down.
463pub fn hours_since_arbitrum(time: u64) -> u32 {
464    let elapsed = time.saturating_sub(ARBITRUM_START_TIME);
465    (elapsed / 3600).min(u32::MAX as u64) as u32
466}
467
468/// Compute program age in seconds from hours since Arbitrum began.
469pub fn hours_to_age(time: u64, hours: u32) -> u64 {
470    let seconds = (hours as u64).saturating_mul(3600);
471    let activated_at = ARBITRUM_START_TIME.saturating_add(seconds);
472    time.saturating_sub(activated_at)
473}
474
475fn days_to_seconds(days: u16) -> u64 {
476    (days as u64) * 24 * 3600
477}
478
479fn div_ceil(a: u64, b: u64) -> u64 {
480    a.div_ceil(b)
481}