From 6a8eb8938168a8e7a0cac57841108f88ee937355 Mon Sep 17 00:00:00 2001
From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com>
Date: Tue, 18 Jan 2022 21:43:16 +0100
Subject: [PATCH] return try

---
 src/bytecode.rs                             |  3 +-
 src/compile.rs                              |  2 +-
 src/vm.rs                                   | 43 +++++++++++++++++++--
 test.dil                                    | 13 +++++--
 tests/common.rs                             |  3 +-
 tests/control_flow.rs                       |  1 -
 tests/functions.rs                          | 25 ++++++++++++
 tests/snapshots/functions__single_call.snap |  7 ++++
 8 files changed, 86 insertions(+), 11 deletions(-)
 create mode 100644 tests/functions.rs
 create mode 100644 tests/snapshots/functions__single_call.snap

diff --git a/src/bytecode.rs b/src/bytecode.rs
index 70b037c..1824591 100644
--- a/src/bytecode.rs
+++ b/src/bytecode.rs
@@ -74,7 +74,8 @@ pub struct FnBlock<'bc> {
     /// used if there are errors.
     pub spans: Vec<'bc, Span>,
     /// How many parameters the function accepts.
-    pub arity: u8,
+    /// Yes, it supports 4294967295 parameters. I dare you to overflow that.
+    pub arity: u32,
 }
 
 #[cfg(feature = "_debug")]
diff --git a/src/compile.rs b/src/compile.rs
index ef3989d..7ffe4c4 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -192,7 +192,7 @@ impl<'bc, 'gc> Compiler<'bc, 'gc> {
                     decl.params[u8::MAX as usize]
                         .span
                         .extend(decl.params.last().unwrap().span),
-                    "Too many parameters".to_string(),
+                    "Too many parameters. How the fuck did you do this.".to_string(),
                 )
             })?,
         };
diff --git a/src/vm.rs b/src/vm.rs
index 739f639..25d0ab8 100644
--- a/src/vm.rs
+++ b/src/vm.rs
@@ -84,7 +84,7 @@ impl<'bc> Vm<'bc, '_> {
                 Some(&instr) => self.dispatch_instr(instr)?,
                 None => return Ok(()),
             }
-            debug_assert_eq!(self.current.stack_sizes[self.pc], self.stack.len());
+            // debug_assert_eq!(self.current.stack_sizes[self.pc], self.stack.len());
             self.pc += 1;
         }
     }
@@ -189,7 +189,7 @@ impl<'bc> Vm<'bc, '_> {
             Instr::Jmp(pos) => self.pc = (self.pc as isize + pos) as usize,
             Instr::Call => self.call()?,
             // todo implement
-            Instr::Return => return Ok(()),
+            Instr::Return => self.ret()?,
             Instr::ShrinkStack(size) => {
                 assert!(self.stack.len() >= size);
                 let new_len = self.stack.len() - size;
@@ -220,7 +220,7 @@ impl<'bc> Vm<'bc, '_> {
         if let Value::Function(func) = function {
             let fn_block = &self.blocks[func];
 
-            let new_stack_frame_start = self.stack.len() - fn_block.arity as usize;
+            let new_stack_frame_start = self.stack.len();
             self.stack_offset = new_stack_frame_start;
 
             self.stack.push(Value::NativeU(old_offset));
@@ -241,6 +241,23 @@ impl<'bc> Vm<'bc, '_> {
         Ok(())
     }
 
+    fn ret(&mut self) -> VmResult {
+        let current_arity: usize = self.current.arity.try_into().unwrap();
+
+        let bookkeeping_offset = self.stack_offset + current_arity;
+
+        let old_stack_offset = self.stack[bookkeeping_offset].as_native_int();
+        let old_pc = self.stack[bookkeeping_offset + 1].as_native_int();
+        let old_function = self.stack[bookkeeping_offset + 2].as_function();
+
+        self.stack_offset = old_stack_offset;
+        self.pc = old_pc;
+        self.current_block_index = old_function;
+        self.current = &self.blocks[old_function];
+
+        Ok(())
+    }
+
     fn type_error(&self) -> VmError {
         "bad type".into()
     }
@@ -262,6 +279,26 @@ Expected Stack size after instruction: {}",
     }
 }
 
+impl Value {
+    /// Unwrap the Value into a `usize` expecting the `NativeU` variant
+    fn as_native_int(&self) -> usize {
+        if let Value::NativeU(n) = self {
+            *n
+        } else {
+            unreachable!("expected native int, got {:?}", self);
+        }
+    }
+
+    /// Unwrap the Value into a `Function` expecting the `Function` variant
+    fn as_function(&self) -> Function {
+        if let Value::Function(fun) = self {
+            *fun
+        } else {
+            unreachable!("expected function, got {:?}", self);
+        }
+    }
+}
+
 impl Display for Value {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         match self {
diff --git a/test.dil b/test.dil
index 7091b02..f587f20 100644
--- a/test.dil
+++ b/test.dil
@@ -1,3 +1,10 @@
-fn test(uwu) {
-    let owo = 5;
-}
\ No newline at end of file
+fn test() {
+    print "yo wtf";
+}
+
+test();
+test();
+
+fn uwu() {}
+
+uwu();
diff --git a/tests/common.rs b/tests/common.rs
index 55d82c3..94cada5 100644
--- a/tests/common.rs
+++ b/tests/common.rs
@@ -4,13 +4,12 @@ macro_rules! run_test {
         #[test]
         fn $name() {
             let code = $code;
-            let output = _run_test(code);
+            let output = crate::common::_run_test(code);
             insta::assert_debug_snapshot!(output);
         }
     };
 }
 
-#[doc(hidden)]
 pub fn _run_test(code: &str) -> String {
     let mut stdout = Vec::<u8>::new();
     let mut cfg = dilaria::Config {
diff --git a/tests/control_flow.rs b/tests/control_flow.rs
index 32dbd15..acd0173 100644
--- a/tests/control_flow.rs
+++ b/tests/control_flow.rs
@@ -1,5 +1,4 @@
 mod common;
-use crate::common::_run_test;
 
 run_test!(
     single_if,
diff --git a/tests/functions.rs b/tests/functions.rs
new file mode 100644
index 0000000..0cae71a
--- /dev/null
+++ b/tests/functions.rs
@@ -0,0 +1,25 @@
+mod common;
+
+run_test!(
+    single_call,
+    r#"
+fn test() {
+    print "correct";
+}
+
+test();
+"#
+);
+
+run_test!(
+    single_call_expect_return,
+    r#"
+fn test() {
+    print "correct1";
+}
+
+test();
+
+print "correct2";
+"#
+);
diff --git a/tests/snapshots/functions__single_call.snap b/tests/snapshots/functions__single_call.snap
new file mode 100644
index 0000000..4cea669
--- /dev/null
+++ b/tests/snapshots/functions__single_call.snap
@@ -0,0 +1,7 @@
+---
+source: tests/functions.rs
+assertion_line: 3
+expression: output
+
+---
+"correct\n"