bf: add fun optimization to fold +- insns

examples/bf.rs

@@ -188,8 +188,9 @@ fn run_jit(prog: &str) {
     let mut label_stack = Vec::new();
 
     // Generate code for each instruction in the bf program.
-    for insn in vm.imem {
-        match insn {
+    let mut pc = 0;
+    while pc < vm.imem.len() {
+        match vm.imem[pc] {
             '>' => {
                 // TODO: generate runtime bounds check.
                 asm.inc(dmem_idx);
@@ -199,10 +200,46 @@ fn run_jit(prog: &str) {
                 asm.dec(dmem_idx);
             }
             '+' => {
-                asm.inc(MemOp8::from(MemOp::IndirectBaseIndex(dmem_base, dmem_idx)));
+                // Apply optimization to fold consecutive '+' instructions to a
+                // single add instruction during compile time.
+
+                match vm.imem[pc..].iter().take_while(|&&i| i.eq(&'+')).count() {
+                    1 => asm.inc(MemOp8::from(MemOp::IndirectBaseIndex(dmem_base, dmem_idx))),
+                    cnt if cnt <= i8::MAX as usize => {
+                        // For add m64, imm8, the immediate is sign-extend and
+                        // hence treated as signed.
+                        asm.add(
+                            MemOp::IndirectBaseIndex(dmem_base, dmem_idx),
+                            Imm8::from(cnt as u8),
+                        );
+
+                        // Advance pc, but account for pc increment at the end
+                        // of the loop.
+                        pc += cnt - 1;
+                    }
+                    cnt @ _ => unimplemented!("cnt={cnt} oob, add with larger imm"),
+                }
             }
             '-' => {
-                asm.dec(MemOp8::from(MemOp::IndirectBaseIndex(dmem_base, dmem_idx)));
+                // Apply optimization to fold consecutive '-' instructions to a
+                // single sub instruction during compile time.
+
+                match vm.imem[pc..].iter().take_while(|&&i| i.eq(&'-')).count() {
+                    1 => asm.dec(MemOp8::from(MemOp::IndirectBaseIndex(dmem_base, dmem_idx))),
+                    cnt if cnt <= i8::MAX as usize => {
+                        // For sub m64, imm8, the immediate is sign-extend and
+                        // hence treated as signed.
+                        asm.sub(
+                            MemOp::IndirectBaseIndex(dmem_base, dmem_idx),
+                            Imm8::from(cnt as u8),
+                        );
+
+                        // Advance pc, but account for pc increment at the end
+                        // of the loop.
+                        pc += cnt - 1;
+                    }
+                    cnt @ _ => unimplemented!("cnt={cnt} oob, sub with larger imm"),
+                }
             }
             '.' => {
                 // Load data memory from active cell into di register, which is
@@ -254,6 +291,9 @@ fn run_jit(prog: &str) {
             }
             _ => unreachable!(),
         }
+
+        // Increment pc to next instruction.
+        pc += 1;
     }
 
     // Return from bf program.

examples/gen.rs

@@ -0,0 +1,45 @@
+//! Add example.
+//!
+//! Jit compile a function at runtime (generate native host code) which calls a function defined in
+//! the example based on the SystemV abi to demonstrate the [`juicebox_asm`] crate.
+
+#[cfg(not(any(target_arch = "x86_64", target_os = "linux")))]
+compile_error!("Only supported on x86_64 with SystemV abi");
+
+use juicebox_asm::insn::*;
+use juicebox_asm::{Asm, Imm16, Imm64, Imm8, Label, Reg16, Reg64::*, Runtime};
+use juicebox_asm::{MemOp, MemOp16, MemOp32, MemOp64, MemOp8};
+
+extern "C" fn add(a: u32, b: u32) -> u32 {
+    a + b
+}
+
+fn main() {
+    let mut asm1 = Asm::new();
+    asm1.xor(r11, r11);
+
+    assert!(0x7fu8 <= i8::MAX as u8);
+
+    asm1.add(MemOp::IndirectBaseIndex(rdi, r11), Imm8::from(0x10u8));
+    asm1.ret();
+
+    let code = asm1.into_code();
+    std::fs::write("jit.asm", &code).unwrap();
+
+    let mut asm2 = Asm::new();
+    asm2.mov(Reg16::cx, Imm16::from(1u16));
+    asm2.mov(MemOp::Indirect(rdi), Reg16::cx);
+    asm2.ret();
+
+    let mut rt = Runtime::new();
+    let fn1 = unsafe { rt.add_code::<extern "C" fn(*mut u64)>(code) };
+    let fn2 = unsafe { rt.add_code::<extern "C" fn(*mut u64)>(asm2.into_code()) };
+
+    let mut data = 0xffff_ffff_ffff_ff00u64;
+    fn1(&mut data as *mut u64);
+    println!("data={:x}", data);
+
+    let mut data = 0xffff_ffff_ffff_ffffu64;
+    fn2(&mut data as *mut u64);
+    println!("data={:x}", data);
+}