examples: add new example demonstrating a simple VM with a jit compiler

1 files changed, 697 insertions, 0 deletions

diff --git a/examples/tiny_vm.rs b/examples/tiny_vm.rs
new file mode 100644
index 0000000..51a8d44
--- /dev/null
+++ b/examples/tiny_vm.rs
@@ -0,0 +1,697 @@
+//! TinyVm example.
+//!
+//! This example introduces as simple 16 bit virtual machine the [`TinyVm`]. The VM consits of
+//! three registers defined in [`TinyReg`], a separate _data_ and _insutrction_ memory and a small
+//! set of instructions [`TinyInsn`], sufficient to implement a guest program to compute the
+//! fiibonacci sequence.
+//!
+//! The `TinyVm` implements a simple _just-in-time (JIT)_ compiler to demonstrate the
+//! [`juicebox_asm`] crate. Additionally, it implements a reference _interpreter_.
+//!
+//! ```
+//! fn main() {
+//!   let mut prog = Vec::new();
+//!   prog.push(TinyInsn::LoadImm(TinyReg::A, 100));
+//!   prog.push(TinyInsn::Add(TinyReg::B, TinyReg::A));
+//!   prog.push(TinyInsn::Addi(TinyReg::C, 100));
+//!   prog.push(TinyInsn::Halt);
+//!
+//!   let mut vm = TinyVm::new(prog);
+//!   vm.interp();
+//!
+//!   assert_eq!(100, vm.read_reg(TinyReg::A));
+//!   assert_eq!(100, vm.read_reg(TinyReg::B));
+//!   assert_eq!(100, vm.read_reg(TinyReg::C));
+//!   assert_eq!(4, vm.icnt);
+//!   assert_eq!(4, vm.pc);
+//!
+//!   vm.pc = 0;
+//!   vm.jit();
+//!
+//!   assert_eq!(100, vm.read_reg(TinyReg::A));
+//!   assert_eq!(200, vm.read_reg(TinyReg::B));
+//!   assert_eq!(200, vm.read_reg(TinyReg::C));
+//!   assert_eq!(8, vm.icnt);
+//!   assert_eq!(4, vm.pc);
+//! }
+//! ```
+
+#[cfg(not(any(target_arch = "x86_64", target_os = "linux")))]
+compile_error!("Only supported on x86_64 with SystemV abi");
+
+use juicebox_asm::prelude::*;
+use juicebox_asm::Runtime;
+
+/// A guest physical address.
+pub struct PhysAddr(pub u16);
+
+impl Into<usize> for PhysAddr {
+    fn into(self) -> usize {
+        self.0 as usize
+    }
+}
+
+/// The registers for the [`TinyVm`].
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub enum TinyReg {
+    A,
+    B,
+    C,
+}
+
+impl TinyReg {
+    #[inline]
+    fn idx(&self) -> usize {
+        *self as usize
+    }
+}
+
+/// The instructions for the [`TinyVm`].
+#[derive(Debug, PartialEq, Clone, Copy)]
+pub enum TinyInsn {
+    /// Halt the VM.
+    Halt,
+    /// Load the immediate value into the register `reg = imm`.
+    LoadImm(TinyReg, u16),
+    /// Load a value from the memory (absolute addressing) into the register `reg = mem[imm]`.
+    Load(TinyReg, u16),
+    /// Store a value from the register into the memory (absolute addressing) `mem[imm] = reg`.
+    Store(TinyReg, u16),
+    /// Add the register to the register `reg1 += reg2`.
+    Add(TinyReg, TinyReg),
+    /// Add the immediate to the register `reg += imm`.
+    Addi(TinyReg, i16),
+    /// Jump unconditional (absolute addressing) `pc = disp`.
+    Branch(usize),
+    /// Jump if the register is zero (absolute addressing) `pc = (reg == 0) ? disp : pc++`.
+    BranchZero(TinyReg, usize),
+}
+
+/// Value returned from a [`JitFn`].
+#[repr(C)]
+struct JitRet(u64, u64);
+
+/// Function signature defining the simple JIT ABI used in this example.
+/// A `JitFn` represents the entry point to a jit compiled _basic block_ of the guest software.
+///
+/// ```text
+/// JIT entry:
+///     arg0: pointer to guest registers
+///     arg1: pointer to guest data memory
+///
+/// JIT exit:
+///      JitRet(0, N): Halt instruction, executed N instructions.
+///      JitRet(N, R): N!=0
+///                    End of basic block, executed N instructions,
+///                    must re-enter at `pc = R`.
+/// ```
+type JitFn = extern "C" fn(*mut u16, *mut u8) -> JitRet;
+
+/// The `TinyVm` virtual machine state.
+pub struct TinyVm {
+    /// Data memory, covering full 16 bit guest address space.
+    ///
+    /// For simplicity add additional trailing 1 byte to support an unaligned access to 0xffff
+    /// without any special handling.
+    dmem: [u8; 0x1_0000 + 1],
+    /// Instruction memory.
+    imem: Vec<TinyInsn>,
+    /// VM registers.
+    regs: [u16; 3],
+    /// VM program counter.
+    pc: usize,
+    /// VM executed instruction counter (perf counter).
+    icnt: usize,
+
+    // -- JIT state.
+    /// Mapping of guest PCs to jitted host code (`JitFn`). This mapping is filled when guest
+    /// _basic blocks_ are jitted.
+    jit_cache: Vec<Option<JitFn>>,
+    /// JIT runtime maintaining the host pages containing the jitted guest code.
+    rt: Runtime,
+}
+
+impl TinyVm {
+    /// Create a new [`TinyVm`] and initialize the instruction memory from `code`.
+    pub fn new(code: Vec<TinyInsn>) -> Self {
+        let mut jit_cache = Vec::with_capacity(code.len());
+        jit_cache.resize(code.len(), None);
+
+        TinyVm {
+            dmem: [0; 0x1_0000 + 1],
+            imem: code,
+            regs: [0; 3],
+            pc: 0,
+            icnt: 0,
+            // -- JIT state.
+            jit_cache,
+            rt: Runtime::new(),
+        }
+    }
+
+    /// Read guest register.
+    #[inline]
+    pub fn read_reg(&self, reg: TinyReg) -> u16 {
+        self.regs[reg.idx()]
+    }
+
+    /// Write guest register.
+    #[inline]
+    pub fn write_reg(&mut self, reg: TinyReg, val: u16) {
+        self.regs[reg.idx()] = val;
+    }
+
+    /// Read guest data memory.
+    #[inline]
+    pub fn read_mem(&self, paddr: PhysAddr) -> u16 {
+        // dmem covers whole 16 bit address space + 1 byte for unaligned access at 0xffff.
+        let bytes = self.dmem[paddr.into()..][..2].try_into().unwrap();
+        u16::from_le_bytes(bytes)
+    }
+
+    /// Write guest data memory.
+    #[inline]
+    pub fn write_mem(&mut self, paddr: PhysAddr, val: u16) {
+        let bytes = val.to_le_bytes();
+        self.dmem[paddr.into()..][..2].copy_from_slice(&bytes);
+    }
+
+    /// Dump the VM state to stdout.
+    pub fn dump(&self) {
+        println!("-- TinyVm state --");
+        println!("  ICNT: {}", self.icnt);
+        println!("  PC  : {:02x}", self.pc - 1);
+        println!(
+            "  A:{:04x} B:{:04x} C:{:04x}",
+            self.read_reg(TinyReg::A),
+            self.read_reg(TinyReg::B),
+            self.read_reg(TinyReg::C),
+        );
+    }
+
+    /// Run in interpreter mode until the next [`TinyInsn::Halt`] instruction is hit.
+    pub fn interp(&mut self) {
+        'outer: loop {
+            let insn = self.imem[self.pc];
+            //println!("[0x{:02x}] {:?}", self.pc, insn);
+
+            self.pc = self.pc.wrapping_add(1);
+            self.icnt += 1;
+
+            match insn {
+                TinyInsn::Halt => {
+                    break 'outer;
+                }
+                TinyInsn::LoadImm(a, imm) => {
+                    self.write_reg(a, imm);
+                }
+                TinyInsn::Load(a, addr) => {
+                    let val = self.read_mem(PhysAddr(addr));
+                    self.write_reg(a, val);
+                }
+                TinyInsn::Store(a, addr) => {
+                    let val = self.read_reg(a);
+                    self.write_mem(PhysAddr(addr), val);
+                }
+                TinyInsn::Add(a, b) => {
+                    let res = self.read_reg(a).wrapping_add(self.read_reg(b));
+                    self.write_reg(a, res);
+                }
+                TinyInsn::Addi(a, imm) => {
+                    let res = self.read_reg(a).wrapping_add(imm as u16);
+                    self.write_reg(a, res);
+                }
+                TinyInsn::Branch(disp) => {
+                    self.pc = disp;
+                }
+                TinyInsn::BranchZero(a, disp) => {
+                    if self.read_reg(a) == 0 {
+                        self.pc = disp;
+                    }
+                }
+            }
+        }
+    }
+
+    /// Run in JIT mode until the next [`TinyInsn::Halt`] instruction is hit. Translate guest
+    /// _basic blocks_ on demand.
+    pub fn jit(&mut self) {
+        'outer: loop {
+            if let Some(bb_fn) = self.jit_cache[self.pc] {
+                match bb_fn(self.regs.as_mut_ptr(), self.dmem.as_mut_ptr()) {
+                    JitRet(0, insn) => {
+                        self.pc += insn as usize;
+                        self.icnt += insn as usize;
+                        break 'outer;
+                    }
+                    JitRet(insn, reenter_pc) => {
+                        self.pc = reenter_pc as usize;
+                        self.icnt += insn as usize;
+                    }
+                }
+            } else {
+                let bb_fn = self.translate_next_bb();
+                self.jit_cache[self.pc] = Some(bb_fn);
+                //println!("[0x{:02x}] translated bb at {:p}", self.pc, bb_fn);
+            }
+        }
+    }
+
+    /// Translate the bb at the current pc and return a JitFn pointer to it.
+    fn translate_next_bb(&mut self) -> JitFn {
+        let mut bb = Asm::new();
+        let mut pc = self.pc;
+
+        'outer: loop {
+            let insn = self.imem[pc];
+
+            pc = pc.wrapping_add(1);
+
+            // JIT ABI.
+            //   enter
+            //     rdi => regs
+            //     rsi => dmem
+            //   exit
+            //     rax => JitRet(0,
+            //     rdx =>        1)
+
+            // Generate memory operand into regs for guest register.
+            let reg_op = |r: TinyReg| {
+                MemOp::IndirectDisp(Reg64::rdi, (r.idx() * 2).try_into().expect("only 3 regs"))
+            };
+
+            // Generate memory operand into dmem for guest phys address.
+            let mem_op = |paddr: u16| MemOp::IndirectDisp(Reg64::rsi, paddr.into());
+
+            // Compute instructions in translated basic block.
+            let bb_icnt = || -> u64 { (pc - self.pc).try_into().unwrap() };
+
+            let reenter_pc = |pc: usize| -> u64 { pc.try_into().unwrap() };
+
+            match insn {
+                TinyInsn::Halt => {
+                    bb.mov(Reg64::rax, Imm64::from(0));
+                    bb.mov(Reg64::rdx, Imm64::from(bb_icnt()));
+                    bb.ret();
+                    break 'outer;
+                }
+                TinyInsn::LoadImm(a, imm) => {
+                    bb.mov(reg_op(a), Imm16::from(imm));
+                }
+                TinyInsn::Load(a, addr) => {
+                    bb.mov(Reg16::ax, mem_op(addr));
+                    bb.mov(reg_op(a), Reg16::ax);
+                }
+                TinyInsn::Store(a, addr) => {
+                    bb.mov(Reg16::ax, reg_op(a));
+                    bb.mov(mem_op(addr), Reg16::ax);
+                }
+                TinyInsn::Add(a, b) => {
+                    bb.mov(Reg16::ax, reg_op(b));
+                    bb.add(reg_op(a), Reg16::ax);
+                }
+                TinyInsn::Addi(a, imm) => {
+                    bb.add(reg_op(a), Imm16::from(imm));
+                }
+                TinyInsn::Branch(disp) => {
+                    bb.mov(Reg64::rax, Imm64::from(bb_icnt()));
+                    bb.mov(Reg64::rdx, Imm64::from(reenter_pc(disp)));
+                    bb.ret();
+                    break 'outer;
+                }
+                TinyInsn::BranchZero(a, disp) => {
+                    bb.cmp(reg_op(a), Imm16::from(0u16));
+                    bb.mov(Reg64::rax, Imm64::from(bb_icnt()));
+                    bb.mov(Reg64::rdx, Imm64::from(reenter_pc(disp)));
+
+                    let mut skip_next_pc = Label::new();
+                    // If register is zero, skip setting next pc as reenter pc.
+                    bb.jz(&mut skip_next_pc);
+                    bb.mov(Reg64::rdx, Imm64::from(reenter_pc(pc)));
+                    bb.bind(&mut skip_next_pc);
+                    bb.ret();
+                    break 'outer;
+                }
+            }
+        }
+
+        unsafe { self.rt.add_code::<JitFn>(bb.into_code()) }
+    }
+}
+
+/// A minial fixup utility to implement jump labels when constructing guest programs.
+pub struct Fixup {
+    pc: usize,
+}
+
+impl Fixup {
+    /// Create a new `Fixup` at the current pc.
+    pub fn new(pc: usize) -> Self {
+        Fixup { pc }
+    }
+
+    /// Bind the `Fixup` to the current location of `prog` and resolve the `Fixup`.
+    pub fn bind(self, prog: &mut Vec<TinyInsn>) {
+        let plen = prog.len();
+        let insn = prog.get_mut(self.pc).expect(&format!(
+            "Trying to apply Fixup, but Fixup is out of range pc={} prog.len={}",
+            self.pc, plen
+        ));
+
+        match insn {
+            TinyInsn::Branch(disp) | TinyInsn::BranchZero(_, disp) => {
+                *disp = plen;
+            }
+            _ => {
+                unimplemented!("Trying to fixup non-branch instruction '{:?}'", *insn);
+            }
+        }
+    }
+}
+
+/// Generate a guest program to compute the fiibonacci sequence for `n`.
+pub fn make_tinyvm_fib(start_n: u16) -> Vec<TinyInsn> {
+    // Reference implementation:
+    //
+    // int fib(int n)
+    //   int tmp = 0;
+    //   int prv = 1;
+    //   int sum = 0;
+    // loop:
+    //   if (n == 0) goto end;
+    //   tmp = sum;
+    //   sum += prv;
+    //   prv = tmp;
+    //   --n;
+    //   goto loop;
+    // end:
+    //   return sum;
+
+    // Variables live in memory, bin to fixed addresses.
+    let tmp = 0u16;
+    let prv = 2u16;
+    let sum = 4u16;
+    // Loop counter mapped to register.
+    let n = TinyReg::C;
+
+    let mut prog = Vec::with_capacity(32);
+
+    // n = start_n
+    prog.push(TinyInsn::LoadImm(n, start_n));
+
+    // tmp = sum = 0
+    prog.push(TinyInsn::LoadImm(TinyReg::A, 0));
+    prog.push(TinyInsn::Store(TinyReg::A, tmp));
+    prog.push(TinyInsn::Store(TinyReg::A, sum));
+
+    // prv = 1
+    prog.push(TinyInsn::LoadImm(TinyReg::A, 1));
+    prog.push(TinyInsn::Store(TinyReg::A, prv));
+
+    // Create loop_start label.
+    let loop_start = prog.len();
+
+    // Create fixup to capture PC that need to be patched later.
+    let end_fixup = Fixup::new(prog.len());
+
+    // if (n == 0) goto end
+    prog.push(TinyInsn::BranchZero(n, 0xdead));
+
+    // tmp = sum
+    prog.push(TinyInsn::Load(TinyReg::A, sum));
+    prog.push(TinyInsn::Store(TinyReg::A, tmp));
+
+    // sum += prv
+    prog.push(TinyInsn::Load(TinyReg::B, prv));
+    prog.push(TinyInsn::Add(TinyReg::A, TinyReg::B));
+    prog.push(TinyInsn::Store(TinyReg::A, sum));
+
+    // prv = tmp
+    prog.push(TinyInsn::Load(TinyReg::A, tmp));
+    prog.push(TinyInsn::Store(TinyReg::A, prv));
+
+    // --n
+    prog.push(TinyInsn::Addi(n, -1));
+
+    // goto loop_start
+    prog.push(TinyInsn::Branch(loop_start));
+
+    // Bind end fixup to current PC, to patch branch to jump to here.
+    end_fixup.bind(&mut prog);
+
+    // TinyReg::A = sum
+    prog.push(TinyInsn::Load(TinyReg::A, sum));
+
+    // Halt the VM.
+    prog.push(TinyInsn::Halt);
+
+    prog
+}
+
+/// Generate a test program for the jit.
+pub fn make_tinyvm_jit_test() -> Vec<TinyInsn> {
+    let mut prog = Vec::with_capacity(32);
+
+    prog.push(TinyInsn::Branch(1));
+    prog.push(TinyInsn::LoadImm(TinyReg::A, 0x0010));
+    prog.push(TinyInsn::LoadImm(TinyReg::B, 0x0));
+
+    let start = prog.len();
+    let end = Fixup::new(prog.len());
+    prog.push(TinyInsn::BranchZero(TinyReg::A, 0xdead));
+    prog.push(TinyInsn::LoadImm(TinyReg::C, 0x1));
+    prog.push(TinyInsn::Add(TinyReg::B, TinyReg::C));
+    prog.push(TinyInsn::Addi(TinyReg::A, -1));
+    prog.push(TinyInsn::Branch(start));
+    end.bind(&mut prog);
+    prog.push(TinyInsn::LoadImm(TinyReg::A, 0xabcd));
+    prog.push(TinyInsn::Store(TinyReg::A, 0xffff));
+    prog.push(TinyInsn::Load(TinyReg::C, 0xffff));
+    prog.push(TinyInsn::Halt);
+    prog.push(TinyInsn::Halt);
+    prog.push(TinyInsn::Halt);
+    prog.push(TinyInsn::Halt);
+
+    prog
+}
+
+/// Generate a simple count down loop to crunch some instructions.
+pub fn make_tinyvm_jit_perf() -> Vec<TinyInsn> {
+    let mut prog = Vec::with_capacity(32);
+
+    prog.push(TinyInsn::Halt);
+    prog.push(TinyInsn::LoadImm(TinyReg::A, 0xffff));
+    prog.push(TinyInsn::LoadImm(TinyReg::B, 1));
+    prog.push(TinyInsn::LoadImm(TinyReg::C, 2));
+    prog.push(TinyInsn::Addi(TinyReg::A, -1));
+    prog.push(TinyInsn::BranchZero(TinyReg::A, 0));
+    prog.push(TinyInsn::Branch(2));
+    prog
+}
+
+fn main() {
+    let use_jit = match std::env::args().nth(1) {
+        Some(a) if a == "-h" || a == "--help" => {
+            println!("Usage: tiny_vm [mode]");
+            println!("");
+            println!("Options:");
+            println!("    mode    if mode is 'jit' then run in jit mode, else in interpreter mode");
+            std::process::exit(0);
+        }
+        Some(a) if a == "jit" => true,
+        _ => false,
+    };
+
+    let mut vm = TinyVm::new(make_tinyvm_fib(42));
+
+    if use_jit {
+        println!("Run in jit mode..");
+        vm.jit();
+    } else {
+        println!("Run in interpreter mode..");
+        vm.interp();
+    }
+    vm.dump();
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    fn fib_rs(n: u64) -> u64 {
+        if n < 2 {
+            n
+        } else {
+            let mut fib_n_m1 = 0;
+            let mut fib_n = 1;
+            for _ in 1..n {
+                let tmp = fib_n + fib_n_m1;
+                fib_n_m1 = fib_n;
+                fib_n = tmp;
+            }
+            fib_n
+        }
+    }
+
+    #[test]
+    fn test_fib_interp() {
+        for n in 0..92 {
+            let mut vm = TinyVm::new(make_tinyvm_fib(n));
+            vm.interp();
+
+            assert_eq!((fib_rs(n as u64) & 0xffff) as u16, vm.read_reg(TinyReg::A));
+        }
+    }
+
+    #[test]
+    fn test_fib_jit() {
+        for n in 0..92 {
+            let mut vm = TinyVm::new(make_tinyvm_fib(n));
+            vm.jit();
+
+            assert_eq!((fib_rs(n as u64) & 0xffff) as u16, vm.read_reg(TinyReg::A));
+        }
+    }
+
+    #[test]
+    fn test_fib_icnt() {
+        let mut vm1 = TinyVm::new(make_tinyvm_fib(91));
+        vm1.interp();
+        let mut vm2 = TinyVm::new(make_tinyvm_fib(91));
+        vm2.jit();
+
+        assert_eq!(vm1.icnt, vm2.icnt);
+        assert_eq!(vm1.pc, vm2.pc);
+    }
+
+    #[test]
+    fn test_jit_load_imm() {
+        let mut prog = Vec::new();
+        prog.push(TinyInsn::LoadImm(TinyReg::A, 0x1111));
+        prog.push(TinyInsn::LoadImm(TinyReg::B, 0x2222));
+        prog.push(TinyInsn::LoadImm(TinyReg::C, 0x3333));
+        prog.push(TinyInsn::Halt);
+
+        let mut vm = TinyVm::new(prog);
+        vm.jit();
+
+        assert_eq!(0x1111, vm.read_reg(TinyReg::A));
+        assert_eq!(0x2222, vm.read_reg(TinyReg::B));
+        assert_eq!(0x3333, vm.read_reg(TinyReg::C));
+
+        assert_eq!(4, vm.icnt);
+        assert_eq!(4, vm.pc);
+    }
+
+    #[test]
+    fn test_jit_add() {
+        let mut prog = Vec::new();
+        prog.push(TinyInsn::LoadImm(TinyReg::A, 0));
+        prog.push(TinyInsn::Addi(TinyReg::A, 123));
+
+        prog.push(TinyInsn::LoadImm(TinyReg::B, 100));
+        prog.push(TinyInsn::LoadImm(TinyReg::C, 200));
+        prog.push(TinyInsn::Add(TinyReg::B, TinyReg::C));
+        prog.push(TinyInsn::Halt);
+
+        let mut vm = TinyVm::new(prog);
+        vm.jit();
+
+        assert_eq!(123, vm.read_reg(TinyReg::A));
+        assert_eq!(300, vm.read_reg(TinyReg::B));
+        assert_eq!(200, vm.read_reg(TinyReg::C));
+
+        assert_eq!(6, vm.icnt);
+        assert_eq!(6, vm.pc);
+    }
+
+    #[test]
+    fn test_jit_load_store() {
+        let mut prog = Vec::new();
+        prog.push(TinyInsn::Load(TinyReg::A, 0xffff));
+
+        prog.push(TinyInsn::LoadImm(TinyReg::B, 0xf00d));
+        prog.push(TinyInsn::Store(TinyReg::B, 0x8000));
+        prog.push(TinyInsn::Halt);
+
+        let mut vm = TinyVm::new(prog);
+        vm.write_mem(PhysAddr(0xffff), 0xaabb);
+        vm.jit();
+
+        assert_eq!(0xaabb, vm.read_reg(TinyReg::A));
+        assert_eq!(0xf00d, vm.read_mem(PhysAddr(0x8000)));
+
+        assert_eq!(4, vm.icnt);
+        assert_eq!(4, vm.pc);
+    }
+
+    #[test]
+    fn test_jit_branch() {
+        let mut prog = Vec::new();
+        prog.push(TinyInsn::Branch(2));
+        prog.push(TinyInsn::Halt);
+        prog.push(TinyInsn::Branch(6));
+        prog.push(TinyInsn::LoadImm(TinyReg::A, 1));
+        prog.push(TinyInsn::LoadImm(TinyReg::B, 2));
+        prog.push(TinyInsn::LoadImm(TinyReg::C, 3));
+        prog.push(TinyInsn::Branch(1));
+
+        let mut vm = TinyVm::new(prog);
+        vm.jit();
+
+        assert_eq!(0, vm.read_reg(TinyReg::A));
+        assert_eq!(0, vm.read_reg(TinyReg::B));
+        assert_eq!(0, vm.read_reg(TinyReg::C));
+
+        assert_eq!(4, vm.icnt);
+        assert_eq!(2, vm.pc);
+    }
+
+    #[test]
+    fn test_jit_branch_zero() {
+        let mut prog = Vec::new();
+        prog.push(TinyInsn::LoadImm(TinyReg::A, 1));
+        prog.push(TinyInsn::BranchZero(TinyReg::A, 5));
+        prog.push(TinyInsn::LoadImm(TinyReg::A, 0));
+        prog.push(TinyInsn::BranchZero(TinyReg::A, 5));
+        prog.push(TinyInsn::LoadImm(TinyReg::B, 22));
+        prog.push(TinyInsn::Halt);
+
+        let mut vm = TinyVm::new(prog);
+        vm.jit();
+
+        assert_eq!(0, vm.read_reg(TinyReg::A));
+        assert_eq!(0, vm.read_reg(TinyReg::B));
+        assert_eq!(0, vm.read_reg(TinyReg::C));
+
+        assert_eq!(5, vm.icnt);
+        assert_eq!(6, vm.pc);
+    }
+
+    #[test]
+    fn test_mixed() {
+        let mut prog = Vec::new();
+        prog.push(TinyInsn::LoadImm(TinyReg::A, 100));
+        prog.push(TinyInsn::Add(TinyReg::B, TinyReg::A));
+        prog.push(TinyInsn::Addi(TinyReg::C, 100));
+        prog.push(TinyInsn::Halt);
+
+        let mut vm = TinyVm::new(prog);
+        vm.interp();
+
+        assert_eq!(100, vm.read_reg(TinyReg::A));
+        assert_eq!(100, vm.read_reg(TinyReg::B));
+        assert_eq!(100, vm.read_reg(TinyReg::C));
+        assert_eq!(4, vm.icnt);
+        assert_eq!(4, vm.pc);
+
+        vm.pc = 0;
+        vm.jit();
+
+        assert_eq!(100, vm.read_reg(TinyReg::A));
+        assert_eq!(200, vm.read_reg(TinyReg::B));
+        assert_eq!(200, vm.read_reg(TinyReg::C));
+        assert_eq!(8, vm.icnt);
+        assert_eq!(4, vm.pc);
+    }
+}