Merge copy instructions if possible

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

@@ -323,6 +323,44 @@ impl InstrEncoder {
         self.instrs.push(instr)
     }
 
+    /// Tries to merge `copy(result, value)` if the last instruction is a matching copy instruction.
+    ///
+    /// - Returns `None` if merging of the copy instruction was not possible.
+    /// - Returns the `Instr` of the merged `copy2` instruction if merging was successful.
+    fn merge_copy_instrs(&mut self, result: Register, value: TypedProvider) -> Option<Instr> {
+        let TypedProvider::Register(value) = value else {
+            // Case: cannot merge copies with immediate values at the moment.
+            //
+            // Note: we could implement this but it would require us to allocate
+            //       function local constants which we want to avoid generally.
+            return None;
+        };
+        let Some(last_instr) = self.last_instr else {
+            // There is no last instruction, e.g. when ending a `block`.
+            return None;
+        };
+        let Instruction::Copy {
+            result: last_result,
+            value: last_value,
+        } = *self.instrs.get(last_instr)
+        else {
+            // Case: last instruction was not a copy instruction, so we cannot merge anything.
+            return None;
+        };
+        if !(result == last_result.next() || result == last_result.prev()) {
+            // Case: cannot merge copy instructions as `copy2` since result registers are not contiguous.
+            return None;
+        }
+        let (merged_result, value0, value1) = if last_result < result {
+            (last_result, last_value, value)
+        } else {
+            (result, value, last_value)
+        };
+        let merged_copy = Instruction::copy2(RegisterSpan::new(merged_result), value0, value1);
+        *self.instrs.get_mut(last_instr) = merged_copy;
+        Some(last_instr)
+    }
+
     /// Encode a `copy result <- value` instruction.
     ///
     /// # Note
@@ -345,6 +383,9 @@ impl InstrEncoder {
             let cref = stack.alloc_const(value.into())?;
             Ok(Instruction::copy(result, cref))
         }
+        if let Some(merged_instr) = self.merge_copy_instrs(result, value) {
+            return Ok(Some(merged_instr));
+        }
         let instr = match value {
             TypedProvider::Register(value) => {
                 if result == value {

crates/wasmi/src/engine/translator/tests/op/copy.rs

@@ -0,0 +1,68 @@
+use super::*;
+use crate::engine::bytecode::RegisterSpan;
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn merge_2_copy_instrs_0() {
+    let wasm = r"
+        (module
+            (func (param i32 i32 i32 i32 i32)
+                (local.set 0 (local.get 2)) ;; copy 0 <- 2
+                (local.set 1 (local.get 4)) ;; copy 1 <- 4
+            )
+        )";
+    TranslationTest::from_wat(wasm)
+        .expect_func_instrs([
+            Instruction::copy2(
+                RegisterSpan::new(Register::from_i16(0)),
+                Register::from_i16(2),
+                Register::from_i16(4),
+            ),
+            Instruction::Return,
+        ])
+        .run()
+}
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn merge_2_copy_instrs_1() {
+    let wasm = r"
+        (module
+            (func (param i32 i32 i32 i32 i32)
+                (local.set 0 (local.get 4)) ;; copy 1 <- 4
+                (local.set 1 (local.get 2)) ;; copy 0 <- 2
+            )
+        )";
+    TranslationTest::from_wat(wasm)
+        .expect_func_instrs([
+            Instruction::copy2(
+                RegisterSpan::new(Register::from_i16(0)),
+                Register::from_i16(4),
+                Register::from_i16(2),
+            ),
+            Instruction::Return,
+        ])
+        .run()
+}
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn merge_2_copy_instrs_2() {
+    let wasm = r"
+        (module
+            (func (param i32 i32 i32 i32 i32)
+                (local.set 1 (local.get 4)) ;; copy 1 <- 4
+                (local.set 0 (local.get 2)) ;; copy 0 <- 2
+            )
+        )";
+    TranslationTest::from_wat(wasm)
+        .expect_func_instrs([
+            Instruction::copy2(
+                RegisterSpan::new(Register::from_i16(0)),
+                Register::from_i16(2),
+                Register::from_i16(4),
+            ),
+            Instruction::Return,
+        ])
+        .run()
+}

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

@@ -6,6 +6,7 @@ mod br_table;
 mod call;
 mod cmp;
 mod cmp_br;
+mod copy;
 mod global_get;
 mod global_set;
 mod i32_eqz;

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

@@ -170,6 +170,7 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
             &self.alloc.buffer.providers[..],
             fuel_info,
         )?;
+        self.alloc.instr_encoder.reset_last_instr();
         // Create loop header label and immediately pin it.
         let stack_height = BlockHeight::new(self.engine(), self.alloc.stack.height(), block_type)?;
         let header = self.alloc.instr_encoder.new_label();
@@ -811,7 +812,6 @@ impl<'a> VisitOperator<'a> for FuncTranslator {
     fn visit_local_get(&mut self, local_index: u32) -> Self::Output {
         bail_unreachable!(self);
         self.alloc.stack.push_local(local_index)?;
-        self.alloc.instr_encoder.reset_last_instr();
         Ok(())
     }