From 415a919678e54efd2dcf187a7fda1c2df5a1e112 Mon Sep 17 00:00:00 2001
From: Robin Freyler <robin.freyler@gmail.com>
Date: Sat, 14 Sep 2024 12:26:42 +0200
Subject: [PATCH] Fix invalid local preservation overwrite (#1177)

* properly assert for audit_1_execution test case

* fix invalid local preservation overwrite
---
 .../src/engine/translator/instr_encoder.rs    |  4 +--
 .../src/engine/translator/tests/fuzz/mod.rs   | 35 +++++++++++++++++++
 .../translator/tests/fuzz/wat/audit_2.wat     | 15 ++++++++
 3 files changed, 52 insertions(+), 2 deletions(-)
 create mode 100644 crates/wasmi/src/engine/translator/tests/fuzz/wat/audit_2.wat

diff --git a/crates/wasmi/src/engine/translator/instr_encoder.rs b/crates/wasmi/src/engine/translator/instr_encoder.rs
index ca474fd0d5..ac29784fb6 100644
--- a/crates/wasmi/src/engine/translator/instr_encoder.rs
+++ b/crates/wasmi/src/engine/translator/instr_encoder.rs
@@ -796,10 +796,10 @@ impl InstrEncoder {
         };
         if matches!(
             stack.get_register_space(returned_value),
-            RegisterSpace::Local
+            RegisterSpace::Local | RegisterSpace::Preserve
         ) {
             // Can only apply the optimization if the returned value of `last_instr`
-            // is _NOT_ itself a local register due to observable behavior.
+            // is _NOT_ itself a local register due to observable behavior or already preserved.
             return fallback_case(self, stack, local, value, preserved, fuel_info);
         }
         let Some(last_instr) = self.last_instr else {
diff --git a/crates/wasmi/src/engine/translator/tests/fuzz/mod.rs b/crates/wasmi/src/engine/translator/tests/fuzz/mod.rs
index 99c50b3680..bfaab8fcd3 100644
--- a/crates/wasmi/src/engine/translator/tests/fuzz/mod.rs
+++ b/crates/wasmi/src/engine/translator/tests/fuzz/mod.rs
@@ -7,6 +7,7 @@ use crate::{
         bytecode::{BranchOffset, BranchOffset16, GlobalIdx, RegisterSpan},
         EngineFunc,
     },
+    Val,
 };
 
 #[test]
@@ -510,3 +511,37 @@ fn fuzz_regression_17() {
         ])
         .run()
 }
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn audit_2_codegen() {
+    let wasm = include_str!("wat/audit_2.wat");
+    TranslationTest::from_wat(wasm)
+        .expect_func_instrs([
+            Instruction::copy(2, 0),
+            Instruction::copy(0, 2),
+            Instruction::copy(1, 0),
+            Instruction::return_many(2, 1, 0),
+            Instruction::register(0),
+        ])
+        .run()
+}
+
+#[test]
+#[cfg_attr(miri, ignore)]
+fn audit_2_execution() {
+    use crate::{Engine, Instance, Store};
+    let wat = include_str!("wat/audit_2.wat");
+    let wasm = wat::parse_str(wat).unwrap();
+    let engine = Engine::default();
+    let mut store = <Store<()>>::new(&engine, ());
+    let module = Module::new(&engine, &wasm[..]).unwrap();
+    let instance = Instance::new(&mut store, &module, &[]).unwrap();
+    let func = instance.get_func(&store, "").unwrap();
+    let inputs = [Val::I32(1)];
+    let mut results = [0_i32; 4].map(Val::from);
+    let expected = [1_i32; 4];
+    func.call(&mut store, &inputs[..], &mut results[..])
+        .unwrap();
+    assert_eq!(results.map(|v| v.i32().unwrap()), expected,);
+}
diff --git a/crates/wasmi/src/engine/translator/tests/fuzz/wat/audit_2.wat b/crates/wasmi/src/engine/translator/tests/fuzz/wat/audit_2.wat
new file mode 100644
index 0000000000..9ff97200f6
--- /dev/null
+++ b/crates/wasmi/src/engine/translator/tests/fuzz/wat/audit_2.wat
@@ -0,0 +1,15 @@
+(module ;; different result on main than on Wasmtime
+  (func (export "") (param i32) (result i32 i32 i32 i32)
+    local.get 0
+    local.get 0
+    block (param i32 i32)
+      local.tee 0
+      block (param i32 i32)
+        local.get 0
+        local.get 0
+        br 2 ;; returns
+      end
+    end
+    unreachable
+  )
+)