//! TinyVm example.
//!
//! This example introduces a simple 16 bit virtual machine the [`TinyVm`]. The VM consists of
//! three registers defined in [`TinyReg`], a separate _data_ and _instruction_ memory and a small
//! set of instructions [`TinyInsn`], sufficient to implement a guest program to compute the
//! Fibonacci 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);
//! }
//! ```

use juicebox_asm::insn::*;
use juicebox_asm::Runtime;
use juicebox_asm::{Asm, Imm16, Imm64, Label, MemOp, Reg16, Reg64};

/// 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(),
            // Confifigure the runtime to generates perf meta data.
            //rt: Runtime::with_profile(),
        }
    }

    /// 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 {
            let bb_fn = if let Some(bb_fn) = self.jit_cache[self.pc] {
                bb_fn
            } 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);
                bb_fn
            };

            match bb_fn(self.regs.as_mut_ptr(), self.dmem.as_mut_ptr()) {
                // HALT instruction hit.
                JitRet(0, insn) => {
                    self.pc += insn as usize;
                    self.icnt += insn as usize;
                    break 'outer;
                }
                // End of basic block, re-enter.
                JitRet(insn, reenter_pc) => {
                    self.pc = reenter_pc as usize;
                    self.icnt += insn as usize;
                }
            }
        }
    }

    #[cfg(all(any(target_arch = "x86_64", target_os = "linux")))]
    /// 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: JitFn -> JitRet
            //
            // According to SystemV abi:
            //   enter
            //     rdi => regs
            //     rsi => dmem
            //   exit
            //     rax => JitRet.0
            //     rdx => JitRet.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);
    }
}