dispatch.rs

#[repr(u16)]
#[derive(Copy, Clone)]
pub enum Instruction {
    I32Add = Self::I32_ADD_OPCODE,
    I32AddImm = Self::I32_ADD_IMM_OPCODE,
    I32Sub = Self::I32_SUB_OPCODE,
    I32SubImm = Self::I32_SUB_IMM_OPCODE,
    I32Mul = Self::I32_MUL_OPCODE,
    I32MulImm = Self::I32_MUL_IMM_OPCODE,
    I64Add = Self::I64_ADD_OPCODE,
    I64AddImm = Self::I64_ADD_IMM_OPCODE,
    // etc..
}

impl Instruction {
    const I32_ADD_OPCODE: u16 = 0;
    const I32_ADD_IMM_OPCODE: u16 = 1;
    const I32_SUB_OPCODE: u16 = 2;
    const I32_SUB_IMM_OPCODE: u16 = 3;
    const I32_MUL_OPCODE: u16 = 4;
    const I32_MUL_IMM_OPCODE: u16 = 5;
    const I64_ADD_OPCODE: u16 = 6;
    const I64_ADD_IMM_OPCODE: u16 = 7;
}

impl From<u16> for Instruction {
    fn from(value: u16) -> Self {
        match value {
            Self::I32_ADD_OPCODE => Self::I32Add,
            Self::I32_ADD_IMM_OPCODE => Self::I32AddImm,
            Self::I32_SUB_OPCODE => Self::I32Sub,
            Self::I32_SUB_IMM_OPCODE => Self::I32SubImm,
            Self::I32_MUL_OPCODE => Self::I32Mul,
            Self::I32_MUL_IMM_OPCODE => Self::I32MulImm,
            _ => unsafe { core::hint::unreachable_unchecked() },
        }
    }
}

impl Instruction {
    #[no_mangle]
    pub fn execute(&self, ctx: &mut Context, ip: InstructionPtr) -> Result<(), Trap> {
        match self {
            Self::I32Add => execute_i32_add(ctx, ip),
            Self::I32AddImm => execute_i32_add_imm(ctx, ip),
            Self::I32Sub => execute_i32_sub(ctx, ip),
            Self::I32SubImm => execute_i32_sub_imm(ctx, ip),
            Self::I32Mul => execute_i32_mul(ctx, ip),
            Self::I32MulImm => execute_i32_mul_imm(ctx, ip),
            Self::I64Add => execute_i64_add(ctx, ip),
            Self::I64AddImm => execute_i64_add_imm(ctx, ip),
        }
    }
}

#[derive(Copy, Clone)]
pub enum Trap {
    DivisionByZero,
    InvalidInteger,
    StackOverflow,
    SignatureMismatch,
    OutOfBounds,
}

#[derive(Copy, Clone)]
pub struct Register(pub u16);

pub struct Context {
    stack: Vec<u64>,
}

impl Context {
    pub fn get_register(&self, reg: Register) -> u64 {
        * unsafe { self.stack.get_unchecked(reg.0 as usize) }
    }

    pub fn set_register(&mut self, reg: Register, value: u64) {
        *unsafe { self.stack.get_unchecked_mut(reg.0 as usize) } = value;
    }
}

#[no_mangle]
fn execute_i32_add(ctx: &mut Context, mut ip: InstructionPtr) -> Result<(), Trap> {
    let res = ip.decode_register();
    let lhs = ip.decode_register();
    let rhs = ip.decode_register();
    let lval = ctx.get_register(lhs) as i32;
    let rval = ctx.get_register(rhs) as i32;
    let result = lval.wrapping_add(rval);
    ctx.set_register(res, result as u64);
    Ok(())
}

#[no_mangle]
fn execute_i32_add_imm(ctx: &mut Context, mut ip: InstructionPtr) -> Result<(), Trap> {
    let res = ip.decode_register();
    let lhs = ip.decode_register();
    let lval = ctx.get_register(ip.decode_register()) as i32;
    let rval = ip.decode_i32();
    let result = lval.wrapping_add(rval);
    ctx.set_register(res, result as u64);
    Ok(())
}

#[no_mangle]
fn execute_i32_sub(ctx: &mut Context, mut ip: InstructionPtr) -> Result<(), Trap> {
    let res = ip.decode_register();
    let lhs = ip.decode_register();
    let rhs = ip.decode_register();
    let lval = ctx.get_register(lhs) as i32;
    let rval = ctx.get_register(rhs) as i32;
    let result = lval.wrapping_sub(rval);
    ctx.set_register(res, result as u64);
    Ok(())
}

#[no_mangle]
fn execute_i32_sub_imm(ctx: &mut Context, mut ip: InstructionPtr) -> Result<(), Trap> {
    let res = ip.decode_register();
    let lhs = ip.decode_register();
    let lval = ctx.get_register(ip.decode_register()) as i32;
    let rval = ip.decode_i32();
    let result = lval.wrapping_sub(rval);
    ctx.set_register(res, result as u64);
    Ok(())
}

#[no_mangle]
fn execute_i32_mul(ctx: &mut Context, mut ip: InstructionPtr) -> Result<(), Trap> {
    let res = ip.decode_register();
    let lhs = ip.decode_register();
    let rhs = ip.decode_register();
    let lval = ctx.get_register(lhs) as i32;
    let rval = ctx.get_register(rhs) as i32;
    let result = lval.wrapping_mul(rval);
    ctx.set_register(res, result as u64);
    Ok(())
}

#[no_mangle]
fn execute_i32_mul_imm(ctx: &mut Context, mut ip: InstructionPtr) -> Result<(), Trap> {
    let res = ip.decode_register();
    let lhs = ip.decode_register();
    let lval = ctx.get_register(ip.decode_register()) as i32;
    let rval = ip.decode_i32();
    let result = lval.wrapping_mul(rval);
    ctx.set_register(res, result as u64);
    Ok(())
}

#[no_mangle]
fn execute_i64_add(ctx: &mut Context, mut ip: InstructionPtr) -> Result<(), Trap> {
    let res = ip.decode_register();
    let lhs = ip.decode_register();
    let rhs = ip.decode_register();
    let lval = ctx.get_register(lhs) as i64;
    let rval = ctx.get_register(rhs) as i64;
    let result = lval.wrapping_add(rval);
    ctx.set_register(res, result as u64);
    Ok(())
}

#[no_mangle]
fn execute_i64_add_imm(ctx: &mut Context, mut ip: InstructionPtr) -> Result<(), Trap> {
    let res = ip.decode_register();
    let lhs = ip.decode_register();
    let lval = ctx.get_register(ip.decode_register()) as i64;
    let rval = ip.decode_i64();
    let result = lval.wrapping_add(rval);
    ctx.set_register(res, result as u64);
    Ok(())
}

pub struct InstructionEncoder {
    buffer: Vec<u8>,
}

impl InstructionEncoder {
    // #[inline]
    pub fn encode_bytes<const N: usize>(&mut self, bytes: [u8; N]) {
        self.buffer.extend(bytes);
    }

    // #[inline]
    pub fn encode_u16(&mut self, value: u16) {
        self.encode_bytes(value.to_le_bytes());
    }

    // #[inline]
    pub fn encode_instr(&mut self, instr: Instruction) {
        self.encode_u16(instr as u16);
    }
}

pub struct InstructionPtr {
    ptr: *mut u8,
}

impl InstructionPtr {
    #[inline]
    pub fn decode_bytes<const N: usize>(&mut self) -> [u8; N] {
        let result = unsafe {
            std::ptr::read_unaligned::<[u8; N]>(self.ptr as _)
        };
        self.ptr = unsafe { self.ptr.add(N) };
        result
    }

    #[no_mangle]
    #[inline]
    pub fn decode_u16(&mut self) -> u16 {
        u16::from_le_bytes(self.decode_bytes::<2>())
    }

    #[no_mangle]
    #[inline]
    pub fn decode_i32(&mut self) -> i32 {
        i32::from_le_bytes(self.decode_bytes::<4>())
    }

    #[no_mangle]
    #[inline]
    pub fn decode_i64(&mut self) -> i64 {
        i64::from_le_bytes(self.decode_bytes::<8>())
    }

    #[no_mangle]
    pub fn decode_instr(&mut self) -> Instruction {
        Instruction::from(self.decode_u16())
    }

    #[no_mangle]
    pub fn decode_register(&mut self) -> Register {
        Register(self.decode_u16())
    }
}