arb_stylus/
native.rs

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