zkvmc-interpreter 0.0.1

use std::ptr::NonNull;

use zkvmc_backend::{Backend, Cond, Signed, Width};
use zkvmc_context::Context;
use zkvmc_core::{
    addr2frame::Frame,
    memory::{read_bytes, read_word, write_bytes},
};

pub struct Interpreter<'a, const DEBUG: bool> {
    pub(crate) ctx: &'a mut Context,
    frames: Vec<(u32, u32, Option<Frame>)>,
}

impl<'a, const DEBUG: bool> Interpreter<'a, DEBUG> {
    pub(crate) fn new(ctx: &'a mut Context) -> Self {
        Self {
            ctx,
            frames: Vec::new(),
        }
    }
}

impl<const DEBUG: bool> Drop for Interpreter<'_, DEBUG> {
    fn drop(&mut self) {
        if !DEBUG {
            return;
        }
        for (pc, target, frame) in self.frames.iter() {
            match frame {
                Some(frame) => {
                    println!(
                        "{:x}->{:x}\t{}@{}:{}",
                        pc, target, frame.name, frame.filename, frame.lineno
                    );
                }
                None => {
                    println!("{:x}->{:x}\t<unknown>", pc, target);
                }
            }
        }
    }
}

impl<const DEBUG: bool> Backend for Interpreter<'_, DEBUG> {
    type Value = u32;

    fn get_reg(&mut self, index: u32) -> Self::Value {
        self.ctx.regs[index as usize]
    }

    fn set_reg(&mut self, index: u32, val: Self::Value) {
        if index != 0 {
            self.ctx.regs[index as usize] = val;
        }
    }

    fn get_pc(&mut self) -> u32 {
        self.ctx.pc
    }

    fn set_next_pc(&mut self, pc: u32) {
        self.ctx.next_pc = pc;
    }

    fn fetch_u32(&mut self, addr: u32) -> u32 {
        unsafe { read_word(self.ctx.mem.as_ptr(), addr) }
    }

    fn const_value(&mut self, val: u32) -> Self::Value {
        val
    }

    fn load(&mut self, width: Width, signed: Signed, addr: Self::Value) -> Self::Value {
        match (width, signed) {
            (Width::Byte, Signed::Unsigned) => unsafe {
                read_bytes(self.ctx.mem.as_ptr(), addr, 1)[0] as u32
            },
            (Width::Byte, Signed::Signed) => unsafe {
                read_bytes(self.ctx.mem.as_ptr(), addr, 1)[0] as i8 as u32
            },
            (Width::Half, Signed::Unsigned) => u16::from_le_bytes(unsafe {
                read_bytes(self.ctx.mem.as_ptr(), addr, 2)
                    .try_into()
                    .unwrap()
            }) as u32,
            (Width::Half, Signed::Signed) => i16::from_le_bytes(unsafe {
                read_bytes(self.ctx.mem.as_ptr(), addr, 2)
                    .try_into()
                    .unwrap()
            }) as u32,
            (Width::Word, _) => unsafe { read_word(self.ctx.mem.as_ptr(), addr) },
        }
    }

    fn store(&mut self, width: Width, addr: Self::Value, val: Self::Value) {
        unsafe {
            write_bytes(
                self.ctx.mem.as_ptr(),
                addr,
                &val.to_le_bytes()[..width as usize],
            )
        };
    }

    fn add(&mut self, a: Self::Value, b: Self::Value) -> Self::Value {
        a.wrapping_add(b)
    }

    fn sub(&mut self, a: Self::Value, b: Self::Value) -> Self::Value {
        a.wrapping_sub(b)
    }

    fn and(&mut self, a: Self::Value, b: Self::Value) -> Self::Value {
        a & b
    }

    fn or(&mut self, a: Self::Value, b: Self::Value) -> Self::Value {
        a | b
    }

    fn xor(&mut self, a: Self::Value, b: Self::Value) -> Self::Value {
        a ^ b
    }

    fn shift_left(&mut self, val: Self::Value, shift: Self::Value) -> Self::Value {
        val.wrapping_shl(shift)
    }

    fn shift_right(&mut self, signed: Signed, val: Self::Value, shift: Self::Value) -> Self::Value {
        match signed {
            Signed::Unsigned => val.wrapping_shr(shift),
            Signed::Signed => (val as i32).wrapping_shr(shift) as u32,
        }
    }

    fn jump(&mut self, target: Self::Value) {
        self.set_next_pc(target);
        if DEBUG {
            self.frames.push((self.ctx.pc, target, unsafe {
                self.ctx.addition.as_ref().lookup_pc(target)
            }));
        }
    }

    fn ecall(&mut self) {
        (self.ctx.ecall_handler)(NonNull::new(self.ctx).unwrap());
    }

    fn ebreak(&mut self) {
        (self.ctx.ebreak_handler)(NonNull::new(self.ctx).unwrap());
    }

    fn undefined(&mut self, instr: u32) {
        (self.ctx.undefined_handler)(NonNull::new(self.ctx).unwrap(), instr);
    }

    fn if_else_val(
        &mut self,
        cond: Cond,
        a: Self::Value,
        b: Self::Value,
        if_true: impl FnOnce(&mut Self) -> Self::Value,
        if_false: impl FnOnce(&mut Self) -> Self::Value,
    ) -> Self::Value {
        let cond = match cond {
            Cond::Eq => a == b,
            Cond::Ne => a != b,
            Cond::Ge(Signed::Signed) => a as i32 >= b as i32,
            Cond::Lt(Signed::Signed) => (a as i32) < b as i32,
            Cond::Ge(Signed::Unsigned) => a >= b,
            Cond::Lt(Signed::Unsigned) => a < b,
        };
        if cond {
            if_true(self)
        } else {
            if_false(self)
        }
    }

    fn mul(&mut self, a: Self::Value, b: Self::Value) -> Self::Value {
        a.wrapping_mul(b)
    }

    fn mulh(&mut self, signed: Signed, a: Self::Value, b: Self::Value) -> Self::Value {
        match signed {
            Signed::Unsigned => ((a as u64).wrapping_mul(b as u64) >> 32) as u32,
            Signed::Signed => ((a as i32 as i64).wrapping_mul(b as i32 as i64) >> 32) as u32,
        }
    }

    fn mulhsu(&mut self, a: Self::Value, b: Self::Value) -> Self::Value {
        ((a as i32 as i64).wrapping_mul(b as i64) >> 32) as u32
    }

    fn div(&mut self, signed: Signed, a: Self::Value, b: Self::Value) -> Self::Value {
        match signed {
            Signed::Unsigned => {
                if b == 0 {
                    return 0xffff_ffff;
                }
                a.wrapping_div(b)
            }
            Signed::Signed => {
                if b == 0 {
                    return 0xffff_ffff;
                }
                (a as i32).wrapping_div(b as i32) as u32
            }
        }
    }

    fn rem(&mut self, signed: Signed, a: Self::Value, b: Self::Value) -> Self::Value {
        match signed {
            Signed::Unsigned => {
                if b == 0 {
                    return a;
                }
                a.wrapping_rem(b)
            }
            Signed::Signed => {
                if b == 0 {
                    return a;
                }
                (a as i32).wrapping_rem(b as i32) as u32
            }
        }
    }
}