fogtix/crates/fogtix-bytecode/src/instr.rs

use crate::consts::{MathBinOp, MathUnOp, OpType};
use core::fmt;

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Instr {
    // control flow
    /// pops the descriptor table pointer from the stack,
    /// then hands execution off to the descriptor table
    CallRemote(u16),

    /// calls the destination $0 inside of the current module
    CallLocal(u64),

    /// calls the destination $0 inside of the current module after the
    /// next `Return`, by inserting the call into the call stack,
    /// but continuing execution normally.
    /// A classic `Jump(i)` is then just `[CallLDefer(i), Return]`.
    CallLDefer(u64),

    /// pops the top stack value and if it is non-zero,
    /// jumps to the specified destination inside of the current module
    JumpCond(u64),

    /// pops the current call frame
    Return,

    // basic stack operations
    /// pushes the given value to the stack
    Push(u64),

    /// pops the top $0+1 values from the stack
    Pop(u16),

    /// duplicates the top-$0 stack value, pushing it to the top of the stack
    Dup(u16),

    /// swaps the top stack value with the top-1-$0 stack value
    /// (`Swap(0)` swaps the two top-most stack values)
    Swap(u16),

    /// shifts the top stack value by $0 ($0 is a signed integer)
    Shift(i8),

    // indirect stack operations
    /// duplicates the top-(*top) stack value, pushing it to the top of the stack
    /// this does not change the stack height
    /// (the top value is replaced by the value it points to)
    DupFrom,

    /// basic math operations
    DoMath1(MathUnOp),

    /// binary math operations, stack top is like: `... lhs rhs $`
    DoMath2(MathBinOp),
}

impl Instr {
    #[inline]
    pub fn typ(&self) -> OpType {
        match self {
            Instr::CallRemote(_) => OpType::CallRemote,
            Instr::CallLocal(_) => OpType::CallLocal,
            Instr::CallLDefer(_) => OpType::CallLDefer,
            Instr::JumpCond(_) => OpType::JumpCond,
            Instr::Return => OpType::Return,
            Instr::Push(_) => OpType::Push,
            Instr::Pop(_) => OpType::Pop,
            Instr::Dup(_) => OpType::Dup,
            Instr::Swap(_) => OpType::Swap,
            Instr::Shift(_) => OpType::Shift,
            Instr::DupFrom => OpType::DupFrom,
            Instr::DoMath1(_) => OpType::DoMath1,
            Instr::DoMath2(_) => OpType::DoMath2,
        }
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct ParseError;

impl fmt::Display for ParseError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "unable to parse instruction from buffer")
    }
}

#[cfg(feature = "std")]
impl std::error::Error for ParseError {}

impl From<()> for ParseError {
    fn from(_: ()) -> Self {
        Self
    }
}

impl<'a> crate::Parse<'a> for Instr {
    type Err = ParseError;

    fn parse(inp: &'a [u8]) -> Result<(&'a [u8], Self), ParseError> {
        if inp.is_empty() {
            return Err(ParseError);
        }
        let (inp, otyp) = OpType::parse(inp)?;
        use crate::ParseExt;
        match otyp {
            OpType::CallRemote => u16::parse_map(inp, Instr::CallRemote),
            OpType::CallLocal => u64::parse_map(inp, Instr::CallLocal),
            OpType::CallLDefer => u64::parse_map(inp, Instr::CallLDefer),
            OpType::JumpCond => u64::parse_map(inp, Instr::JumpCond),
            OpType::Push => u64::parse_map(inp, Instr::Push),
            OpType::Pop => u16::parse_map(inp, Instr::Pop),
            OpType::Dup => u16::parse_map(inp, Instr::Dup),
            OpType::Swap => u16::parse_map(inp, Instr::Swap),
            OpType::Shift => i8::parse_map(inp, Instr::Shift),
            OpType::DoMath1 => MathUnOp::parse_map(inp, Instr::DoMath1),
            OpType::DoMath2 => MathBinOp::parse_map(inp, Instr::DoMath2),
            OpType::Return => Ok((inp, Instr::Return)),
            OpType::DupFrom => Ok((inp, Instr::DupFrom)),
        }
        .map_err(|()| ParseError)
    }
}

#[cfg(any(test, feature = "std"))]
impl Instr {
    pub fn write_to<W: std::io::Write>(&self, mut writer: W) -> std::io::Result<()> {
        use int_enum::IntEnum;
        writer.write_all(&self.typ().int_value().to_be_bytes())?;
        match self {
            Instr::CallRemote(val) => writer.write_all(&val.to_be_bytes()),
            Instr::CallLocal(val) | Instr::CallLDefer(val) | Instr::JumpCond(val) => {
                writer.write_all(&val.to_be_bytes())
            }
            Instr::Push(val) => writer.write_all(&val.to_be_bytes()),
            Instr::Pop(val) | Instr::Dup(val) | Instr::Swap(val) => {
                writer.write_all(&val.to_be_bytes())
            }
            Instr::Shift(shv) => writer.write_all(&shv.to_be_bytes()),
            Instr::DoMath1(val) => writer.write_all(&val.int_value().to_be_bytes()),
            Instr::DoMath2(val) => writer.write_all(&val.int_value().to_be_bytes()),
            Instr::Return | Instr::DupFrom => Ok(()),
        }
    }
}

pub fn next_instr(m: &[u8], pos: usize) -> Option<Result<(usize, Instr), &[u8]>> {
    use crate::Parse;
    m.get(pos..).map(|nxti_arr| match Instr::parse(nxti_arr) {
        Err(_) => Err(&nxti_arr[..core::cmp::min(nxti_arr.len(), 20)]),
        Ok((ptr, i)) => Ok((nxti_arr.len() - ptr.len(), i)),
    })
}

pub fn is_call2jump(m: &[u8], pos: usize) -> bool {
    // tail-call optimization (would otherwise require much more opcodes)
    matches!(next_instr(m, pos), Some(Ok((_, Instr::Return))))
}

#[cfg(test)]
mod tests {
    use super::*;

    proptest::proptest! {
        #![proptest_config(proptest::prelude::ProptestConfig::with_cases(4096))]

        #[test]
        fn doesnt_crash(inp in proptest::collection::vec(0..=u8::MAX, 0..1024)) {
            let mut inp = &inp[..];
            while let Ok((xinp, v)) = <Instr as crate::Parse<'_>>::parse(inp) {
                let mut buf = alloc::vec::Vec::with_capacity(inp.len() - xinp.len());
                v.write_to(&mut buf).unwrap();
                assert_eq!(buf[..], inp[..inp.len() - xinp.len()]);
                inp = xinp;
            }
        }
    }
}