arbos/programs/
mod.rs

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