stack_vm.rs

use num_derive::FromPrimitive;
use num_traits::FromPrimitive;
use std::collections::HashMap;

use Instruction::*;
#[derive(FromPrimitive, Clone, Debug)]
enum Instruction {
  // stack
  LIT,
  DROP,
  DUP,

  // ret stack
  RPUSH,
  RPOP,

  // arith
  ADD,
  SUB,
  MUL,

  // comparison
  LT,
  EQ,

  // jumps
  JMP,
  JMPT,
  JMPF,
  CALL,
  RET,

  // heap
  LOAD,
  STORE,

  //misc
  HALT,
  NOOP,
}

struct VM {
  stack: Vec<u32>,
  ret_stack: Vec<u32>,
  call_stack: Vec<u32>,
  heap: Vec<u32>,
  code: Vec<u8>,
  pc: u32,
  is_halted: bool,
}

impl VM {
  fn new(code: Vec<u8>) -> VM {
    VM {
      stack: Vec::new(),
      ret_stack: Vec::new(),
      call_stack: Vec::new(),
      heap: vec![0; 1000],
      code,
      pc: 0,
      is_halted: false,
    }
  }

  fn pop(&mut self) -> u32 {
    self.stack.pop().expect("empty stack on POP")
  }

  fn push(&mut self, val: u32) {
    self.stack.push(val);
  }

  fn rpush(&mut self) {
    let val = self.pop();
    self.ret_stack.push(val);
  }

  fn rpop(&mut self) {
    let val = self.ret_stack.pop().expect("empty ret_stack on POP");
    self.push(val);
  }

  fn consume(&mut self) -> u8 {
    let val = self.code[self.pc as usize];
    self.pc += 1;
    val
  }

  fn consume_u32(&mut self) -> u32 {
    u32_join(
      self.consume(),
      self.consume(),
      self.consume(),
      self.consume(),
    )
  }

  fn load(&mut self, addr: u32) -> u32 {
    self.heap[addr as usize]
  }

  fn store(&mut self, addr: u32, val: u32) {
    self.heap[addr as usize] = val;
  }

  fn step(&mut self) {
    let op = FromPrimitive::from_u8(self.consume()).expect("unknown opcode");
    println!(
      "OP: {:?},  \tSTACK: {:?},  \tRSTACK: {:?}",
      op, self.stack, self.ret_stack
    );
    match op {
      // stack
      LIT => {
        let val = self.consume_u32();
        self.push(val);
      }
      DROP => {
        self.pop();
        ()
      }
      DUP => {
        let val = self.pop();
        self.push(val);
        self.push(val)
      }

      // ret stack
      RPUSH => self.rpush(),
      RPOP => self.rpop(),

      // arith
      ADD => {
        let b = self.pop();
        let a = self.pop();
        self.push(a + b);
      }
      SUB => {
        let b = self.pop();
        let a = self.pop();
        self.push(a - b);
      }
      MUL => {
        let b = self.pop();
        let a = self.pop();
        self.push(a * b);
      }

      // comparison
      LT => {
        let b = self.pop();
        let a = self.pop();
        self.push(if a < b { 1 } else { 0 })
      }
      EQ => {
        let b = self.pop();
        let a = self.pop();
        self.push(if a == b { 1 } else { 0 })
      }

      // jumps
      JMP => {
        let dst = self.consume_u32();
        self.pc = dst;
      }
      JMPT => {
        let dst = self.consume_u32();
        let val = self.pop();
        if val != 0 {
          self.pc = dst;
        }
      }
      JMPF => {
        let dst = self.consume_u32();
        let val = self.pop();
        if val == 0 {
          self.pc = dst;
        }
      }
      CALL => {
        let dst = self.consume_u32();
        self.call_stack.push(self.pc);
        self.pc = dst;
      }
      RET => {
        let dst = self.call_stack.pop().expect("no frames in callstack");
        self.pc = dst;
      }

      // heap
      LOAD => {
        let addr = self.consume_u32();
        let val = self.load(addr);
        self.push(val);
      }
      STORE => {
        let addr = self.consume_u32();
        let val = self.pop();
        self.store(addr, val);
      }

      // misc
      HALT => self.is_halted = true,
      NOOP => (),
    }
  }

  fn run(&mut self) {
    while !self.is_halted {
      self.step();
    }
  }
}

use Builder::*;
#[derive(Clone)]
enum Builder {
  OP(Instruction),
  U32(u32),
  REF(String),
  LABEL(String),
}

fn u32_split(val: u32) -> (u8, u8, u8, u8) {
  (
    ((val >> 24) & 0xFF) as u8,
    ((val >> 16) & 0xFF) as u8,
    ((val >> 08) & 0xFF) as u8,
    ((val >> 00) & 0xFF) as u8,
  )
}

fn u32_join(a: u8, b: u8, c: u8, d: u8) -> u32 {
  (a as u32) << 24 | (b as u32) << 16 | (c as u32) << 8 | (d as u32)
}

fn push_u32(vec: &mut Vec<u8>, val: u32) {
  let (a, b, c, d) = u32_split(val);
  vec.push(a);
  vec.push(b);
  vec.push(c);
  vec.push(d);
}

impl Builder {
  fn build(bs: Vec<Builder>) -> Vec<u8> {
    let mut result = Vec::new();
    let mut labels = HashMap::new();
    let mut index: u32 = 0;

    for b in bs.clone() {
      match b {
        OP(_) => index += 1,
        REF(_) => index += 4,
        U32(_) => index += 4,
        LABEL(l) => {
          labels.insert(l, index);
          ()
        }
      }
    }

    for b in bs {
      match b {
        OP(i) => result.push(i as u8),
        U32(x) => push_u32(&mut result, x),
        REF(l) => push_u32(&mut result, *labels.get(&l).expect("unknown label")),
        LABEL(_) => (),
      }
    }

    result
  }
}

fn decompile(code: Vec<u8>) {
  let mut i = 0;
  while i < code.len() {
    match FromPrimitive::from_u8(code[i]).expect("unknown opcode") {
      LIT => {
        println!(
          "{:4}: LIT {}",
          i,
          u32_join(code[i + 1], code[i + 2], code[i + 3], code[i + 4])
        );
        i += 4;
      }
      DROP => println!("{:4}: DROP", i),
      DUP => println!("{:4}: DUP", i),
      RPUSH => println!("{:4}: RPUSH", i),
      RPOP => println!("{:4}: RPOP", i),
      ADD => println!("{:4}: ADD", i),
      SUB => println!("{:4}: SUB", i),
      MUL => println!("{:4}: MUL", i),
      LT => println!("{:4}: LT", i),
      EQ => println!("{:4}: EQ", i),
      JMP => {
        println!(
          "{:4}: JMP {}",
          i,
          u32_join(code[i + 1], code[i + 2], code[i + 3], code[i + 4])
        );
        i += 4;
      }
      JMPT => {
        println!(
          "{:4}: JMPT {}",
          i,
          u32_join(code[i + 1], code[i + 2], code[i + 3], code[i + 4])
        );
        i += 4;
      }
      JMPF => {
        println!(
          "{:4}: JMPF {}",
          i,
          u32_join(code[i + 1], code[i + 2], code[i + 3], code[i + 4])
        );
        i += 4;
      }
      CALL => {
        println!(
          "{:4}: CALL {}",
          i,
          u32_join(code[i + 1], code[i + 2], code[i + 3], code[i + 4])
        );
        i += 4;
      }
      RET => println!("{:4}: RET", i),
      LOAD => {
        println!(
          "{:4}: LOAD {}",
          i,
          u32_join(code[i + 1], code[i + 2], code[i + 3], code[i + 4])
        );
        i += 4;
      }
      STORE => {
        println!(
          "{:4}: STORE {}",
          i,
          u32_join(code[i + 1], code[i + 2], code[i + 3], code[i + 4])
        );
        i += 4;
      }
      HALT => println!("{:4}: HALT", i),
      NOOP => println!("{:4}: NOOP", i),
    }

    i += 1;
  }
}

fn main() {
  #[rustfmt::skip]
  let example = vec![
    OP(LIT), U32(10),
    OP(CALL), REF(String::from("FACT")),
    OP(STORE), U32(0),
    OP(HALT),
    
    LABEL(String::from("FACT")),
    OP(DUP),
    
    LABEL(String::from("FACT_LOOP")),
    OP(LIT), U32(1), OP(SUB),
    OP(DUP), OP(LIT), U32(0), OP(EQ), OP(JMPT), REF(String::from("FACT_RET")),
    OP(DUP), OP(RPUSH), OP(MUL), OP(RPOP),
    OP(JMP), REF(String::from("FACT_LOOP")),

    LABEL(String::from("FACT_RET")),
    OP(DROP),
    OP(RET),
  ];

  let compiled = Builder::build(example);

  println!("{:?}", decompile(compiled.clone()));

  let mut vm = VM::new(compiled);
  vm.run();

  println!("result: {}", vm.heap[0]);
}