Add lazy validation compilation mode (#854)

* rename CompilationMode::Lazy -> LazyTranslation

* add CompilationMode::Lazy

This new mode allows for partial Wasm module validation and defers both Wasm translation and validation to first use.

* add support for --compilation-mode in Wasmi CLI

* adjust benchmarks for new compilation modes

* fix internal doc link

crates/cli/src/args.rs

@@ -1,5 +1,5 @@
 use anyhow::{Context, Error, Result};
-use clap::Parser;
+use clap::{Parser, ValueEnum};
 use std::{
     ffi::OsStr,
     net::SocketAddr,
@@ -83,8 +83,8 @@ pub struct Args {
     invoke: Option<String>,
 
     /// Enable lazy Wasm compilation.
-    #[clap(long = "lazy")]
-    lazy: bool,
+    #[clap(long = "compilation-mode", value_enum, default_value_t=CompilationMode::Eager)]
+    compilation_mode: CompilationMode,
 
     /// Enable execution fiel metering with N units of fuel.
     ///
@@ -97,6 +97,25 @@ pub struct Args {
     func_args: Vec<String>,
 }
 
+/// The chosen Wasmi compilation mode.
+#[derive(Debug, Default, Copy, Clone, ValueEnum)]
+enum CompilationMode {
+    #[default]
+    Eager,
+    LazyTranslation,
+    Lazy,
+}
+
+impl From<CompilationMode> for wasmi::CompilationMode {
+    fn from(mode: CompilationMode) -> Self {
+        match mode {
+            CompilationMode::Eager => Self::Eager,
+            CompilationMode::LazyTranslation => Self::LazyTranslation,
+            CompilationMode::Lazy => Self::Lazy,
+        }
+    }
+}
+
 impl Args {
     /// Returns the Wasm file path given to the CLI app.
     pub fn wasm_file(&self) -> &Path {
@@ -119,8 +138,8 @@ impl Args {
     }
 
     /// Returns `true` if lazy Wasm compilation is enabled.
-    pub fn lazy(&self) -> bool {
-        self.lazy
+    pub fn compilation_mode(&self) -> wasmi::CompilationMode {
+        self.compilation_mode.into()
     }
 
     /// Pre-opens all directories given in `--dir` and returns them for use by the [`WasiCtx`].

crates/cli/src/context.rs

@@ -27,19 +27,15 @@ impl Context {
         wasm_file: &Path,
         wasi_ctx: WasiCtx,
         fuel: Option<u64>,
-        lazy: bool,
+        compilation_mode: CompilationMode,
     ) -> Result<Self, Error> {
         let mut config = Config::default();
         config.wasm_tail_call(true);
         config.wasm_extended_const(true);
         if fuel.is_some() {
             config.consume_fuel(true);
         }
-        let mode = match lazy {
-            true => CompilationMode::Lazy,
-            false => CompilationMode::Eager,
-        };
-        config.compilation_mode(mode);
+        config.compilation_mode(compilation_mode);
         let engine = wasmi::Engine::new(&config);
         let wasm_bytes = utils::read_wasm_or_wat(wasm_file)?;
         let module = wasmi::Module::new(&engine, &mut &wasm_bytes[..]).map_err(|error| {

crates/cli/src/main.rs

@@ -20,7 +20,7 @@ fn main() -> Result<()> {
     let args = Args::parse();
     let wasm_file = args.wasm_file();
     let wasi_ctx = args.wasi_context()?;
-    let mut ctx = Context::new(wasm_file, wasi_ctx, args.fuel(), args.lazy())?;
+    let mut ctx = Context::new(wasm_file, wasi_ctx, args.fuel(), args.compilation_mode())?;
     let (func_name, func) = get_invoked_func(&args, &ctx)?;
     let ty = func.ty(ctx.store());
     let func_args = utils::decode_func_args(&ty, args.func_args())?;

crates/wasmi/benches/benches.rs

@@ -10,7 +10,18 @@ use self::bench::{
 use bench::bench_config;
 use core::{slice, time::Duration};
 use criterion::{criterion_group, criterion_main, Bencher, Criterion};
-use wasmi::{core::TrapCode, Engine, Extern, Func, Linker, Memory, Module, Store, Value};
+use wasmi::{
+    core::TrapCode,
+    CompilationMode,
+    Engine,
+    Extern,
+    Func,
+    Linker,
+    Memory,
+    Module,
+    Store,
+    Value,
+};
 use wasmi_core::{Pages, ValueType, F32, F64};
 
 criterion_group!(
@@ -106,14 +117,6 @@ enum Validation {
     Unchecked,
 }
 
-/// How to translate a Wasm module.
-enum CompilationMode {
-    /// Eagerly compiles Wasm function bodies.
-    Eager,
-    /// Lazily compiles Wasm function bodies.
-    Lazy,
-}
-
 fn bench_translate_for(
     c: &mut Criterion,
     name: &str,
@@ -128,6 +131,7 @@ fn bench_translate_for(
     };
     let mode_id = match mode {
         CompilationMode::Eager => "eager",
+        CompilationMode::LazyTranslation => "lazy-translation",
         CompilationMode::Lazy => "lazy",
     };
     let fuel_id = match fuel_metering {
@@ -140,14 +144,7 @@ fn bench_translate_for(
         if matches!(fuel_metering, FuelMetering::Enabled) {
             config.consume_fuel(true);
         }
-        match mode {
-            CompilationMode::Eager => {
-                config.compilation_mode(wasmi::CompilationMode::Eager);
-            }
-            CompilationMode::Lazy => {
-                config.compilation_mode(wasmi::CompilationMode::Lazy);
-            }
-        }
+        config.compilation_mode(mode);
         let create_module = match validation {
             Validation::Checked => {
                 |engine: &Engine, bytes: &[u8]| -> Module { Module::new(engine, bytes).unwrap() }
@@ -187,15 +184,15 @@ fn bench_translate_for_all(c: &mut Criterion, name: &str, path: &str) {
         name,
         path,
         Validation::Checked,
-        CompilationMode::Lazy,
+        CompilationMode::LazyTranslation,
         FuelMetering::Disabled,
     );
     bench_translate_for(
         c,
         name,
         path,
-        Validation::Unchecked,
-        CompilationMode::Eager,
+        Validation::Checked,
+        CompilationMode::Lazy,
         FuelMetering::Disabled,
     );
     bench_translate_for(
@@ -204,14 +201,6 @@ fn bench_translate_for_all(c: &mut Criterion, name: &str, path: &str) {
         path,
         Validation::Unchecked,
         CompilationMode::Eager,
-        FuelMetering::Enabled,
-    );
-    bench_translate_for(
-        c,
-        name,
-        path,
-        Validation::Unchecked,
-        CompilationMode::Lazy,
         FuelMetering::Disabled,
     );
 }

crates/wasmi/src/engine/code_map.rs

@@ -5,17 +5,18 @@
 //! This is the data structure specialized to handle compiled
 //! register machine based bytecode functions.
 
-use super::{FuncTranslationDriver, FuncTranslator};
+use super::{FuncTranslationDriver, FuncTranslator, ValidatingFuncTranslator};
 use crate::{
     core::UntypedValue,
     engine::bytecode::Instruction,
     module::{FuncIdx, ModuleHeader},
     Error,
 };
 use alloc::boxed::Box;
-use core::{mem, ops, slice};
+use core::{fmt, mem, ops, slice};
 use spin::RwLock;
 use wasmi_arena::{Arena, ArenaIndex};
+use wasmparser::{FuncToValidate, ValidatorResources};
 
 /// A reference to a compiled function stored in the [`CodeMap`] of an [`Engine`](crate::Engine).
 #[derive(Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
@@ -95,7 +96,6 @@ impl InternalFuncEntity {
 }
 
 /// An internal uncompiled function entity.
-#[derive(Debug)]
 pub struct UncompiledFuncEntity {
     /// The index of the function within the `module`.
     func_idx: FuncIdx,
@@ -106,6 +106,21 @@ pub struct UncompiledFuncEntity {
     /// This is required for Wasm module related information in order
     /// to compile the Wasm function body.
     module: ModuleHeader,
+    /// Optional Wasm validation information.
+    ///
+    /// This is `Some` if the [`UncompiledFuncEntity`] is to be validated upon compilation.
+    func_to_validate: Option<FuncToValidate<ValidatorResources>>,
+}
+
+impl fmt::Debug for UncompiledFuncEntity {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        f.debug_struct("UncompiledFuncEntity")
+            .field("func_idx", &self.func_idx)
+            .field("bytes", &self.bytes)
+            .field("module", &self.module)
+            .field("validate", &self.func_to_validate.is_some())
+            .finish()
+    }
 }
 
 /// A boxed byte slice that stores up to 30 bytes inline.
@@ -304,6 +319,7 @@ impl CodeMap {
         func: CompiledFunc,
         bytes: &[u8],
         module: &ModuleHeader,
+        func_to_validate: Option<FuncToValidate<ValidatorResources>>,
     ) {
         let Some(func) = self.entities.get_mut(func).map(RwLock::get_mut) else {
             panic!("tried to initialize invalid compiled func: {func:?}")
@@ -315,6 +331,7 @@ impl CodeMap {
             func_idx,
             bytes,
             module,
+            func_to_validate,
         });
     }
 
@@ -372,17 +389,33 @@ impl CodeMap {
                     module.engine()
                 )
             };
-            let allocs = engine.get_translation_allocs();
-            let translator = FuncTranslator::new(func_idx, module, allocs)?;
-            let allocs = FuncTranslationDriver::new(0, &bytes[..], translator)?.translate(
-                |compiled_func| {
-                    *func = InternalFuncEntity::Compiled(compiled_func);
-                },
-            )?;
+            match uncompiled.func_to_validate.take() {
+                Some(func_to_validate) => {
+                    let allocs = engine.get_allocs();
+                    let translator = FuncTranslator::new(func_idx, module, allocs.0)?;
+                    let validator = func_to_validate.into_validator(allocs.1);
+                    let translator = ValidatingFuncTranslator::new(validator, translator)?;
+                    let allocs = FuncTranslationDriver::new(0, &bytes[..], translator)?.translate(
+                        |compiled_func| {
+                            *func = InternalFuncEntity::Compiled(compiled_func);
+                        },
+                    )?;
+                    engine.recycle_allocs(allocs.translation, allocs.validation);
+                }
+                None => {
+                    let allocs = engine.get_translation_allocs();
+                    let translator = FuncTranslator::new(func_idx, module, allocs)?;
+                    let allocs = FuncTranslationDriver::new(0, &bytes[..], translator)?.translate(
+                        |compiled_func| {
+                            *func = InternalFuncEntity::Compiled(compiled_func);
+                        },
+                    )?;
+                    engine.recycle_translation_allocs(allocs);
+                }
+            }
             // TODO: In case translation of `func` fails it is going to be recompiled over and over again
             //       for every threads which might be very costly. A status flag that indicates compilation
             //       failure might be required to fix this.
-            engine.recycle_translation_allocs(allocs);
             // Note: Leaking a read-lock will deny write access to all future accesses.
             //       This is fine since function entities are immutable once compiled.
             let func = spin::RwLockReadGuard::leak(func.downgrade())

crates/wasmi/src/engine/config.rs

@@ -202,7 +202,15 @@ pub enum CompilationMode {
     /// The Wasm code is compiled eagerly to `wasmi` bytecode.
     #[default]
     Eager,
-    /// The Wasm code is compiled lazily on use to `wasmi` bytecode.
+    /// The Wasm code is validated eagerly and translated lazily on first use.
+    LazyTranslation,
+    /// The Wasm code is validated and translated lazily on first use.
+    ///
+    /// # Note
+    ///
+    /// This configuration might be removed in the future since it results in
+    /// partial Wasm module validation which is a controversial topic.
+    /// Read more here: <https://github.com/WebAssembly/design/issues/1464>
     Lazy,
 }
 

crates/wasmi/src/engine/mod.rs

@@ -241,17 +241,18 @@ impl Engine {
                     .translate(|func_entity| self.inner.init_func(compiled_func, func_entity))?;
                 self.inner.recycle_translation_allocs(allocs);
             }
-            (CompilationMode::Lazy, Some(func_to_validate)) => {
+            (CompilationMode::LazyTranslation, Some(func_to_validate)) => {
                 let allocs = self.inner.get_validation_allocs();
-                let translator = LazyFuncTranslator::new(func_index, compiled_func, module);
+                let translator = LazyFuncTranslator::new(func_index, compiled_func, module, None);
                 let validator = func_to_validate.into_validator(allocs);
                 let translator = ValidatingFuncTranslator::new(validator, translator)?;
                 let allocs = FuncTranslationDriver::new(offset, bytes, translator)?
                     .translate(|func_entity| self.inner.init_func(compiled_func, func_entity))?;
                 self.inner.recycle_validation_allocs(allocs.validation);
             }
-            (CompilationMode::Lazy, None) => {
-                let translator = LazyFuncTranslator::new(func_index, compiled_func, module);
+            (CompilationMode::Lazy | CompilationMode::LazyTranslation, func_to_validate) => {
+                let translator =
+                    LazyFuncTranslator::new(func_index, compiled_func, module, func_to_validate);
                 FuncTranslationDriver::new(offset, bytes, translator)?
                     .translate(|func_entity| self.inner.init_func(compiled_func, func_entity))?;
             }
@@ -264,11 +265,25 @@ impl Engine {
         self.inner.get_translation_allocs()
     }
 
+    /// Returns reusable [`FuncTranslatorAllocations`] and [`FuncValidatorAllocations`] from the [`Engine`].
+    pub(crate) fn get_allocs(&self) -> (FuncTranslatorAllocations, FuncValidatorAllocations) {
+        self.inner.get_allocs()
+    }
+
     /// Recycles the given [`FuncTranslatorAllocations`] in the [`Engine`].
     pub(crate) fn recycle_translation_allocs(&self, allocs: FuncTranslatorAllocations) {
         self.inner.recycle_translation_allocs(allocs)
     }
 
+    /// Recycles the given [`FuncTranslatorAllocations`] and [`FuncValidatorAllocations`] in the [`Engine`].
+    pub(crate) fn recycle_allocs(
+        &self,
+        translation: FuncTranslatorAllocations,
+        validation: FuncValidatorAllocations,
+    ) {
+        self.inner.recycle_allocs(translation, validation)
+    }
+
     /// Initializes the uninitialized [`CompiledFunc`] for the [`Engine`].
     ///
     /// # Note
@@ -286,8 +301,10 @@ impl Engine {
         func: CompiledFunc,
         bytes: &[u8],
         module: &ModuleHeader,
+        func_to_validate: Option<FuncToValidate<ValidatorResources>>,
     ) {
-        self.inner.init_lazy_func(func_idx, func, bytes, module)
+        self.inner
+            .init_lazy_func(func_idx, func, bytes, module, func_to_validate)
     }
 
     /// Resolves the [`CompiledFunc`] to the underlying `wasmi` bytecode instructions.
@@ -704,11 +721,12 @@ impl EngineInner {
         func: CompiledFunc,
         bytes: &[u8],
         module: &ModuleHeader,
+        func_to_validate: Option<FuncToValidate<ValidatorResources>>,
     ) {
         self.res
             .write()
             .code_map
-            .init_lazy_func(func_idx, func, bytes, module)
+            .init_lazy_func(func_idx, func, bytes, module, func_to_validate)
     }
 
     /// Resolves the [`InternalFuncEntity`] for [`CompiledFunc`] and applies `f` to it.