fogtix/crates/fogtix-vm/src/lib.rs

use async_trait::async_trait;
use fogtix_bytecode::{consts, Atom, Instr, Parse, Pointer, Value as BcValue};
use std::sync::{Arc, RwLock};

pub type Module = Arc<dyn ModuleKind>;

pub trait ModuleKind: Send + Sync {
    fn as_slice(&self) -> &[u8];
}

impl<T: AsRef<[u8]> + Send + Sync + ?Sized> ModuleKind for T {
    #[inline(always)]
    fn as_slice(&self) -> &[u8] {
        self.as_ref()
    }
}

#[derive(Clone)]
pub struct InstrPtr {
    pub m: Module,
    pub pos: usize,
}

fn next_instr(m: &Module, pos: usize) -> Option<Result<(usize, Instr<'_>), &[u8]>> {
    m.as_slice()
        .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)),
        })
}

impl InstrPtr {
    #[inline(always)]
    pub fn next_instr(&self) -> Option<Result<(usize, Instr<'_>), &[u8]>> {
        next_instr(&self.m, self.pos)
    }

    fn is_call2jump(&self) -> bool {
        // tail-call optimization (would otherwise require much more opcodes)
        matches!(self.next_instr(), Some(Ok((_, Instr::Return))))
    }
}

#[async_trait]
pub trait Origin: Send + Sync + core::fmt::Debug {
    async fn call(&self, p: &Pointer, stack: &mut Vec<StackEntValue>) -> InstrPtr;
}

pub type KnownOrigins = Arc<[RwLock<Option<Box<dyn Origin>>>; 0x10000]>;

fn make_large_arc_slice<const N: usize, T, F: Fn() -> T>(f: F) -> Arc<[T; N]> {
    // see also: https://stackoverflow.com/a/68122278
    // license: CC BY-SA 4.0; author: "Johannes Maria Frank"
    let bxs = core::iter::repeat_with(f).take(N).collect::<Arc<[T]>>();
    let ptr = Arc::into_raw(bxs) as *mut [T; N];
    unsafe { Arc::from_raw(ptr) }
}

pub fn make_default_origins() -> KnownOrigins {
    make_large_arc_slice::<0x10000, RwLock<Option<Box<dyn Origin>>>, _>(|| RwLock::new(None))
}

#[derive(Clone, Debug)]
pub enum StackEntValue {
    Bytes(Vec<u8>),
    Int(u64),
    Atom(Atom),
    Pointer(Pointer),
}

#[derive(Debug, thiserror::Error)]
pub enum Error {
    #[error("`{0}` arrived at non-jump target `{1}` @ {2}")]
    InvalidJumpTarget(&'static str, String, usize),

    #[error("instruction pointer out-of-bounds")]
    InstrpOutOfBounds,

    #[error("reached unparsable instruction: {0:?}")]
    UnparsableInstruction(Vec<u8>),

    #[error("out of fuel")]
    OutOfFuel,

    #[error("not enough operands on stack")]
    NotEnoughStacked,

    #[error("call-r called with invalid/poisoned origin {0:x}")]
    InvalidOrigin(u16),

    #[error("operand from stack doesn't have correct data type: expected={expected}, got={got}")]
    StackedInvalidType { expected: &'static str, got: String },

    #[error("tried to divide by zero")]
    DivisionByZero,
}

pub struct Process {
    pub korigs: KnownOrigins,
    pub stack: Vec<StackEntValue>,
    pub callstack: Vec<InstrPtr>,
    pub instrp: InstrPtr,
}

fn verify_jumptarget_explicit(
    previptr: usize,
    jinstr: &'static str,
    jtip: &InstrPtr,
) -> Result<(), Error> {
    if let Some(Ok((_, trgi))) = jtip.next_instr() {
        if trgi != Instr::Label {
            return Err(Error::InvalidJumpTarget(
                jinstr,
                format!("{:?}", trgi),
                previptr,
            ));
        }
    }
    Ok(())
}

fn u128_binmath(a: u128, b: u128, mbo: consts::MathBinOp) -> Result<StackEntValue, Error> {
    use consts::MathBinOp as B;
    Ok(StackEntValue::Atom(Atom::from(match mbo {
        B::Concat => {
            return Err(Error::StackedInvalidType {
                expected: "bytes,bytes",
                got: "[int|atom, int|atom]".to_string(),
            })
        }
        B::Lt => return Ok(StackEntValue::Int(if a < b { 1 } else { 0 })),
        B::Eq => return Ok(StackEntValue::Int(if a < b { 1 } else { 0 })),
        B::And => a & b,
        B::Or => a | b,
        B::Xor => a ^ b,
        B::Add => a.wrapping_add(b),
        B::Sub => a.wrapping_sub(b),
        B::Mul => a.wrapping_mul(b),
        B::Div => match a.checked_div(b) {
            Some(x) => x,
            None => return Err(Error::DivisionByZero),
        },
        B::Rem => match a.checked_rem(b) {
            Some(x) => x,
            None => return Err(Error::DivisionByZero),
        },
    })))
}

impl Process {
    fn verify_jumptarget(&self, previptr: usize, jinstr: &'static str) -> Result<(), Error> {
        verify_jumptarget_explicit(previptr, jinstr, &self.instrp)
    }

    fn stpop(&mut self) -> Result<StackEntValue, Error> {
        self.stack.pop().ok_or(Error::NotEnoughStacked)
    }

    pub async fn run(&mut self, mut fuel: Option<&mut u64>) -> Result<(), Error> {
        loop {
            let previptr = self.instrp.pos;
            tracing::trace!("previptr = {}", previptr);
            if let Some(ref mut x) = fuel {
                if **x == 0 {
                    return Err(Error::OutOfFuel);
                }
                **x -= 1;
            }
            let (nxtidelta, nxti) = match next_instr(&self.instrp.m, self.instrp.pos) {
                None => return Err(Error::InstrpOutOfBounds),
                Some(Err(code)) => return Err(Error::UnparsableInstruction(code.to_vec())),
                Some(Ok(x)) => x,
            };
            self.instrp.pos += nxtidelta;
            match nxti {
                Instr::Label => {}
                Instr::CallRemote => {
                    if !self.instrp.is_call2jump() {
                        self.callstack.push(self.instrp.clone());
                    }
                    let mut ptr = match self.stpop()? {
                        StackEntValue::Pointer(ptr) => ptr,
                        x => {
                            return Err(Error::StackedInvalidType {
                                expected: "pointer",
                                got: format!("{:?}", x),
                            })
                        }
                    };
                    let argc = match self.stpop()? {
                        StackEntValue::Int(i) => i,
                        x => {
                            return Err(Error::StackedInvalidType {
                                expected: "int",
                                got: format!("{:?}", x),
                            })
                        }
                    };
                    let ssl = self.stack.len();
                    let argstart = match usize::try_from(argc)
                        .ok()
                        .and_then(|argc| ssl.checked_sub(argc))
                    {
                        Some(x) => x,
                        None => return Err(Error::NotEnoughStacked),
                    };
                    let mut args = self.stack.drain(argstart..).collect();
                    let origin_id = ptr.origin();
                    ptr.set_origin(0);
                    self.instrp = match self.korigs[usize::from(origin_id)].read() {
                        Ok(origin) => match origin.as_ref() {
                            Some(origin) => origin.call(&ptr, &mut args).await,
                            None => return Err(Error::InvalidOrigin(origin_id)),
                        },
                        Err(_) => return Err(Error::InvalidOrigin(origin_id)),
                    };
                    self.verify_jumptarget(previptr, "call-r")?;
                    let rargc = u64::try_from(args.len()).unwrap();
                    self.stack.extend(args);
                    self.stack.push(StackEntValue::Int(rargc));
                }
                Instr::CallLocal(x) => {
                    if !self.instrp.is_call2jump() {
                        self.callstack.push(self.instrp.clone());
                    }
                    self.instrp.pos = match x.try_into() {
                        Ok(y) => y,
                        Err(_) => return Err(Error::InstrpOutOfBounds),
                    };
                    self.verify_jumptarget(previptr, "call-l")?;
                }
                Instr::CallLDefer(x) => match x.try_into() {
                    Ok(pos) => {
                        let jtip = InstrPtr {
                            m: self.instrp.m.clone(),
                            pos,
                        };
                        verify_jumptarget_explicit(previptr, "call-l-defer", &jtip)?;
                        if self.instrp.is_call2jump() {
                            self.instrp = jtip;
                        } else {
                            self.callstack.push(jtip);
                        }
                    }
                    Err(_) => return Err(Error::InstrpOutOfBounds),
                },
                Instr::JumpCond(x) => {
                    let x: usize = match x.try_into() {
                        Ok(y) => y,
                        Err(_) => return Err(Error::InstrpOutOfBounds),
                    };
                    let doit = match self.stpop()? {
                        StackEntValue::Int(i) => i != 0,
                        StackEntValue::Atom(a) => a != Atom([0; 16]),
                        z => {
                            return Err(Error::StackedInvalidType {
                                expected: "atom | int",
                                got: format!("{:?}", z),
                            })
                        }
                    };
                    if doit {
                        self.instrp.pos = x;
                        self.verify_jumptarget(previptr, "jump-cond")?;
                    }
                }
                Instr::Return => match self.callstack.pop() {
                    Some(x) => self.instrp = x,
                    None => break Ok(()),
                },
                Instr::Push(v) => {
                    self.stack.push(match v {
                        BcValue::Bytes(v) => StackEntValue::Bytes(v.to_vec()),
                        BcValue::Int(i) => StackEntValue::Int(i),
                        BcValue::Atom(a) => StackEntValue::Atom(a),
                        BcValue::Pointer(p) => StackEntValue::Pointer(p),
                    });
                }
                Instr::Pop(_) if self.stack.is_empty() => return Err(Error::NotEnoughStacked),
                Instr::Pop(cnt) => {
                    let ssl = self.stack.len() - 1;
                    let cnt = usize::from(cnt);
                    if cnt >= ssl {
                        self.stack = Vec::new();
                    } else {
                        self.stack.truncate(ssl - cnt - 1);
                    }
                }
                Instr::Dup(delta) => {
                    let x = match self.stack.len().checked_sub(usize::from(delta) + 1) {
                        None => return Err(Error::NotEnoughStacked),
                        // SAFETY: the value x is always smaller than the stack height
                        Some(x) => self.stack[x].clone(),
                    };
                    self.stack.push(x);
                }
                Instr::Swap(delta) => {
                    let ssl = self.stack.len();
                    let (y, z) = match ssl.checked_sub(usize::from(delta) + 2) {
                        None => return Err(Error::NotEnoughStacked),
                        Some(ltrg) => self.stack[ltrg..].split_at_mut(1),
                    };
                    core::mem::swap(&mut y[0], &mut z[0]);
                }
                Instr::ADecon => match self.stpop()? {
                    StackEntValue::Atom(atom) => {
                        let (b, a) = atom.into();
                        self.stack.push(StackEntValue::Int(b));
                        self.stack.push(StackEntValue::Int(a));
                    }
                    x => {
                        return Err(Error::StackedInvalidType {
                            expected: "atom",
                            got: format!("{:?}", x),
                        })
                    }
                },
                Instr::DoMath1(ubo) => {
                    use consts::MathUnOp as U;
                    use StackEntValue as V;
                    let x = self.stpop()?;
                    self.stack.push(match (x, ubo) {
                        (V::Int(x), U::Not) => V::Int(!x),
                        (V::Atom(x), U::Not) => V::Atom(Atom::from(!u128::from(x))),
                        (V::Bytes(mut x), U::Not) => {
                            x.iter_mut().for_each(|i| *i = !*i);
                            V::Bytes(x)
                        }
                        (x, _) => {
                            return Err(Error::StackedInvalidType {
                                expected: "int|atom|bytes",
                                got: format!("{:?}", x),
                            })
                        }
                    });
                }
                Instr::DoMath2(mbo) => {
                    let b = self.stpop()?;
                    let a = self.stpop()?;
                    use consts::MathBinOp as B;
                    use StackEntValue as V;
                    self.stack.push(match (a, b) {
                        (V::Int(a), V::Int(b)) => match mbo {
                            B::Concat => V::Atom(Atom::from((a, b))),
                            B::Lt => V::Int(if a < b { 1 } else { 0 }),
                            B::Eq => V::Int(if a == b { 1 } else { 0 }),
                            B::And => V::Int(a & b),
                            B::Or => V::Int(a | b),
                            B::Xor => V::Int(a ^ b),
                            B::Add => V::Int(a.wrapping_add(b)),
                            B::Sub => V::Int(a.wrapping_sub(b)),
                            B::Mul => {
                                V::Atom(Atom::from(u128::from(a).wrapping_mul(u128::from(b))))
                            }
                            B::Div => V::Int(match a.checked_div(b) {
                                Some(x) => x,
                                None => return Err(Error::DivisionByZero),
                            }),
                            B::Rem => V::Int(match a.checked_rem(b) {
                                Some(x) => x,
                                None => return Err(Error::DivisionByZero),
                            }),
                        },
                        (V::Atom(a), V::Atom(b)) => {
                            u128_binmath(u128::from(a), u128::from(b), mbo)?
                        }
                        (V::Int(a), V::Atom(b)) => u128_binmath(u128::from(a), u128::from(b), mbo)?,
                        (V::Atom(a), V::Int(b)) => u128_binmath(u128::from(a), u128::from(b), mbo)?,
                        (V::Bytes(a), V::Bytes(b)) => match mbo {
                            B::Concat => {
                                let mut x = Vec::with_capacity(a.len() + b.len());
                                x.extend_from_slice(&a);
                                x.extend_from_slice(&b);
                                V::Bytes(x)
                            }
                            B::Eq => V::Int(if a == b { 1 } else { 0 }),
                            B::And => {
                                let mut x: Vec<_> =
                                    a.iter().zip(b.iter()).map(|(w, v)| w & v).collect();
                                let minp = x.len();
                                x.extend_from_slice(
                                    &(if a.len() < b.len() { b } else { a })[minp..],
                                );
                                V::Bytes(x)
                            }
                            B::Or => {
                                let mut x: Vec<_> =
                                    a.iter().zip(b.iter()).map(|(w, v)| w | v).collect();
                                let minp = x.len();
                                x.extend_from_slice(
                                    &(if a.len() < b.len() { b } else { a })[minp..],
                                );
                                V::Bytes(x)
                            }
                            B::Xor => {
                                let mut x: Vec<_> =
                                    a.iter().zip(b.iter()).map(|(w, v)| w ^ v).collect();
                                let minp = x.len();
                                x.extend_from_slice(
                                    &(if a.len() < b.len() { b } else { a })[minp..],
                                );
                                V::Bytes(x)
                            }
                            _ => {
                                return Err(Error::StackedInvalidType {
                                    expected: "[int|atom, int|atom]",
                                    got: "[bytes, bytes]".to_string(),
                                })
                            }
                        },
                        _ if mbo == B::Eq => V::Int(0),
                        (a, b) => {
                            return Err(Error::StackedInvalidType {
                                expected: "[int|atom, int|atom] | [bytes,bytes]",
                                got: format!("{:?}, {:?}", a, b),
                            })
                        }
                    });
                }
            }
        }
    }
}

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

    #[test]
    fn stack_item_size() {
        assert_eq!(core::mem::size_of::<StackEntValue>(), 32);
    }

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

        #[test]
        fn doesnt_crash(inp in proptest::collection::vec(0..=u8::MAX, 0..1024)) {
            use once_cell::sync::Lazy;
            // we use a static here to avoid allocating the very large array on every iteration
            static KORIGS: Lazy<KnownOrigins> = Lazy::new(|| make_default_origins());
            let inp: Arc<Vec<u8>> = Arc::new(inp);
            let module: Module = inp;
            let mut p = Process {
                korigs: Arc::clone(&*KORIGS),
                stack: Vec::new(),
                callstack: Vec::new(),
                instrp: InstrPtr {
                    m: module,
                    pos: 0,
                },
            };
            let _ = futures_lite::future::block_on(p.run(Some(&mut 2048)));
        }
    }
}