Allow local.set optimization with active preservation (#792)

* allow local.set optimization with active preservation

* apply rustfmt

crates/wasmi/src/engine/func_builder/inst_builder.rs

@@ -48,6 +48,15 @@ impl Instr {
     pub fn into_u32(self) -> u32 {
         self.0
     }
+
+    /// Returns the absolute distance between `self` and `other`.
+    ///
+    /// - Returns `0` if `self == other`.
+    /// - Returns `1` if `self` is adjacent to `other` in the sequence of instructions.
+    /// - etc..
+    pub fn distance(self, other: Self) -> u32 {
+        self.0.abs_diff(other.0)
+    }
 }
 
 /// The relative depth of a Wasm branching target.

crates/wasmi/src/engine/regmach/tests/op/local_set.rs

@@ -267,18 +267,48 @@ fn preserve_result_1() {
     );
     TranslationTest::new(wasm)
         .expect_func_instrs([
+            Instruction::copy(Register::from_i16(3), Register::from_i16(0)),
             Instruction::i32_add(
-                Register::from_i16(2),
+                Register::from_i16(0),
                 Register::from_i16(0),
                 Register::from_i16(1),
             ),
-            Instruction::copy(Register::from_i16(3), Register::from_i16(0)),
-            Instruction::copy(Register::from_i16(0), Register::from_i16(2)),
             Instruction::return_reg(Register::from_i16(3)),
         ])
         .run()
 }
 
+#[test]
+#[cfg_attr(miri, ignore)]
+fn preserve_result_2() {
+    let wasm = wat2wasm(
+        r#"
+        (module
+            (func (param $lhs i32) (param $rhs i32) (result i32 i32)
+                (local.get 0)
+                (local.get 0)
+                (local.set 0
+                    (i32.add
+                        (local.get $lhs)
+                        (local.get $rhs)
+                    )
+                )
+            )
+        )"#,
+    );
+    TranslationTest::new(wasm)
+        .expect_func_instrs([
+            Instruction::copy(Register::from_i16(3), Register::from_i16(0)),
+            Instruction::i32_add(
+                Register::from_i16(0),
+                Register::from_i16(0),
+                Register::from_i16(1),
+            ),
+            Instruction::return_reg2(3, 3),
+        ])
+        .run()
+}
+
 #[test]
 #[cfg_attr(miri, ignore)]
 fn preserve_multiple_0() {

crates/wasmi/src/engine/regmach/translator/instr_encoder.rs

@@ -77,6 +77,35 @@ impl InstrSequence {
         Ok(instr)
     }
 
+    /// Pushes an [`Instruction`] before the [`Instruction`] at [`Instr`].
+    ///
+    /// Returns the [`Instr`] of the [`Instruction`] that was at [`Instr`] before this operation.
+    ///
+    /// # Note
+    ///
+    /// - This operation might be costly. Callers are advised to only insert
+    ///   instructions near the end of the sequence in order to avoid massive
+    ///   copy overhead since all following instructions are required to be
+    ///   shifted in memory.
+    /// - The `instr` will refer to the inserted [`Instruction`] after this operation.
+    ///
+    /// # Errors
+    ///
+    /// If there are too many instructions in the instruction sequence.
+    fn push_before(
+        &mut self,
+        instr: Instr,
+        instruction: Instruction,
+    ) -> Result<Instr, TranslationError> {
+        self.instrs.insert(instr.into_usize(), instruction);
+        let shifted_instr = instr
+            .into_u32()
+            .checked_add(1)
+            .map(Instr::from_u32)
+            .unwrap_or_else(|| panic!("pushed to many instructions to a single function"));
+        Ok(shifted_instr)
+    }
+
     /// Returns the [`Instruction`] associated to the [`Instr`] for this [`InstrSequence`].
     ///
     /// # Panics
@@ -619,46 +648,94 @@ impl InstrEncoder {
 
     /// Encode a `local.set` or `local.tee` instruction.
     ///
-    /// # Note
-    ///
     /// This also applies an optimization in that the previous instruction
     /// result is replaced with the `local` [`Register`] instead of encoding
     /// another `copy` instruction if the `local.set` or `local.tee` belongs
     /// to the same basic block.
-    pub fn encode_local_set(
+    ///
+    /// # Note
+    ///
+    /// - If `value` is a [`Register`] it usually is equal to the
+    ///   result [`Register`] of the previous instruction.
+    pub fn encode_local_set_v2(
         &mut self,
         stack: &mut ValueStack,
         res: &ModuleResources,
         local: Register,
-        value: Register,
+        value: TypedProvider,
+        preserved: Option<Register>,
     ) -> Result<(), TranslationError> {
-        /// Fallback for when we need to encode a `copy` instruction to encode the `local.set` or `local.tee`.
-        fn fallback_copy(
+        fn fallback_case(
             this: &mut InstrEncoder,
+            stack: &mut ValueStack,
             local: Register,
-            value: Register,
+            value: TypedProvider,
+            preserved: Option<Register>,
         ) -> Result<(), TranslationError> {
-            this.push_instr(Instruction::copy(local, value))?;
+            if let Some(preserved) = preserved {
+                let preserve_instr = this.push_instr(Instruction::copy(preserved, local))?;
+                this.notify_preserved_register(preserve_instr);
+            }
+            this.encode_copy(stack, local, value)?;
             Ok(())
         }
+
+        debug_assert!(matches!(
+            stack.get_register_space(local),
+            RegisterSpace::Local
+        ));
+        let TypedProvider::Register(returned_value) = value else {
+            // Cannot apply the optimization for `local.set C` where `C` is a constant value.
+            return fallback_case(self, stack, local, value, preserved);
+        };
+        if matches!(
+            stack.get_register_space(returned_value),
+            RegisterSpace::Local
+        ) {
+            // Can only apply the optimization if the returned value of `last_instr`
+            // is _NOT_ itself a local register due to observable behavior.
+            return fallback_case(self, stack, local, value, preserved);
+        }
         let Some(last_instr) = self.last_instr else {
-            return fallback_copy(self, local, value);
+            // Can only apply the optimization if there is a previous instruction
+            // to replace its result register instead of emitting a copy.
+            return fallback_case(self, stack, local, value, preserved);
         };
+        if preserved.is_some() && last_instr.distance(self.instrs.next_instr()) >= 4 {
+            // We avoid applying the optimization if the last instruction
+            // has a very large encoding, e.g. for function calls with lots
+            // of parameters. This is because the optimization while also
+            // preserving a local register requires costly shifting all
+            // instruction words of the last instruction.
+            // Thankfully most instructions are small enough.
+            return fallback_case(self, stack, local, value, preserved);
+        }
         let Some(result) = self.instrs.get_mut(last_instr).result_mut(res) else {
-            return fallback_copy(self, local, value);
+            // Can only apply the optimization if the last instruction has exactly one result.
+            return fallback_case(self, stack, local, value, preserved);
         };
-        if matches!(stack.get_register_space(*result), RegisterSpace::Local) {
-            return fallback_copy(self, local, value);
-        }
-        if *result != value {
+        if *result != returned_value {
+            // We only want to apply the optimization if indeed `x` in `local.set n x`
+            // is the same register as the result register of the last instruction.
+            //
             // TODO: Find out in what cases `result != value`. Is this a bug or an edge case?
             //       Generally `result` should be equal to `value` since `value` refers to the
             //       `result` of the previous instruction.
             //       Therefore, instead of an `if` we originally had a `debug_assert`.
             //       (Note: the spidermonkey bench test failed without this change.)
-            return fallback_copy(self, local, value);
+            return fallback_case(self, stack, local, value, preserved);
         }
         *result = local;
+        if let Some(preserved) = preserved {
+            // We were able to apply the optimization.
+            // Preservation requires the copy to be before the optimized last instruction.
+            // Therefore we need to push the preservation `copy` instruction before it.
+            let shifted_last_instr = self
+                .instrs
+                .push_before(last_instr, Instruction::copy(preserved, local))?;
+            self.notify_preserved_register(last_instr);
+            self.last_instr = Some(shifted_last_instr);
+        }
         Ok(())
     }
 

crates/wasmi/src/engine/regmach/translator/visit.rs

@@ -803,53 +803,15 @@ impl<'a> VisitOperator<'a> for FuncTranslator<'a> {
     fn visit_local_set(&mut self, local_index: u32) -> Self::Output {
         bail_unreachable!(self);
         let value = self.alloc.stack.pop();
-        let local_register = Register::try_from(local_index)?;
-        if let Some(register) = self.alloc.stack.preserve_locals(local_index)? {
-            // Case: we need to preserve the `local.get` on the value stack.
-            //
-            // Unfortunately this prevents us from applying the optimization of
-            // switching out the result register of the previous instruction because
-            // we would need to encode a copy in before already encoded
-            // instructions to preserve functional equivalence.
-            //
-            // It is possible to do so, however, a bit complicated. So for now we don't.
-            let preserve_instr = self
-                .alloc
-                .instr_encoder
-                .push_instr(Instruction::copy(register, local_register))?;
-            self.alloc
-                .instr_encoder
-                .notify_preserved_register(preserve_instr);
-            self.alloc
-                .instr_encoder
-                .encode_copy(&mut self.alloc.stack, local_register, value)?;
-            return Ok(());
-        }
-        match value {
-            TypedProvider::Const(_) => {
-                // Case: we set the local variable to a constant value.
-                self.alloc.instr_encoder.encode_copy(
-                    &mut self.alloc.stack,
-                    local_register,
-                    value,
-                )?;
-            }
-            TypedProvider::Register(value) => {
-                // Case: we set the local variable to a register value.
-                //
-                // It could be that the register value is the result of a previous
-                // computation which allows us to exchange the result register of
-                // this previous instruction instead of encoding another `copy`
-                // instruction as an optimization.
-                self.alloc.instr_encoder.encode_local_set(
-                    &mut self.alloc.stack,
-                    &self.res,
-                    local_register,
-                    value,
-                )?;
-            }
-        }
-        self.alloc.instr_encoder.reset_last_instr();
+        let local = Register::try_from(local_index)?;
+        let preserved = self.alloc.stack.preserve_locals(local_index)?;
+        self.alloc.instr_encoder.encode_local_set_v2(
+            &mut self.alloc.stack,
+            &self.res,
+            local,
+            value,
+            preserved,
+        )?;
         Ok(())
     }