Add reproduction for block optimization bug

src/main.rs

@@ -105,6 +105,59 @@ fn inv_middle() -> PyResult<()> {
     Ok(())
 }
 
+fn block_optimization() -> PyResult<()> {
+    let mut vm: Icicle = Icicle::new(
+        "x86_64".to_string(), // architecture
+        true,                 // jit
+        true,                 // jit_mem
+        true,                 // shadow_stack
+        true,                 // recompilation
+        false,                // track_uninitialized
+        true,                 // optimize_instructions
+        true,                 // optimize_block
+        false,                // tracing
+    )?;
+
+    // Memory setup
+    let addr: u64 = 0x140001A73;
+    let heap: u64 = 0x71000;
+    vm.mem_map(heap, 0x1000, MemoryProtection::ReadWrite)?;
+    vm.mem_write(heap + 4, b"\x37\x13\x00\x00".to_vec())?;
+    vm.mem_map(addr & !0xFFF, 0x1000, MemoryProtection::ExecuteRead)?;
+    vm.mem_write(addr, b"\x41\xc1\xea\x07\x41\x83\xe2\x1f\x74\x08\x44\x89\xd0\x48\x89\x54\xc6\x08\x49\x83\xc1\x04\x4c\x89\x0e\x4c\x89\xc9\x44\x8b\x11\x44\x89\xd0\xf7\xd0\x49\x89\xc9\xa8\x03\x0f\x84\x88\xf6\xff\xff\xeb\x4c".to_vec())?;
+
+    // Register setup
+    vm.reg_write("r9", heap)?;
+    vm.reg_write("r10", 0x13)?;
+    vm.reg_write("rip", addr)?;
+    vm.reg_write("rsi", heap + 0x100)?;
+
+    // Step through instructions
+    for i in 0..11 {
+        let rip = vm.reg_read("rip")?;
+        let rcx = vm.reg_read("rcx")?;
+        let r9 = vm.reg_read("r9")?;
+
+        println!("[{}] RIP: {:#x}, RCX: {:#x}, R9: {:#x}", i, rip, rcx, r9);
+
+        if rip == 0x140001A8F {
+            vm.reg_write("r9", 0x13370900)?;
+        }
+
+        vm.step(1);
+
+        if rip == 0x140001A9A {
+            if rcx != r9 {
+                println!("[BUG] expected rcx({:#x}) == r9({:#x})", rcx, r9);
+            } else {
+                println!("Everything works!");
+            }
+        }
+    }
+
+    Ok(())
+}
+
 fn main() {
     #![allow(unused_must_use)]
     println!("=== NX (block start) ===");
@@ -115,4 +168,6 @@ fn main() {
     inv_start();
     println!("=== Invalid instruction (block middle) ===");
     inv_middle();
+    println!("=== Block optimization bug ===");
+    block_optimization();
 }

tests/blockopt.py

@@ -0,0 +1,36 @@
+from icicle import *
+
+def main():
+    emu = Icicle("x86_64", jit=True, optimize_block=True, tracing=True)
+
+    instructions = bytes.fromhex("41 C1 EA 07 41 83 E2 1F 74 08 44 89 D0 48 89 54 C6 08 49 83 C1 04 4C 89 0E 4C 89 C9 44 8B 11 44 89 D0 F7 D0 49 89 C9 A8 03 0F 84 88 F6 FF FF EB 4C".replace(" ", ""))
+
+    addr = 0x140001A73
+    heap = 0x71000
+    emu.mem_map(heap, 0x1000, MemoryProtection.ReadWrite)
+    emu.mem_write(heap + 4, 0x1337.to_bytes(4, "little"))
+    emu.mem_map(addr & ~0xFFF, 0x1000, MemoryProtection.ExecuteRead)
+    emu.mem_write(addr, instructions)
+    emu.reg_write("r9", heap)
+    emu.reg_write("r10", 0x13)
+    emu.reg_write("rip", addr)
+    emu.reg_write("rsi", heap + 0x100)
+
+    for i in range(11):
+        rip = emu.reg_read("rip")
+        rcx = emu.reg_read("rcx")
+        r9 = emu.reg_read("r9")
+        print(f"[{i}] RIP: {hex(rip)}, RCX: {hex(rcx)}, R9: {hex(r9)}")
+
+        if rip == 0x140001A8F:
+            emu.reg_write("r9", 0x13370900)
+
+        emu.step(1)
+
+        if rip == 0x140001A9A:
+            assert rcx == r9, f"expected rcx({hex(rcx)}) == r9({hex(r9)})"
+
+    print("Everything works!")
+
+if __name__ == "__main__":
+    main()