arb_stylus/
native.rs

1use arbos::programs::types::EvmData;
2use eyre::Result;
3use std::ops::{Deref, DerefMut};
4use wasmer::{
5    imports, Function, FunctionEnv, Instance, Memory, Module, Store, TypedFunction, Value,
6};
7
8use crate::{
9    cache::InitCache,
10    config::{CompileConfig, PricingParams, StylusConfig},
11    env::{MeterData, WasmEnv},
12    evm_api::EvmApi,
13    host,
14    ink::Ink,
15    meter::{
16        DepthCheckedMachine, GasMeteredMachine, MachineMeter, MeteredMachine, STYLUS_INK_LEFT,
17        STYLUS_INK_STATUS, STYLUS_STACK_LEFT,
18    },
19};
20
21/// A native WASM instance ready for execution.
22#[derive(Debug)]
23pub struct NativeInstance<E: EvmApi> {
24    pub instance: Instance,
25    pub store: Store,
26    pub env: FunctionEnv<WasmEnv<E>>,
27}
28
29impl<E: EvmApi> NativeInstance<E> {
30    pub fn new(instance: Instance, store: Store, env: FunctionEnv<WasmEnv<E>>) -> Self {
31        let mut native = Self {
32            instance,
33            store,
34            env,
35        };
36        if let Some(config) = native.env().config {
37            native.set_stack(config.max_depth);
38        }
39        native
40    }
41
42    pub fn env(&self) -> &WasmEnv<E> {
43        self.env.as_ref(&self.store)
44    }
45
46    pub fn env_mut(&mut self) -> &mut WasmEnv<E> {
47        self.env.as_mut(&mut self.store)
48    }
49
50    pub fn config(&self) -> StylusConfig {
51        self.env().config.expect("no config")
52    }
53
54    pub fn memory(&self) -> Memory {
55        self.env()
56            .memory
57            .as_ref()
58            .expect("WASM memory not initialized")
59            .clone()
60    }
61
62    /// Create from a serialized module with caching.
63    ///
64    /// # Safety
65    ///
66    /// `module` must represent a valid serialized module.
67    pub unsafe fn deserialize_cached(
68        module: &[u8],
69        version: u16,
70        evm: E,
71        evm_data: EvmData,
72        mut long_term_tag: u32,
73        debug: bool,
74    ) -> Result<Self> {
75        let compile = CompileConfig::version(version, debug);
76        let env = WasmEnv::new(compile, None, evm, evm_data);
77        let module_hash = env.evm_data.module_hash;
78        if !env.evm_data.cached {
79            long_term_tag = 0;
80        }
81        if let Some((module, store)) = InitCache::get(module_hash, version, long_term_tag, debug) {
82            return Self::from_module(module, store, env);
83        }
84        let (module, store) =
85            InitCache::insert(module_hash, module, version, long_term_tag, debug)?;
86        Self::from_module(module, store, env)
87    }
88
89    /// Create from WASM bytes (compiles the module).
90    pub fn from_bytes(
91        bytes: impl AsRef<[u8]>,
92        evm_api: E,
93        evm_data: EvmData,
94        compile: &CompileConfig,
95        config: StylusConfig,
96    ) -> Result<Self> {
97        let env = WasmEnv::new(compile.clone(), Some(config), evm_api, evm_data);
98        let store = env.compile.store();
99        let module = Module::new(&store, bytes)?;
100        Self::from_module(module, store, env)
101    }
102
103    pub fn from_module(module: Module, mut store: Store, env: WasmEnv<E>) -> Result<Self> {
104        let debug_funcs = env.compile.debug.debug_funcs;
105        let func_env = FunctionEnv::new(&mut store, env);
106
107        macro_rules! func {
108            ($func:expr) => {
109                Function::new_typed_with_env(&mut store, &func_env, $func)
110            };
111        }
112
113        let mut import_object = imports! {
114            "vm_hooks" => {
115                "read_args" => func!(host::read_args::<E>),
116                "write_result" => func!(host::write_result::<E>),
117                "exit_early" => func!(host::exit_early::<E>),
118                "storage_load_bytes32" => func!(host::storage_load_bytes32::<E>),
119                "storage_cache_bytes32" => func!(host::storage_cache_bytes32::<E>),
120                "storage_flush_cache" => func!(host::storage_flush_cache::<E>),
121                "transient_load_bytes32" => func!(host::transient_load_bytes32::<E>),
122                "transient_store_bytes32" => func!(host::transient_store_bytes32::<E>),
123                "call_contract" => func!(host::call_contract::<E>),
124                "delegate_call_contract" => func!(host::delegate_call_contract::<E>),
125                "static_call_contract" => func!(host::static_call_contract::<E>),
126                "create1" => func!(host::create1::<E>),
127                "create2" => func!(host::create2::<E>),
128                "read_return_data" => func!(host::read_return_data::<E>),
129                "return_data_size" => func!(host::return_data_size::<E>),
130                "emit_log" => func!(host::emit_log::<E>),
131                "account_balance" => func!(host::account_balance::<E>),
132                "account_code" => func!(host::account_code::<E>),
133                "account_codehash" => func!(host::account_codehash::<E>),
134                "account_code_size" => func!(host::account_code_size::<E>),
135                "evm_gas_left" => func!(host::evm_gas_left::<E>),
136                "evm_ink_left" => func!(host::evm_ink_left::<E>),
137                "block_basefee" => func!(host::block_basefee::<E>),
138                "chainid" => func!(host::chainid::<E>),
139                "block_coinbase" => func!(host::block_coinbase::<E>),
140                "block_gas_limit" => func!(host::block_gas_limit::<E>),
141                "block_number" => func!(host::block_number::<E>),
142                "block_timestamp" => func!(host::block_timestamp::<E>),
143                "contract_address" => func!(host::contract_address::<E>),
144                "math_div" => func!(host::math_div::<E>),
145                "math_mod" => func!(host::math_mod::<E>),
146                "math_pow" => func!(host::math_pow::<E>),
147                "math_add_mod" => func!(host::math_add_mod::<E>),
148                "math_mul_mod" => func!(host::math_mul_mod::<E>),
149                "msg_reentrant" => func!(host::msg_reentrant::<E>),
150                "msg_sender" => func!(host::msg_sender::<E>),
151                "msg_value" => func!(host::msg_value::<E>),
152                "tx_gas_price" => func!(host::tx_gas_price::<E>),
153                "tx_ink_price" => func!(host::tx_ink_price::<E>),
154                "tx_origin" => func!(host::tx_origin::<E>),
155                "pay_for_memory_grow" => func!(host::pay_for_memory_grow::<E>),
156                "native_keccak256" => func!(host::native_keccak256::<E>),
157            },
158        };
159
160        if debug_funcs {
161            import_object.define("console", "log_txt", func!(host::console_log_text::<E>));
162            import_object.define("console", "log_i32", func!(host::console_log::<E, u32>));
163            import_object.define("console", "log_i64", func!(host::console_log::<E, u64>));
164            import_object.define("console", "log_f32", func!(host::console_log::<E, f32>));
165            import_object.define("console", "log_f64", func!(host::console_log::<E, f64>));
166            import_object.define("console", "tee_i32", func!(host::console_tee::<E, u32>));
167            import_object.define("console", "tee_i64", func!(host::console_tee::<E, u64>));
168            import_object.define("console", "tee_f32", func!(host::console_tee::<E, f32>));
169            import_object.define("console", "tee_f64", func!(host::console_tee::<E, f64>));
170            import_object.define("debug", "null_host", func!(host::null_host::<E>));
171            import_object.define(
172                "debug",
173                "start_benchmark",
174                func!(host::start_benchmark::<E>),
175            );
176            import_object.define("debug", "end_benchmark", func!(host::end_benchmark::<E>));
177        }
178
179        let instance = Instance::new(&mut store, &module, &import_object)?;
180        let memory = instance.exports.get_memory("memory")?.clone();
181
182        let env = func_env.as_mut(&mut store);
183        env.memory = Some(memory);
184
185        let mut native = Self::new(instance, store, func_env);
186        native.set_meter_data();
187        Ok(native)
188    }
189
190    pub fn set_meter_data(&mut self) {
191        self.env_mut().meter = Some(MeterData::new());
192    }
193
194    /// Sync meter data from WASM globals.
195    pub(crate) fn sync_meter_from_globals(&mut self) {
196        let mut ink_val = 0u64;
197        let mut status_val = 0u32;
198        {
199            let store = &mut self.store;
200            let exports = &self.instance.exports;
201            if let Ok(ink_left) = exports.get_global(STYLUS_INK_LEFT) {
202                if let Value::I64(v) = ink_left.get(store) {
203                    ink_val = v as u64;
204                }
205            }
206            if let Ok(ink_status) = exports.get_global(STYLUS_INK_STATUS) {
207                if let Value::I32(v) = ink_status.get(store) {
208                    status_val = v as u32;
209                }
210            }
211        }
212        if let Some(meter) = self.env_mut().meter.as_mut() {
213            meter.set_ink(Ink(ink_val));
214            meter.set_status(status_val);
215        }
216    }
217
218    /// Push meter data to WASM globals.
219    pub(crate) fn sync_meter_to_globals(&mut self) {
220        let meter_data = self.env().meter.as_ref().map(|m| (m.ink(), m.status()));
221        if let Some((ink, status)) = meter_data {
222            let store = &mut self.store;
223            let exports = &self.instance.exports;
224            if let Ok(g) = exports.get_global(STYLUS_INK_LEFT) {
225                let _ = g.set(store, Value::I64(ink.0 as i64));
226            }
227            if let Ok(g) = exports.get_global(STYLUS_INK_STATUS) {
228                let _ = g.set(store, Value::I32(status as i32));
229            }
230        }
231    }
232
233    pub fn get_global<T>(&mut self, name: &str) -> Result<T>
234    where
235        T: TryFrom<Value>,
236        T::Error: std::fmt::Debug,
237    {
238        let store = &mut self.store;
239        let global = self
240            .instance
241            .exports
242            .get_global(name)
243            .map_err(|_| eyre::eyre!("global {name} does not exist"))?;
244        global
245            .get(store)
246            .try_into()
247            .map_err(|_| eyre::eyre!("global {name} has wrong type"))
248    }
249
250    pub fn set_global<T>(&mut self, name: &str, value: T) -> Result<()>
251    where
252        T: Into<Value>,
253    {
254        let store = &mut self.store;
255        let global = self
256            .instance
257            .exports
258            .get_global(name)
259            .map_err(|_| eyre::eyre!("global {name} does not exist"))?;
260        global
261            .set(store, value.into())
262            .map_err(|e| eyre::eyre!("{e}"))
263    }
264
265    pub fn call_func<R>(&mut self, func: TypedFunction<(), R>, ink: Ink) -> Result<R>
266    where
267        R: wasmer::WasmTypeList,
268    {
269        self.set_ink(ink);
270        self.sync_meter_to_globals();
271        let result = func.call(&mut self.store)?;
272        self.sync_meter_from_globals();
273        Ok(result)
274    }
275}
276
277impl<E: EvmApi> Deref for NativeInstance<E> {
278    type Target = Instance;
279    fn deref(&self) -> &Self::Target {
280        &self.instance
281    }
282}
283
284impl<E: EvmApi> DerefMut for NativeInstance<E> {
285    fn deref_mut(&mut self) -> &mut Self::Target {
286        &mut self.instance
287    }
288}
289
290impl<E: EvmApi> MeteredMachine for NativeInstance<E> {
291    fn ink_left(&self) -> MachineMeter {
292        let vm = self.env().meter();
293        match vm.status() {
294            0 => MachineMeter::Ready(vm.ink()),
295            _ => MachineMeter::Exhausted,
296        }
297    }
298
299    fn set_meter(&mut self, meter: MachineMeter) {
300        let vm = self.env_mut().meter_mut();
301        vm.set_ink(meter.ink());
302        vm.set_status(meter.status());
303    }
304}
305
306impl<E: EvmApi> GasMeteredMachine for NativeInstance<E> {
307    fn pricing(&self) -> PricingParams {
308        self.env()
309            .config
310            .expect("Stylus config not initialized")
311            .pricing
312    }
313}
314
315impl<E: EvmApi> DepthCheckedMachine for NativeInstance<E> {
316    fn stack_left(&mut self) -> u32 {
317        self.get_global(STYLUS_STACK_LEFT).unwrap_or(0)
318    }
319
320    fn set_stack(&mut self, size: u32) {
321        let _ = self.set_global(STYLUS_STACK_LEFT, size);
322    }
323}
324
325/// Compile WASM bytes into a serialized module.
326pub fn compile_module(wasm: &[u8], version: u16, debug: bool) -> Result<Vec<u8>> {
327    let compile = CompileConfig::version(version, debug);
328    let store = compile.store();
329    let module = Module::new(&store, wasm)?;
330    let serialized = module.serialize()?;
331    Ok(serialized.to_vec())
332}
arb_stylus/native.rs

arb_stylus/
native.rs
