Merge pull request #4 from slothlang/vm-basics

Basic VM

1 files changed, 599 insertions, 0 deletions

diff --git a/crates/sloth_vm/src/vm.rs b/crates/sloth_vm/src/vm.rs
new file mode 100644
index 0000000..2d58c3d
--- /dev/null
+++ b/crates/sloth_vm/src/vm.rs
@@ -0,0 +1,599 @@
+use std::mem::MaybeUninit;
+
+use sloth_bytecode::Opcode;
+
+use crate::value::{Function, Object, ObjectType, Primitive};
+use crate::{native, ObjectMap, Stack};
+
+#[derive(Clone, Copy)]
+pub struct CallFrame {
+    pointer: usize,
+    stack_offset: usize,
+    function: *const Function, // TODO: Safety
+}
+
+impl CallFrame {
+    fn new(stack_offset: usize, function: &Function) -> Self {
+        Self {
+            pointer: 0,
+            stack_offset,
+            function: function as *const _,
+        }
+    }
+
+    #[inline]
+    fn function(&self) -> &Function {
+        unsafe { &*self.function }
+    }
+}
+
+const CALL_STACK_SIZE: usize = 1024;
+
+pub struct CallStack {
+    top: usize,
+    frames: [MaybeUninit<CallFrame>; CALL_STACK_SIZE],
+}
+
+impl Default for CallStack {
+    fn default() -> Self {
+        Self {
+            top: 0,
+            frames: [MaybeUninit::uninit(); CALL_STACK_SIZE],
+        }
+    }
+}
+
+impl CallStack {
+    fn push(&mut self, frame: CallFrame) {
+        self.frames[self.top].write(frame);
+        self.top += 1;
+    }
+
+    fn pop(&mut self) {
+        self.top -= 1;
+    }
+
+    fn peek(&self) -> &CallFrame {
+        unsafe { self.frames[self.top - 1].assume_init_ref() }
+    }
+
+    fn peek_mut(&mut self) -> &mut CallFrame {
+        unsafe { self.frames[self.top - 1].assume_init_mut() }
+    }
+}
+
+pub struct VM {
+    stack: Stack,
+    call_stack: CallStack,
+    objects: ObjectMap,
+}
+
+impl Default for VM {
+    fn default() -> Self {
+        Self::init(ObjectMap::default())
+    }
+}
+
+impl VM {
+    pub fn init(objects: ObjectMap) -> Self {
+        Self {
+            stack: Stack::default(),
+            call_stack: CallStack::default(),
+            objects,
+        }
+    }
+
+    pub fn new(objects: ObjectMap, mut root: Function) -> Self {
+        let mut this = Self::init(objects);
+
+        // Allocating the root function
+        root.chunk.code.push(Opcode::Halt as u8);
+        this.call_stack.push(CallFrame::new(0, &root));
+        this.objects
+            .allocate(Object::new(ObjectType::Function(root)));
+
+        this
+    }
+
+    pub fn step(&mut self) -> bool {
+        use Primitive::*;
+
+        let opcode = self.read_u8();
+
+        match Opcode::from_u8(opcode) {
+            Opcode::Constant => {
+                let idx = self.read_u16() as usize;
+                let value = self.call_stack.peek().function().chunk.constants[idx];
+
+                self.stack.push(value);
+            }
+
+            Opcode::Dup => {
+                let value = self.stack.pop();
+                self.stack.push(value);
+                self.stack.push(value);
+            }
+            Opcode::Pop => {
+                self.stack.pop();
+            }
+            Opcode::GetLocal => {
+                let idx = self.read_u16() as usize;
+                let value = self.stack[self.call_stack.peek().stack_offset + idx];
+
+                self.stack.push(value);
+            }
+            Opcode::SetLocal => {
+                let idx = self.read_u16() as usize;
+                let value = self.stack.pop();
+
+                self.stack[self.call_stack.peek().stack_offset + idx] = value;
+            }
+
+            Opcode::Add => {
+                let value = match self.stack.pop2() {
+                    (Integer(lhs), Integer(rhs)) => Integer(lhs + rhs),
+                    (Float(lhs), Float(rhs)) => Float(lhs + rhs),
+                    _ => panic!(),
+                };
+
+                self.stack.push(value);
+            }
+            Opcode::Sub => {
+                let value = match self.stack.pop2() {
+                    (Integer(lhs), Integer(rhs)) => Integer(lhs - rhs),
+                    (Float(lhs), Float(rhs)) => Float(lhs - rhs),
+                    _ => panic!(),
+                };
+
+                self.stack.push(value);
+            }
+            Opcode::Mul => {
+                let value = match self.stack.pop2() {
+                    (Integer(lhs), Integer(rhs)) => Integer(lhs * rhs),
+                    (Float(lhs), Float(rhs)) => Float(lhs * rhs),
+                    _ => panic!(),
+                };
+
+                self.stack.push(value);
+            }
+            Opcode::Div => {
+                let value = match self.stack.pop2() {
+                    (Integer(_), Integer(0)) => panic!("Divide by 0"),
+                    (Integer(lhs), Integer(rhs)) => Integer(lhs / rhs),
+                    (Float(lhs), Float(rhs)) => Float(lhs / rhs),
+                    _ => panic!(),
+                };
+
+                self.stack.push(value);
+            }
+            Opcode::Mod => {
+                let value = match self.stack.pop2() {
+                    (Integer(lhs), Integer(rhs)) => Integer(lhs % rhs),
+                    (Float(lhs), Float(rhs)) => Float(lhs % rhs),
+                    _ => panic!(),
+                };
+
+                self.stack.push(value);
+            }
+
+            Opcode::Eq => {
+                let value = match self.stack.pop2() {
+                    (Integer(lhs), Integer(rhs)) => Bool(lhs == rhs),
+                    (Float(lhs), Float(rhs)) => Bool(lhs == rhs),
+                    (Bool(lhs), Bool(rhs)) => Bool(lhs == rhs),
+                    (Object(lhs), Object(rhs)) => Bool(lhs == rhs),
+                    (Empty, Empty) => Bool(true),
+                    _ => Bool(false),
+                };
+
+                self.stack.push(value);
+            }
+            Opcode::Ne => {
+                let value = match self.stack.pop2() {
+                    (Integer(lhs), Integer(rhs)) => Bool(lhs != rhs),
+                    (Float(lhs), Float(rhs)) => Bool(lhs != rhs),
+                    (Bool(lhs), Bool(rhs)) => Bool(lhs != rhs),
+                    (Object(lhs), Object(rhs)) => Bool(lhs != rhs),
+                    (Empty, Empty) => Bool(false),
+                    _ => Bool(false),
+                };
+
+                self.stack.push(value);
+            }
+
+            Opcode::Jump => {
+                let to = self.read_u16();
+                self.call_stack.peek_mut().pointer = to as usize;
+            }
+            Opcode::JumpIf => {
+                let to = self.read_u16();
+                let value = self.stack.pop();
+
+                if let Bool(true) = value {
+                    self.call_stack.peek_mut().pointer = to as usize;
+                }
+            }
+
+            Opcode::Call => {
+                let Primitive::Object(ptr) = self.stack.pop() else {
+                    panic!("Last element on stack was not an object");
+                };
+
+                self.call(ptr as usize);
+            }
+
+            Opcode::Return => {
+                self.call_return();
+            }
+
+            Opcode::Halt => return false,
+
+            opcode => unimplemented!("Opcode {:?} unimplemented", opcode),
+        }
+
+        true
+    }
+
+    pub fn run(&mut self) {
+        while self.step() {}
+    }
+
+    pub fn call(&mut self, ptr: usize) {
+        let Some(obj) = self.objects.get(ptr) else {
+            panic!("Pointer referenced nothing");
+        };
+
+        match &obj.typ {
+            ObjectType::Function(function) => {
+                // Add a callstack entry for the function
+                let offset = self.stack.top - (function.arity as usize);
+                self.call_stack.push(CallFrame::new(offset, function));
+            }
+            ObjectType::NativeFunction(function) => {
+                let mut args = Vec::with_capacity(function.arity as usize);
+                for _ in 0..function.arity {
+                    args.push(self.stack.pop());
+                }
+
+                let name = function.name;
+                let returns_value = function.returns_value;
+
+                let internal = function.function;
+                let result = internal(self, &args);
+
+                match result {
+                    Ok(value) => {
+                        if returns_value {
+                            self.stack.push(value);
+                        }
+                    }
+                    Err(error) => match error {
+                        native::Error::InvalidArgument => {
+                            panic!("Invalid argument provided to '{name}'");
+                        }
+                        native::Error::Unknown(msg) => {
+                            panic!("Native function '{name}' failed due to '{msg}'");
+                        }
+                    },
+                }
+            }
+            _ => panic!("Object was not a function"),
+        }
+    }
+
+    fn call_return(&mut self) {
+        let function = self.call_stack.peek().function();
+        let stack_offset = self.call_stack.peek().stack_offset;
+
+        let return_value = if function.returns_value {
+            Some(self.stack.pop())
+        } else {
+            None
+        };
+
+        self.stack.top = stack_offset;
+
+        if let Some(return_value) = return_value {
+            self.stack.push(return_value);
+        }
+
+        self.call_stack.pop();
+    }
+
+    fn unwind(&mut self) {
+        unimplemented!("Implement unwinding for error handling");
+    }
+
+    #[inline(always)]
+    fn read_u8(&mut self) -> u8 {
+        let frame = self.call_stack.peek_mut();
+        let function = frame.function();
+        let byte = function.chunk.code[frame.pointer];
+        frame.pointer += 1;
+        byte
+    }
+
+    #[inline(always)]
+    fn read_u16(&mut self) -> u16 {
+        let frame = self.call_stack.peek_mut();
+        let chunk = &frame.function().chunk;
+
+        let bytes = (chunk.code[frame.pointer], chunk.code[frame.pointer + 1]);
+
+        frame.pointer += 2;
+
+        ((bytes.0 as u16) << 8) + (bytes.1 as u16)
+    }
+
+    #[inline(always)]
+    pub fn objects(&self) -> &ObjectMap {
+        &self.objects
+    }
+
+    #[inline(always)]
+    pub fn objects_mut(&mut self) -> &mut ObjectMap {
+        &mut self.objects
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::value::{Function, Object, ObjectType, Primitive};
+    use crate::{sloth_std, Chunk, ObjectMap, VM};
+
+    #[test]
+    fn arithmetic_ops() {
+        // Addition
+        let mut vm = VM::new(
+            ObjectMap::default(),
+            Function::root(Chunk {
+                constants: vec![Primitive::Integer(7)],
+                code: vec![
+                    0x00, 0, 0,    // Load constant from 0
+                    0x10, // Duplicate
+                    0x20, // Add
+                    0xE0,
+                ],
+            }),
+        );
+
+        vm.run();
+        assert_eq!(vm.stack.peek(), Primitive::Integer(14));
+
+        let mut vm = VM::new(
+            ObjectMap::default(),
+            Function::root(Chunk {
+                constants: vec![Primitive::Integer(2), Primitive::Integer(11)],
+                code: vec![
+                    0x00, 0, 0, // Load constant from 0
+                    0x00, 0, 1,    // Load constant from 1
+                    0x20, // Add
+                    0xE0,
+                ],
+            }),
+        );
+
+        vm.run();
+        assert_eq!(vm.stack.peek(), Primitive::Integer(13));
+    }
+
+    #[test]
+    fn basic_function() {
+        let mut vm = VM::new(
+            ObjectMap::from(vec![Object::new(ObjectType::Function(Function {
+                name: Some("add".to_string()),
+                chunk: Chunk {
+                    constants: vec![],
+                    code: vec![0x14, 0, 0, 0x14, 0, 1, 0x20, 0x52],
+                },
+                arity: 2,
+                returns_value: true,
+            }))]),
+            Function::root(Chunk {
+                constants: vec![
+                    Primitive::Integer(6),
+                    Primitive::Integer(3),
+                    Primitive::Object(0),
+                    Primitive::Object(1),
+                    Primitive::Object(2),
+                ],
+                code: vec![
+                    0x00, 0, 0, // Load first function parameter from 0
+                    0x00, 0, 1, // Load second function parameter from 1
+                    0x00, 0, 2,    // Load function constant from 2
+                    0x50, // Call function
+                ],
+            }),
+        );
+
+        vm.run();
+
+        assert_eq!(vm.stack.peek(), Primitive::Integer(9));
+    }
+
+    #[test]
+    fn native_function() {
+        let mut vm = VM::new(
+            ObjectMap::from(vec![
+                Object::new(ObjectType::NativeFunction(sloth_std::rand::GEN_FUNCTION)),
+                Object::new(ObjectType::NativeFunction(
+                    sloth_std::rand::GEN_RANGE_FUNCTION,
+                )),
+            ]),
+            Function::root(Chunk {
+                constants: vec![
+                    Primitive::Object(0),
+                    Primitive::Object(1),
+                    Primitive::Integer(5),
+                    Primitive::Integer(10),
+                ],
+                code: vec![
+                    // First part
+                    0x00, 0, 0,    //
+                    0x50, //
+                    0xE0, //
+                    // Second part
+                    0x00, 0, 3, //
+                    0x00, 0, 2, //
+                    0x00, 0, 1,    //
+                    0x50, //
+                ],
+            }),
+        );
+
+        vm.run();
+
+        assert!({
+            let Primitive::Float(i) = vm.stack.peek() else { panic!(); };
+            (0.0..=1.0).contains(&i)
+        });
+
+        vm.run();
+
+        assert!({
+            let Primitive::Integer(i) = vm.stack.peek() else { panic!(); };
+            (5..10).contains(&i)
+        });
+    }
+
+    #[test]
+    fn fibonacci() {
+        #[rustfmt::skip]
+        let mut vm = VM::new(
+            ObjectMap::default(),
+            Function::root(Chunk {
+                constants: vec![
+                    Primitive::Integer(0),
+                    Primitive::Integer(1),
+                    Primitive::Integer(10),
+                ],
+                code: vec![
+                    // Load variables
+                    0x00, 0, 0,  // 0  Index
+                    0x00, 0, 0,  // 3  Me
+                    0x00, 0, 0,  // 6  Parent
+                    0x00, 0, 1,  // 9  Grandparent
+
+                    // Load parent and grandparent, sum them and put the value in me
+                    0x14, 0, 2,  // 12
+                    0x14, 0, 3,  // 15
+                    0x20,        // 16
+                    0x15, 0, 1,  // 19
+                    
+                    // Set grandparent to parent
+                    0x14, 0, 2,  // 22
+                    0x15, 0, 3,  // 25
+                    
+                    // Set parent to me
+                    0x14, 0, 1,  // 28
+                    0x15, 0, 2,  // 31
+
+                    // Increment Index by 1
+                    0x00, 0, 1,  // 34
+                    0x14, 0, 0,  // 37 Index
+                    0x20,        // 40
+                    0x15, 0, 0,  // 41 Index
+
+                    // Load me
+                    0x14, 0, 1,  // 44
+                    0xE0,        // 47
+                    0x11,        // 48
+
+                    // Repeat until Index is 9
+                    0x00, 0, 2,  // 49 
+                    0x14, 0, 0,  // 52 Index
+                    0x31,        // 55
+                    0x41, 0, 12, // 56
+                ],
+            }),
+        );
+
+        let mut values = Vec::new();
+        for _ in 0..10 {
+            vm.run();
+            values.push(vm.stack.peek());
+        }
+
+        assert_eq!(&values, &[
+            Primitive::Integer(1),
+            Primitive::Integer(1),
+            Primitive::Integer(2),
+            Primitive::Integer(3),
+            Primitive::Integer(5),
+            Primitive::Integer(8),
+            Primitive::Integer(13),
+            Primitive::Integer(21),
+            Primitive::Integer(34),
+            Primitive::Integer(55),
+        ]);
+    }
+
+    #[test]
+    fn fibonacci_recursive() {
+        #[rustfmt::skip]
+        let mut vm = VM::new(
+            ObjectMap::from(vec![Object::new(ObjectType::Function(Function {
+                name: Some("fib".to_owned()),
+                chunk: Chunk {
+                    constants: vec![
+                        Primitive::Object(0),
+                        Primitive::Integer(0),
+                        Primitive::Integer(1),
+                        Primitive::Integer(2),
+                    ],
+                    code: vec![
+                        0x14, 0, 0,  // 0
+                        0x00, 0, 1,  // 3
+                        0x31,        // 6
+                        0x41, 0, 14, // 7
+                        0x00, 0, 1,  // 10
+                        0x52,        // 13
+
+                        0x14, 0, 0,  // 14
+                        0x00, 0, 2,  // 17
+                        0x31,        // 20
+                        0x41, 0, 28, // 21
+                        0x00, 0, 2,  // 24
+                        0x52,        // 27
+
+                        // fib(n - 1)
+                        0x00, 0, 2,  // 28
+                        0x14, 0, 0,  // 31
+                        0x21,        // 34
+                        0x00, 0, 0,  // 35
+                        0x50,        // 38
+                        
+                        // fib(n - 2)
+                        0x00, 0, 3,  // 39
+                        0x14, 0, 0,  // 42
+                        0x21,        // 45
+                        0x00, 0, 0,  // 46
+                        0x50,        // 49
+
+                        // add & return
+                        0x20,        // 50
+                        0x52,        // 51
+                    ],
+                },
+                arity: 1,
+                returns_value: true,
+            }))]),
+            Function::root(Chunk {
+                constants: vec![
+                    Primitive::Object(0),
+                    Primitive::Integer(10),
+                ],
+                code: vec![
+                    // Load n and the function and call it
+                    0x00, 0, 1, // 0
+                    0x00, 0, 0, // 3
+                    0x50,       // 6
+                ],
+            }),
+        );
+
+        vm.run();
+
+        assert_eq!(Primitive::Integer(55), vm.stack.peek());
+    }
+}