Merge pull request #95 from OffchainLabs/onchain-compiled

Onchain compiled

Makefile

@@ -153,6 +153,7 @@ build-node-deps: $(go_source) build-prover-header build-prover-lib build-jit .ma
 test-go-deps: \
 	build-replay-env \
 	$(stylus_test_wasms) \
+	$(arbitrator_stylus_lib) \
 	$(patsubst %,$(arbitrator_cases)/%.wasm, global-state read-inboxmsg-10 global-state-wrapper const)
 
 build-prover-header: $(arbitrator_generated_header)

arbitrator/arbutil/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2021"
 digest = "0.9.0"
 eyre = "0.6.5"
 hex = "0.4.3"
+num-traits = "0.2.17"
 sha3 = "0.10.5"
 siphasher = "0.3.10"
 wasmparser = "0.83"

arbitrator/arbutil/src/math.rs

@@ -1,6 +1,7 @@
 // Copyright 2023, Offchain Labs, Inc.
 // For license information, see https://github.com/nitro/blob/master/LICENSE
 
+use num_traits::{ops::saturating::SaturatingAdd, Zero};
 use std::ops::{BitAnd, Sub};
 
 /// Checks if a number is a power of 2.
@@ -13,3 +14,22 @@ where
     }
     value & (value - 1.into()) == 0.into()
 }
+
+/// Calculates a sum, saturating in cases of overflow.
+pub trait SaturatingSum {
+    type Number;
+
+    fn saturating_sum(self) -> Self::Number;
+}
+
+impl<I, T> SaturatingSum for I
+where
+    I: Iterator<Item = T>,
+    T: SaturatingAdd + Zero,
+{
+    type Number = T;
+
+    fn saturating_sum(self) -> Self::Number {
+        self.fold(T::zero(), |acc, x| acc.saturating_add(&x))
+    }
+}

arbitrator/jit/src/gostack.rs

@@ -11,7 +11,10 @@ use crate::{
 use arbutil::Color;
 use ouroboros::self_referencing;
 use rand_pcg::Pcg32;
-use std::collections::{BTreeSet, BinaryHeap};
+use std::{
+    collections::{BTreeSet, BinaryHeap},
+    fmt::Debug,
+};
 use wasmer::{AsStoreRef, Memory, MemoryView, StoreMut, StoreRef, WasmPtr};
 
 #[self_referencing]
@@ -138,6 +141,10 @@ impl GoStack {
         self.read_u64() as *mut T
     }
 
+    pub unsafe fn read_ref<'a, T>(&mut self) -> &'a T {
+        &*self.read_ptr()
+    }
+
     /// TODO: replace `unbox` with a safe id-based API
     pub fn unbox<T>(&mut self) -> T {
         let ptr: *mut T = self.read_ptr_mut();
@@ -236,9 +243,12 @@ impl GoStack {
         data
     }
 
-    pub fn write_slice(&self, ptr: u64, src: &[u8]) {
-        u32::try_from(ptr).expect("Go pointer not a u32");
-        self.view().write(ptr, src).unwrap();
+    pub fn write_slice<T: TryInto<u32>>(&self, ptr: T, src: &[u8])
+    where
+        T::Error: Debug,
+    {
+        let ptr: u32 = ptr.try_into().expect("Go pointer not a u32");
+        self.view().write(ptr.into(), src).unwrap();
     }
 
     pub fn read_value_slice(&self, mut ptr: u64, len: u64) -> Vec<JsValue> {

arbitrator/jit/src/machine.rs

@@ -21,6 +21,7 @@ use std::{
     io::{self, Write},
     io::{BufReader, BufWriter, ErrorKind, Read},
     net::TcpStream,
+    sync::Arc,
     time::{Duration, Instant},
 };
 
@@ -114,11 +115,10 @@ pub fn create(opts: &Opts, env: WasmEnv) -> (Instance, FunctionEnv<WasmEnv>, Sto
             github!("wavmio.readDelayedInboxMessage") => func!(wavmio::read_delayed_inbox_message),
             github!("wavmio.resolvePreImage") => func!(wavmio::resolve_preimage),
 
-            github!("arbos/programs.compileUserWasmRustImpl") => func!(user::compile_user_wasm),
-            github!("arbos/programs.callUserWasmRustImpl") => func!(user::call_user_wasm),
+            github!("arbos/programs.activateProgramRustImpl") => func!(user::stylus_activate),
+            github!("arbos/programs.callProgramRustImpl") => func!(user::stylus_call),
             github!("arbos/programs.readRustVecLenImpl") => func!(user::read_rust_vec_len),
             github!("arbos/programs.rustVecIntoSliceImpl") => func!(user::rust_vec_into_slice),
-            github!("arbos/programs.rustMachineDropImpl") => func!(user::drop_machine),
             github!("arbos/programs.rustConfigImpl") => func!(user::rust_config_impl),
             github!("arbos/programs.rustEvmDataImpl") => func!(user::evm_data_impl),
 
@@ -193,10 +193,7 @@ impl From<RuntimeError> for Escape {
 pub type WasmEnvMut<'a> = FunctionEnvMut<'a, WasmEnv>;
 pub type Inbox = BTreeMap<u64, Vec<u8>>;
 pub type Oracle = BTreeMap<Bytes32, Vec<u8>>;
-
-/// Represents a mapping of a WASM program codehash and version to the compiled wasm
-/// code itself and its noncanonical program hash.
-pub type UserWasms = HashMap<(Bytes32, u16), (Vec<u8>, Bytes32)>;
+pub type ModuleAsm = Arc<[u8]>;
 
 #[derive(Default)]
 pub struct WasmEnv {
@@ -212,8 +209,8 @@ pub struct WasmEnv {
     pub large_globals: [Bytes32; 2],
     /// An oracle allowing the prover to reverse keccak256
     pub preimages: Oracle,
-    /// A collection of user wasms called during the course of execution
-    pub user_wasms: UserWasms,
+    /// A collection of programs called during the course of execution
+    pub module_asms: HashMap<Bytes32, ModuleAsm>,
     /// The sequencer inbox's messages
     pub sequencer_messages: Inbox,
     /// The delayed inbox's messages

arbitrator/jit/src/socket.rs

@@ -20,11 +20,6 @@ pub fn read_u8<T: Read>(reader: &mut BufReader<T>) -> Result<u8, io::Error> {
     reader.read_exact(&mut buf).map(|_| u8::from_be_bytes(buf))
 }
 
-pub fn read_u16<T: Read>(reader: &mut BufReader<T>) -> Result<u16, io::Error> {
-    let mut buf = [0; 2];
-    reader.read_exact(&mut buf).map(|_| u16::from_be_bytes(buf))
-}
-
 pub fn read_u32<T: Read>(reader: &mut BufReader<T>) -> Result<u32, io::Error> {
     let mut buf = [0; 4];
     reader.read_exact(&mut buf).map(|_| u32::from_be_bytes(buf))
@@ -47,6 +42,10 @@ pub fn read_bytes<T: Read>(reader: &mut BufReader<T>) -> Result<Vec<u8>, io::Err
     Ok(buf)
 }
 
+pub fn read_boxed_slice<T: Read>(reader: &mut BufReader<T>) -> Result<Box<[u8]>, io::Error> {
+    Ok(Vec::into_boxed_slice(read_bytes(reader)?))
+}
+
 pub fn write_u8(writer: &mut BufWriter<TcpStream>, data: u8) -> Result<(), io::Error> {
     let buf = [data; 1];
     writer.write_all(&buf)

arbitrator/jit/src/user/evm_api.rs

@@ -5,7 +5,7 @@
 
 use crate::{
     gostack::GoStack,
-    machine::WasmEnvMut,
+    machine::{ModuleAsm, WasmEnvMut},
     syscall::{DynamicObject, GoValue, JsValue, STYLUS_ID},
 };
 use arbutil::{
@@ -53,7 +53,7 @@ impl JsCallIntoGo for ApiCaller {
 pub(super) fn exec_wasm(
     sp: &mut GoStack,
     mut env: WasmEnvMut,
-    module: Vec<u8>,
+    module: ModuleAsm,
     calldata: Vec<u8>,
     compile: CompileConfig,
     config: StylusConfig,

arbitrator/jit/src/user/mod.rs

@@ -5,60 +5,64 @@ use crate::{
     gostack::GoStack,
     machine::{Escape, MaybeEscape, WasmEnvMut},
     user::evm_api::exec_wasm,
+    wavmio::Bytes32,
 };
 use arbutil::{
     evm::{user::UserOutcome, EvmData},
     format::DebugBytes,
     heapify,
 };
 use prover::{
+    machine::Module,
     programs::{config::PricingParams, prelude::*},
-    Machine,
 };
 use std::mem;
-use stylus::native;
 
 mod evm_api;
 
-/// Compiles and instruments a user wasm.
+/// Instruments and "activates" a user wasm, producing a unique module hash.
+///
+/// Note that this operation costs gas and is limited by the amount supplied via the `gas` pointer.
+/// The amount left is written back at the end of the call.
 ///
 /// # Go side
 ///
+/// The `modHash` and `gas` pointers must not be null.
+///
 /// The Go compiler expects the call to take the form
-///     λ(wasm []byte, pageLimit, version u16, debug u32) (module *Vec<u8>, info WasmInfo, err *Vec<u8>)
+///     λ(wasm []byte, pageLimit, version u16, debug u32, modHash *hash, gas *u64) (footprint u16, err *Vec<u8>)
 ///
 /// These values are placed on the stack as follows
-///     stack:  || wasm... || pageLimit | version | debug || mod ptr || info... || err ptr ||
-///     info:   || footprint | 2 pad | size ||
+///     || wasm... || pageLimit | version | debug || modhash ptr || gas ptr || footprint | 6 pad || err ptr ||
 ///
-pub fn compile_user_wasm(env: WasmEnvMut, sp: u32) {
+pub fn stylus_activate(env: WasmEnvMut, sp: u32) {
     let mut sp = GoStack::simple(sp, &env);
     let wasm = sp.read_go_slice_owned();
     let page_limit = sp.read_u16();
     let version = sp.read_u16();
     let debug = sp.read_bool32();
-    let compile = CompileConfig::version(version, debug);
+    let module_hash = sp.read_go_ptr();
+    let gas = sp.read_go_ptr();
 
     macro_rules! error {
         ($error:expr) => {{
-            let error = $error.wrap_err("failed to compile").debug_bytes();
-            sp.write_nullptr();
-            sp.skip_space(); // skip info
+            let error = $error.wrap_err("failed to activate").debug_bytes();
+            sp.write_u64_raw(gas, 0);
+            sp.write_slice(module_hash, &Bytes32::default());
+            sp.skip_space();
             sp.write_ptr(heapify(error));
             return;
         }};
     }
 
-    let (footprint, size) = match Machine::new_user_stub(&wasm, page_limit, version, debug) {
-        Ok((_, info)) => (info.footprint, info.size),
-        Err(error) => error!(error),
-    };
-    let module = match native::module(&wasm, compile) {
-        Ok(module) => module,
+    let gas_left = &mut sp.read_u64_raw(gas);
+    let (module, pages) = match Module::activate(&wasm, version, page_limit, debug, gas_left) {
+        Ok(result) => result,
         Err(error) => error!(error),
     };
-    sp.write_ptr(heapify(module));
-    sp.write_u16(footprint).skip_u16().write_u32(size); // wasm info
+    sp.write_u64_raw(gas, *gas_left);
+    sp.write_slice(module_hash, &module.hash().0);
+    sp.write_u16(pages).skip_space();
     sp.write_nullptr();
 }
 
@@ -67,32 +71,35 @@ pub fn compile_user_wasm(env: WasmEnvMut, sp: u32) {
 /// # Go side
 ///
 /// The Go compiler expects the call to take the form
-///     λ(
-///           mach *Machine, calldata []byte, params *Configs, evmApi []byte, evmData: *EvmData,
-///           gas *u64, root *[32]byte
-///     ) -> (status byte, out *Vec<u8>)
+///     λ(moduleHash *[32]byte, calldata []byte, params *Configs, evmApi []byte, evmData: *EvmData, gas *u64) (
+///         status byte, out *Vec<u8>,
+///     )
 ///
 /// These values are placed on the stack as follows
-///     || mach || calldata... || params || evmApi... || evmData || gas || root || status | 3 pad | out ptr ||
+///     || modHash || calldata... || params || evmApi... || evmData || gas || status | 7 pad | out ptr ||
 ///
-pub fn call_user_wasm(env: WasmEnvMut, sp: u32) -> MaybeEscape {
+pub fn stylus_call(env: WasmEnvMut, sp: u32) -> MaybeEscape {
     let sp = &mut GoStack::simple(sp, &env);
     use UserOutcome::*;
 
     // move inputs
-    let module: Vec<u8> = sp.unbox();
+    let module_hash = sp.read_bytes32();
     let calldata = sp.read_go_slice_owned();
     let (compile, config): (CompileConfig, StylusConfig) = sp.unbox();
     let evm_api = sp.read_go_slice_owned();
     let evm_data: EvmData = sp.unbox();
+    let gas = sp.read_go_ptr();
 
     // buy ink
     let pricing = config.pricing;
-    let gas = sp.read_go_ptr();
     let ink = pricing.gas_to_ink(sp.read_u64_raw(gas));
 
-    // skip the root since we don't use these
-    sp.skip_u64();
+    let Some(module) = env.data().module_asms.get(&module_hash).cloned() else {
+        return Escape::failure(format!(
+            "module hash {module_hash:?} not found in {:?}",
+            env.data().module_asms.keys()
+        ));
+    };
 
     let result = exec_wasm(
         sp, env, module, calldata, compile, config, evm_api, evm_data, ink,
@@ -122,7 +129,7 @@ pub fn call_user_wasm(env: WasmEnvMut, sp: u32) -> MaybeEscape {
 ///
 pub fn read_rust_vec_len(env: WasmEnvMut, sp: u32) {
     let mut sp = GoStack::simple(sp, &env);
-    let vec: &Vec<u8> = unsafe { &*sp.read_ptr() };
+    let vec: &Vec<u8> = unsafe { sp.read_ref() };
     sp.write_u32(vec.len() as u32);
 }
 
@@ -144,20 +151,6 @@ pub fn rust_vec_into_slice(env: WasmEnvMut, sp: u32) {
     mem::drop(vec)
 }
 
-/// Drops module bytes. Note that in user-host this would be a `Machine`.
-///
-/// # Go side
-///
-/// The Go compiler expects the call to take the form
-///     λ(module *Vec<u8>)
-///
-pub fn drop_machine(env: WasmEnvMut, sp: u32) {
-    let mut sp = GoStack::simple(sp, &env);
-    if let Some(module) = sp.unbox_option::<Vec<u8>>() {
-        mem::drop(module);
-    }
-}
-
 /// Creates a `StylusConfig` from its component parts.
 ///
 /// # Go side

arbitrator/jit/src/wavmio.rs

@@ -314,11 +314,9 @@ fn ready_hostio(env: &mut WasmEnv) -> MaybeEscape {
 
     let programs_count = socket::read_u32(stream)?;
     for _ in 0..programs_count {
-        let codehash = socket::read_bytes32(stream)?;
-        let wasm = socket::read_bytes(stream)?;
-        let hash = socket::read_bytes32(stream)?;
-        let version = socket::read_u16(stream)?;
-        env.user_wasms.insert((codehash, version), (wasm, hash));
+        let module_hash = socket::read_bytes32(stream)?;
+        let module_asm = socket::read_boxed_slice(stream)?;
+        env.module_asms.insert(module_hash, module_asm.into());
     }
 
     if socket::read_u8(stream)? != socket::READY {