Fix a bunch of register-machine wasmi translation bugs (#838)

* add failing translation tests

* add debug_assert

* remove NonZeroUsize runtime check

* improve docs

* fix naming of preservation registers and space

Previously we called this the "storage" space but now we renamed all occurrences to "preservation". This should improve code readability.

* more renamings from storage to preserve

* make some methods private if possible

* cleanup push_local method

* add debug_assert to bump_preserved

* rename storage -> preserve

* improve lifetime tracking of preserved registers

The new systems starts preserved register amounts at 2 instead of 1. This prevents removal of the slot when popping it from the preservation stack. In order to properly recycle preservation registers again, we now check for all previously removed preservation registers if they are truly removed (amount = 1) and remove them before allocating a new preservation register.

* add debug_assert to push_preserved

* add dev docs for new semantics

* add else provider regression test and fix bug

* apply rustfmt

* remove unneeded validation checks

* add dev comment

* add missing call to dec_register_usage

* fix intra doc link

* apply rustfmt

* add another if test with missing else

* simplify fuzz_6.wat test case

* finalize testcases 5 and 6

* fix test cases 5 and 6

crates/wasmi/src/engine/bytecode/provider.rs

@@ -37,6 +37,16 @@ pub enum Provider<T> {
     Const(T),
 }
 
+impl<T> Provider<T> {
+    /// Returns `Some` if `self` is a [`Provider::Register`].
+    pub fn into_register(self) -> Option<Register> {
+        match self {
+            Self::Register(register) => Some(register),
+            Self::Const(_) => None,
+        }
+    }
+}
+
 /// An untyped [`Provider`].
 ///
 /// # Note

crates/wasmi/src/engine/translator/mod.rs

@@ -900,24 +900,28 @@ impl<'parser> FuncTranslator<'parser> {
                 .expect("must have `else` label since `else` is reachable"),
         );
         let if_height = frame.block_height().into_u16() as usize;
+        let else_providers = self.alloc.control_stack.pop_else_providers();
         if has_results {
             // We haven't visited the `else` block and thus the `else`
             // providers are still on the auxiliary stack and need to
             // be popped. We use them to restore the stack to the state
             // when entering the `if` block so that we can properly copy
             // the `else` results to were they are expected.
             self.alloc.stack.trunc(if_height);
-            for provider in self.alloc.control_stack.pop_else_providers() {
+            for provider in else_providers {
                 self.alloc.stack.push_provider(provider)?;
+                if let TypedProvider::Register(register) = provider {
+                    self.alloc.stack.dec_register_usage(register);
+                }
             }
-            self.translate_copy_branch_params(frame.branch_params(self.engine()))?;
+            self.translate_copy_branch_params(frame.branch_params(self.res.engine()))?;
         }
         // After `else` parameters have been copied we can finally pin the `end` label.
         self.alloc.instr_encoder.pin_label(frame.end_label());
         // Without `else` block the code after the `if` is always reachable and
         // thus we need to clean up and prepare the value stack for the following code.
         self.alloc.stack.trunc(if_height);
-        for result in frame.branch_params(self.engine()) {
+        for result in frame.branch_params(self.res.engine()) {
             self.alloc.stack.push_register(result)?;
         }
         self.reachable = true;

crates/wasmi/src/engine/translator/stack/mod.rs

@@ -49,7 +49,7 @@ impl From<TaggedProvider> for TypedProvider {
         match provider {
             TaggedProvider::Local(register)
             | TaggedProvider::Dynamic(register)
-            | TaggedProvider::Storage(register)
+            | TaggedProvider::Preserved(register)
             | TaggedProvider::ConstLocal(register) => Self::Register(register),
             TaggedProvider::ConstValue(value) => Self::Const(value),
         }
@@ -218,12 +218,13 @@ impl ValueStack {
                 self.providers.push_dynamic(reg);
                 return Ok(());
             }
-            RegisterSpace::Storage => {
+            RegisterSpace::Preserve => {
                 // Note: we currently do not call `self.reg_alloc.push_storage()`
                 //       since that API would push always another register on the preservation
                 //       stack instead of trying to bump the amount of already existing
                 //       preservation slots for the same register if possible.
-                self.providers.push_storage(reg);
+                self.reg_alloc.bump_preserved(reg);
+                self.providers.push_preserved(reg);
             }
             RegisterSpace::Local => {
                 self.providers.push_local(reg);
@@ -241,11 +242,11 @@ impl ValueStack {
     ///
     /// If too many registers have been registered.
     pub fn push_local(&mut self, local_index: u32) -> Result<Register, TranslationError> {
-        let index = i16::try_from(local_index)
+        let reg = i16::try_from(local_index)
             .ok()
-            .filter(|&value| value as u16 <= self.reg_alloc.len_locals())
+            .map(Register::from_i16)
+            .filter(|reg| self.reg_alloc.is_local(*reg))
             .ok_or_else(|| TranslationError::new(TranslationErrorInner::RegisterOutOfBounds))?;
-        let reg = Register::from_i16(index);
         self.providers.push_local(reg);
         Ok(reg)
     }
@@ -390,4 +391,24 @@ impl ValueStack {
     pub fn get_register_space(&self, register: Register) -> RegisterSpace {
         self.reg_alloc.register_space(register)
     }
+
+    /// Increase preservation [`Register`] usage.
+    ///
+    /// # Note
+    ///
+    /// - This is mainly used to extend the lifetime of `else` providers on the stack.
+    /// - This does nothing if `register` is not a preservation [`Register`].
+    pub fn inc_register_usage(&mut self, register: Register) {
+        self.reg_alloc.inc_register_usage(register)
+    }
+
+    /// Decrease preservation [`Register`] usage.
+    ///
+    /// # Note
+    ///
+    /// - This is mainly used to shorten the lifetime of `else` providers on the stack.
+    /// - This does nothing if `register` is not a preservation [`Register`].
+    pub fn dec_register_usage(&mut self, register: Register) {
+        self.reg_alloc.dec_register_usage(register)
+    }
 }

crates/wasmi/src/engine/translator/stack/provider.rs

@@ -19,8 +19,8 @@ pub enum TaggedProvider {
     Local(Register),
     /// A register referring to a dynamically allocated register.
     Dynamic(Register),
-    /// A register referring to a storage allocated register.
-    Storage(Register),
+    /// A register referring to a preservation allocated register.
+    Preserved(Register),
     /// An untyped constant value.
     ConstValue(TypedValue),
 }
@@ -45,10 +45,10 @@ impl ProviderStack {
         self.locals.reset();
     }
 
-    /// Preserves `local.get` on the [`ProviderStack`] by shifting to storage space.
+    /// Preserves `local.get` on the [`ProviderStack`] by shifting to the preservation space.
     ///
     /// In case there are `local.get n` with `n == preserve_index` on the [`ProviderStack`]
-    /// there is a [`Register`] on the storage space allocated for them. The [`Register`]
+    /// there is a [`Register`] on the preservation space allocated for them. The [`Register`]
     /// allocated this way is returned. Otherwise `None` is returned.
     pub fn preserve_locals(
         &mut self,
@@ -64,16 +64,16 @@ impl ProviderStack {
             debug_assert!(matches!(provider, TaggedProvider::Local(_)));
             let preserved_register = match preserved {
                 Some(register) => {
-                    reg_alloc.bump_storage(register);
+                    reg_alloc.bump_preserved(register);
                     register
                 }
                 None => {
-                    let register = reg_alloc.push_storage()?;
+                    let register = reg_alloc.push_preserved()?;
                     preserved = Some(register);
                     register
                 }
             };
-            *provider = TaggedProvider::Storage(preserved_register);
+            *provider = TaggedProvider::Preserved(preserved_register);
         }
         Ok(preserved)
     }
@@ -112,10 +112,10 @@ impl ProviderStack {
         self.push(TaggedProvider::Dynamic(reg));
     }
 
-    /// Pushes a storage allocated [`Register`] to the [`ProviderStack`].
-    pub fn push_storage(&mut self, reg: Register) {
+    /// Pushes a preservation allocated [`Register`] to the [`ProviderStack`].
+    pub fn push_preserved(&mut self, reg: Register) {
         debug_assert!(!reg.is_const());
-        self.push(TaggedProvider::Storage(reg));
+        self.push(TaggedProvider::Preserved(reg));
     }
 
     /// Pushes a [`Register`] to the [`ProviderStack`] referring to a function parameter or local variable.