rasp/src/vm/mem.rs

use crate::vm::{addr::*, error::*, inst::*, reg::*};
use std::{convert::TryInto, ops::Index, mem};

pub struct MemCursor<T> {
    pos: Addr,
    mem: T,
}

impl<T> MemCursor<T>
    where T: AsRef<[u8]>
{
    pub fn new(mem: T) -> Self {
        MemCursor { pos: Addr(0), mem }
    }

    pub fn position(&self) -> Addr {
        self.pos
    }

    pub fn set_position<P: Into<Addr>>(&mut self, position: P) {
        self.pos = position.into();
    }

    pub fn next_u8_unchecked(&mut self) -> u8 {
        let val = self[self.pos];
        self.pos += 1u64;
        val
    }

    pub fn next_u8(&mut self) -> Result<u8> {
        self.check_addr(self.pos)
            .map(|_| self.next_u8_unchecked())
    }

    pub fn next_u16_unchecked(&mut self) -> u16 {
        let (int_bytes, _) = self.mem.as_ref()
            .split_at(mem::size_of::<u16>());
        let val = u16::from_le_bytes(int_bytes.try_into().unwrap());
        self.pos += 2u64;
        val
    }

    pub fn next_u16(&mut self) -> Result<u16> {
        self.check_addr(self.pos)
            .map(|_| self.next_u16_unchecked())
    }

    pub fn next_u32_unchecked(&mut self) -> u32 {
        let (int_bytes, _) = self.mem.as_ref()
            .split_at(mem::size_of::<u32>());
        let val = u32::from_le_bytes(int_bytes.try_into().unwrap());
        self.pos += 4u64;
        val
    }

    pub fn next_u32(&mut self) -> Result<u32> {
        self.check_addr(self.pos)
            .map(|_| self.next_u32_unchecked())
    }
    
    pub fn next_u64_unchecked(&mut self) -> u64 {
        let (int_bytes, _) = self.mem.as_ref()
            .split_at(mem::size_of::<u64>());
        let val = u64::from_le_bytes(int_bytes.try_into().unwrap());
        self.pos += 8u64;
        val
    }

    pub fn next_u64(&mut self) -> Result<u64> {
        self.check_addr(self.pos)
            .map(|_| self.next_u64_unchecked())
    }

    pub fn next_addr(&mut self) -> Result<Addr> {
        self.check_addr(self.pos)
            .map(|_| self.next_addr_unchecked())
    }

    pub fn next_addr_unchecked(&mut self) -> Addr {
        Addr(self.next_u64_unchecked())
    }

    pub fn next_inst(&mut self) -> Result<Inst> {
        let op = self.next_u16()?;

        macro_rules! dest_source {
            ($variant:ident) => {{
                let (d, s) = self.next_dest_source()?;
                Ok(Inst::$variant(d, s))
            }};
        }
        macro_rules! source_source {
            ($variant:ident) => {{
                let (s1, s2) = self.next_source_source()?;
                Ok(Inst::$variant(s1, s2))
            }};
        }
        macro_rules! source {
            ($variant:ident) => {{
                let spec = (self.next_u8()? & 0xF0) >> 4;
                let source = self.next_source(spec)?;
                Ok(Inst::$variant(source))
            }};
        }
        match op {
            ADD => dest_source!(Add),
            SUB => dest_source!(Sub),
            MUL => dest_source!(Mul),
            DIV => dest_source!(Div),
            MOD => dest_source!(Mod),
            AND => dest_source!(And),
            OR => dest_source!(Or),
            XOR => dest_source!(Xor),
            SHL => dest_source!(Shl),
            SHR => dest_source!(Shr),
            INEG => dest_source!(INeg),
            INV => dest_source!(Inv),
            NOT => dest_source!(Not),
            CMPEQ => source_source!(CmpEq),
            CMPLT => source_source!(CmpLt),
            JMP => source!(Jmp),
            JZ => source!(Jz),
            JNZ => source!(Jnz),
            MOV => dest_source!(Mov),
            HALT => Ok(Inst::Halt),
            NOP => Ok(Inst::Nop),
            DUMP => Ok(Inst::Dump),
            _ => Err(VmError::IllegalOp { op }),
        }
    }

    fn next_source_source(&mut self) -> Result<(Source, Source)> {
        let spec = self.next_u8()?;
        let s1_spec = (spec & 0xF0) >> 4;
        let s2_spec = spec & 0x0F;
        let s1 = self.next_source(s1_spec)?;
        let s2 = self.next_source(s2_spec)?;
        Ok((s1, s2))
    }

    fn next_dest_source(&mut self) -> Result<(Dest, Source)> {
        let spec = self.next_u8()?;
        let dest_spec = (spec & 0xF0) >> 4;
        let source_spec = spec & 0x0F;
        let dest = self.next_dest(dest_spec)?;
        let source = self.next_source(source_spec)?;
        Ok((dest, source))
    }

    fn next_dest(&mut self, spec: u8) -> Result<Dest> {
        match spec {
            DEST_ADDR64 => Ok(Dest::Addr64(self.next_addr()?)),
            DEST_ADDR32 => Ok(Dest::Addr32(self.next_addr()?)),
            DEST_ADDR16 => Ok(Dest::Addr16(self.next_addr()?)),
            DEST_ADDR8 => Ok(Dest::Addr8(self.next_addr()?)),
            DEST_REG => Ok(Dest::Reg(self.next_reg()?)),
            _ => Err(VmError::IllegalDestSpec { spec }),
        }
    }

    fn next_source(&mut self, spec: u8) -> Result<Source> {
        match spec {
            SOURCE_ADDR64 => Ok(Source::Addr64(self.next_addr()?)),
            SOURCE_ADDR32 => Ok(Source::Addr32(self.next_addr()?)),
            SOURCE_ADDR16 => Ok(Source::Addr16(self.next_addr()?)),
            SOURCE_ADDR8 => Ok(Source::Addr8(self.next_addr()?)),
            SOURCE_REG => Ok(Source::Reg(self.next_reg()?)),
            SOURCE_IMM64 => Ok(Source::Imm(self.next_u64()?)),
            SOURCE_IMM32 => Ok(Source::Imm(self.next_u32()? as u64)),
            SOURCE_IMM16 => Ok(Source::Imm(self.next_u16()? as u64)),
            SOURCE_IMM8 => Ok(Source::Imm(self.next_u8()? as u64)),
            _ => Err(VmError::IllegalSourceSpec { spec }),
        }
    }

    fn next_reg(&mut self) -> Result<Reg> {
        let reg = self.next_u8()?;
        if (reg as usize) >= NUM_REGS {
            Err(VmError::IllegalReg { reg })
        } else {
            Ok(reg)
        }
    }

    fn check_addr(&self, addr: Addr) -> Result<()> {
        if addr > self.mem.as_ref().len() {
            Err(VmError::MemOutOfBounds { addr })
        } else {
            Ok(())
        }
    }
}

////////////////////////////////////////////////////////////////////////////////
// Index impl
////////////////////////////////////////////////////////////////////////////////
impl<T: AsRef<[u8]>> Index<usize> for MemCursor<T> {
    type Output = u8;

    fn index(&self, addr: usize) -> &Self::Output {
        self.mem.as_ref().index(addr)
    }
}

impl<T: AsRef<[u8]>> Index<u64> for MemCursor<T> {
    type Output = u8;

    fn index(&self, addr: u64) -> &Self::Output {
        self.index(addr as usize)
    }
}

impl<T: AsRef<[u8]>> Index<Addr> for MemCursor<T> {
    type Output = u8;

    fn index(&self, addr: Addr) -> &Self::Output {
        self.index(addr.0)
    }
}