Fix lazy translation when running into out of fuel (#1648)

* extend fuel_metering.rs integration tests

* add FuncEntity::set_uncompiled method

This allows to go from Compiling state back to Uncompiled state.

* rename bindings and no longer use mem::take

* rename param: entity -> uncompiled

* differentiate failed lazy compilation from out-of-fuel

Author: Robbepop

crates/wasmi/src/engine/code_map.rs

@@ -15,6 +15,7 @@ use crate::{
     module::{FuncIdx, ModuleHeader},
     Config,
     Error,
+    TrapCode,
 };
 use alloc::boxed::Box;
 use core::{
@@ -366,11 +367,11 @@ impl CodeMap {
         &'a self,
         fuel: Option<&mut Fuel>,
         func: EngineFunc,
-        mut entity: UncompiledFuncEntity,
+        mut uncompiled: UncompiledFuncEntity,
     ) -> Result<CompiledFuncRef<'a>, Error> {
         // Note: it is important that compilation happens without locking the `CodeMap`
         //       since compilation can take a prolonged time.
-        let compiled_func = entity.compile(fuel, &self.features);
+        let compiled_func = uncompiled.compile(fuel, &self.features);
         let mut funcs = self.funcs.lock();
         let Some(entity) = funcs.get_mut(func) else {
             panic!("encountered invalid internal function: {func:?}")
@@ -380,6 +381,10 @@ impl CodeMap {
                 let cref = entity.set_compiled(compiled_func);
                 Ok(self.adjust_cref_lifetime(cref))
             }
+            Err(error) if error.as_trap_code() == Some(TrapCode::OutOfFuel) => {
+                entity.set_uncompiled(uncompiled);
+                Err(error)
+            }
             Err(error) => {
                 entity.set_failed_to_compile();
                 Err(error)
@@ -502,6 +507,28 @@ impl FuncEntity {
         }
     }
 
+    /// Sets the state back to [`UncompiledFuncEntity`] if possible.
+    ///
+    /// # Panics
+    ///
+    /// If the current state was not [`FuncEntity::Compiling`].
+    #[inline]
+    pub fn set_uncompiled(&mut self, uncompiled: UncompiledFuncEntity) {
+        match mem::replace(self, FuncEntity::Uncompiled(uncompiled)) {
+            Self::Compiling => {}
+            unexpected => {
+                // Safety: we just asserted that `self` must be an uncompiled function
+                //         since otherwise we would have returned `None` above.
+                //         Since this is a performance critical path we need to leave out this check.
+                unsafe {
+                    unreachable_unchecked!(
+                        "can only set `Compiling` back to `UncompiledFuncEntity` but found: {unexpected:?}"
+                    )
+                }
+            }
+        }
+    }
+
     /// Sets the [`FuncEntity`] as [`CompiledFuncEntity`].
     ///
     /// Returns a [`CompiledFuncRef`] to the [`CompiledFuncEntity`].
@@ -617,10 +644,10 @@ impl UncompiledFuncEntity {
             VALIDATE_FUEL_PER_BYTE + COMPILE_FUEL_PER_BYTE;
 
         let func_idx = self.func_index;
-        let bytes = mem::take(&mut self.bytes);
+        let wasm_bytes = self.bytes.as_slice();
         let needs_validation = self.validation.is_some();
         let compilation_fuel = |_costs: &FuelCostsProvider| {
-            let len_bytes = bytes.as_slice().len() as u64;
+            let len_bytes = wasm_bytes.len() as u64;
             let fuel_per_byte = match needs_validation {
                 false => COMPILE_FUEL_PER_BYTE,
                 true => VALIDATE_AND_COMPILE_FUEL_PER_BYTE,
@@ -655,7 +682,7 @@ impl UncompiledFuncEntity {
                 };
                 let validator = func_to_validate.into_validator(allocs.1);
                 let translator = ValidatingFuncTranslator::new(validator, translator)?;
-                let allocs = FuncTranslationDriver::new(0, &bytes[..], translator)?.translate(
+                let allocs = FuncTranslationDriver::new(0, wasm_bytes, translator)?.translate(
                     |compiled_func| {
                         result.write(compiled_func);
                     },
@@ -665,7 +692,7 @@ impl UncompiledFuncEntity {
             None => {
                 let allocs = engine.get_translation_allocs();
                 let translator = FuncTranslator::new(func_idx, module, allocs)?;
-                let allocs = FuncTranslationDriver::new(0, &bytes[..], translator)?.translate(
+                let allocs = FuncTranslationDriver::new(0, wasm_bytes, translator)?.translate(
                     |compiled_func| {
                         result.write(compiled_func);
                     },

crates/wasmi/tests/integration/fuel_metering.rs

@@ -1,13 +1,13 @@
 //! Tests to check if wasmi's fuel metering works as intended.
 
 use std::fmt::Debug;
-use wasmi::{Config, Engine, Error, Func, Linker, Module, Store, TrapCode};
+use wasmi::{CompilationMode, Config, Engine, Error, Func, Linker, Module, Store, TrapCode};
 
 /// Setup [`Engine`] and [`Store`] for fuel metering.
-fn test_setup() -> (Store<()>, Linker<()>) {
+fn test_setup(mode: CompilationMode) -> (Store<()>, Linker<()>) {
     let mut config = Config::default();
     config.consume_fuel(true);
-    config.compilation_mode(wasmi::CompilationMode::Eager);
+    config.compilation_mode(mode);
     let engine = Engine::new(&config);
     let store = Store::new(&engine, ());
     let linker = Linker::new(&engine);
@@ -24,8 +24,8 @@ fn create_module(store: &Store<()>, bytes: &[u8]) -> Module {
 }
 
 /// Setup [`Store`] and [`Instance`] for fuel metering.
-fn default_test_setup(wasm: &[u8]) -> (Store<()>, Func) {
-    let (mut store, linker) = test_setup();
+fn default_test_setup(mode: CompilationMode, wasm: &[u8]) -> (Store<()>, Func) {
+    let (mut store, linker) = test_setup(mode);
     let module = create_module(&store, wasm);
     let instance = linker.instantiate_and_start(&mut store, &module).unwrap();
     let func = instance.get_func(&store, "test").unwrap();
@@ -38,37 +38,45 @@ fn default_test_setup(wasm: &[u8]) -> (Store<()>, Func) {
 ///
 /// We just check if the call succeeded, not if the results are correct.
 /// That is to be determined by another kind of test.
+#[track_caller]
 fn assert_success<T>(call_result: Result<T, Error>)
 where
     T: Debug,
 {
-    assert!(call_result.is_ok());
+    if let Err(error) = call_result {
+        panic!("expected `Ok` but got: {error}")
+    }
 }
 
 /// Asserts that the call trapped with [`TrapCode::OutOfFuel`].
+#[track_caller]
 fn assert_out_of_fuel<T>(call_result: Result<T, Error>)
 where
     T: Debug,
 {
-    assert!(matches!(
-        call_result.unwrap_err().as_trap_code(),
-        Some(TrapCode::OutOfFuel)
-    ));
+    let Err(error) = call_result else {
+        panic!("expected `out of fuel` error but got `Ok`")
+    };
+    assert_eq!(
+        error.as_trap_code(),
+        Some(TrapCode::OutOfFuel),
+        "expected `out of fuel` but got: {error}"
+    );
 }
 
-#[test]
-fn metered_i32_add() {
-    let wasm = r#"
-        (module
-            (func (export "test") (param $a i32) (param $b i32) (result i32)
-                (i32.add
-                    (local.get $a)
-                    (local.get $b)
-                )
+const WASM_INPUT: &str = r#"
+    (module
+        (func (export "test") (param $a i32) (param $b i32) (result i32)
+            (i32.add
+                (local.get $a)
+                (local.get $b)
             )
         )
-    "#;
-    let (mut store, func) = default_test_setup(wasm.as_bytes());
+    )
+"#;
+
+fn run_test(mode: CompilationMode, final_fuel: u64) {
+    let (mut store, func) = default_test_setup(mode, WASM_INPUT.as_bytes());
     let func = func.typed::<(i32, i32), i32>(&store).unwrap();
     // No fuel -> no success.
     assert_out_of_fuel(func.call(&mut store, (1, 2)));
@@ -78,7 +86,22 @@ fn metered_i32_add() {
     assert_out_of_fuel(func.call(&mut store, (1, 2)));
     assert_eq!(store.get_fuel().ok(), Some(1));
     // Now add enough fuel, so execution should succeed.
-    store.set_fuel(10).unwrap();
+    store.set_fuel(100).unwrap();
     assert_success(func.call(&mut store, (1, 2)));
-    assert_eq!(store.get_fuel().ok(), Some(7));
+    assert_eq!(store.get_fuel().ok(), Some(final_fuel));
+}
+
+#[test]
+fn metered_i32_add_eager() {
+    run_test(CompilationMode::Eager, 97)
+}
+
+#[test]
+fn metered_i32_add_lazy_translation() {
+    run_test(CompilationMode::LazyTranslation, 48)
+}
+
+#[test]
+fn metered_i32_add_lazy() {
+    run_test(CompilationMode::Lazy, 34)
 }