New Module Hash Model

Makefile

@@ -259,17 +259,17 @@ $(replay_wasm): $(DEP_PREDICATE) $(go_source) .make/solgen
 	mkdir -p `dirname $(replay_wasm)`
 	GOOS=js GOARCH=wasm go build -o $@ ./cmd/replay/...
 
-$(prover_bin): $(DEP_PREDICATE) $(rust_prover_files) $(output_latest)/forward_stub.wasm
+$(prover_bin): $(DEP_PREDICATE) $(rust_prover_files)
 	mkdir -p `dirname $(prover_bin)`
 	cargo build --manifest-path arbitrator/Cargo.toml --release --bin prover ${CARGOFLAGS}
 	install arbitrator/target/release/prover $@
 
-$(arbitrator_stylus_lib): $(DEP_PREDICATE) $(stylus_files) $(output_latest)/forward_stub.wasm
+$(arbitrator_stylus_lib): $(DEP_PREDICATE) $(stylus_files)
 	mkdir -p `dirname $(arbitrator_stylus_lib)`
 	cargo build --manifest-path arbitrator/Cargo.toml --release --lib -p stylus ${CARGOFLAGS}
 	install arbitrator/target/release/libstylus.a $@
 
-$(arbitrator_jit): $(DEP_PREDICATE) .make/cbrotli-lib $(jit_files) $(output_latest)/forward_stub.wasm
+$(arbitrator_jit): $(DEP_PREDICATE) .make/cbrotli-lib $(jit_files)
 	mkdir -p `dirname $(arbitrator_jit)`
 	cargo build --manifest-path arbitrator/Cargo.toml --release -p jit ${CARGOFLAGS}
 	install arbitrator/target/release/jit $@
@@ -352,7 +352,6 @@ $(output_latest)/forward.wasm: $(DEP_PREDICATE) $(wasm_lib)/user-host/forward.wa
 	wat2wasm $(wasm_lib)/user-host/forward.wat -o $@
 
 $(output_latest)/forward_stub.wasm: $(DEP_PREDICATE) $(wasm_lib)/user-host/forward_stub.wat .make/machines
-	mkdir -p $(output_latest)
 	wat2wasm $(wasm_lib)/user-host/forward_stub.wat -o $@
 
 $(output_latest)/machine.wavm.br: $(DEP_PREDICATE) $(prover_bin) $(arbitrator_wasm_libs) $(replay_wasm)

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.rustModuleDropImpl") => func!(user::drop_module),
             github!("arbos/programs.rustConfigImpl") => func!(user::rust_config_impl),
             github!("arbos/programs.rustEvmDataImpl") => func!(user::evm_data_impl),
 
@@ -193,6 +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>>;
+pub type ModuleAsm = Arc<[u8]>;
 
 #[derive(Default)]
 pub struct WasmEnv {
@@ -208,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 compiled_modules: HashMap<Bytes32, Vec<u8>>,
+    /// 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,63 +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::programs::{config::PricingParams, prelude::*};
+use prover::{
+    machine::Module,
+    programs::{config::PricingParams, prelude::*},
+};
 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 (out_hash_ptr, out_hash_len) = sp.read_go_slice();
+    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;
         }};
     }
 
-    if out_hash_len != 32 {
-        error!(eyre::eyre!(
-            "Go attempting to read compiled machine hash into bad buffer length: {out_hash_len}"
-        ));
-    }
-
-    // ensure the wasm compiles during proving
-    let (module, canonical_hash, info) =
-        match native::compile_user_wasm(&wasm, version, page_limit, debug) {
-            Ok(result) => result,
-            Err(error) => error!(error),
-        };
-
-    sp.write_slice(out_hash_ptr, canonical_hash.as_slice());
-    sp.write_ptr(heapify(module));
-    sp.write_u16(info.footprint).skip_u16().write_u32(info.size); // wasm info
+    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_u64_raw(gas, *gas_left);
+    sp.write_slice(module_hash, &module.hash().0);
+    sp.write_u16(pages).skip_space();
     sp.write_nullptr();
 }
 
@@ -70,20 +71,19 @@ pub fn compile_user_wasm(env: WasmEnvMut, sp: u32) {
 /// # Go side
 ///
 /// The Go compiler expects the call to take the form
-///     λ(
-///           hash *common.Hash, 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
-///     || hash || 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 compiled_hash = sp.read_bytes32();
+    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();
@@ -94,15 +94,11 @@ pub fn call_user_wasm(env: WasmEnvMut, sp: u32) -> MaybeEscape {
     let pricing = config.pricing;
     let ink = pricing.gas_to_ink(sp.read_u64_raw(gas));
 
-    let module = match &env.data().compiled_modules.get(&compiled_hash) {
-        None => {
-            return Err(Escape::Failure(format!(
-                "compiled hash requested {:?} not found in {:?}",
-                compiled_hash,
-                env.data().compiled_modules.keys()
-            )))
-        }
-        Some(module) => (*module).clone(),
+    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(
@@ -133,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);
 }
 
@@ -155,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_module(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