From e9cad4b49e016258b8b1a709b2fb71c5628e808d Mon Sep 17 00:00:00 2001
From: Nilstrieb <48135649+Nilstrieb@users.noreply.github.com>
Date: Fri, 7 Jan 2022 21:45:00 +0100
Subject: [PATCH] try to fix stack things

---
 src/bytecode.rs |  3 +++
 src/compile.rs  | 20 ++++++++++++++++++--
 src/vm.rs       |  9 +++++++--
 test.dil        |  8 ++++++--
 4 files changed, 34 insertions(+), 6 deletions(-)

diff --git a/src/bytecode.rs b/src/bytecode.rs
index eb09dcf..3672dcd 100644
--- a/src/bytecode.rs
+++ b/src/bytecode.rs
@@ -85,4 +85,7 @@ pub enum Instr {
     JmpFalse(isize),
     /// Same as `JmpFalse`, but unconditional
     Jmp(isize),
+
+    /// Shrinks the stack by `usize` elements, should always be emitted before backwards jumps
+    ShrinkStack(usize),
 }
diff --git a/src/compile.rs b/src/compile.rs
index bb0f362..0bbb79e 100644
--- a/src/compile.rs
+++ b/src/compile.rs
@@ -207,13 +207,19 @@ impl<'bc, 'gc> Compiler<'bc, 'gc> {
           */
 
         let first_stmt_idx = self.code_len();
+        let pre_loop_stack_size = self.current_stack_size();
 
         self.loop_nesting += 1;
 
         self.compile_block(ast_block)?;
 
-        let jmp_offset = self.back_jmp_offset(first_stmt_idx);
+        self.push_instr(
+            Instr::ShrinkStack(self.current_stack_size() - pre_loop_stack_size),
+            StackChange::None,
+            span,
+        );
 
+        let jmp_offset = self.back_jmp_offset(first_stmt_idx);
         self.push_instr(Instr::Jmp(jmp_offset), StackChange::None, span);
 
         self.end_loop();
@@ -231,7 +237,7 @@ impl<'bc, 'gc> Compiler<'bc, 'gc> {
           */
 
         let cond_index = self.code_len();
-
+        let pre_loop_stack_size = self.current_stack_size();
         self.loop_nesting += 1;
 
         self.compile_expr(&while_stmt.cond)?;
@@ -241,6 +247,11 @@ impl<'bc, 'gc> Compiler<'bc, 'gc> {
 
         self.compile_block(&while_stmt.body)?;
 
+        self.push_instr(
+            Instr::ShrinkStack(self.current_stack_size() - pre_loop_stack_size),
+            StackChange::None,
+            while_stmt.span,
+        );
         let jmp_offset = self.back_jmp_offset(cond_index);
         self.push_instr(Instr::Jmp(jmp_offset), StackChange::None, while_stmt.span);
 
@@ -393,6 +404,11 @@ impl<'bc, 'gc> Compiler<'bc, 'gc> {
         block.code.len() as isize
     }
 
+    fn current_stack_size(&self) -> usize {
+        let block = &self.blocks[self.current_block];
+        block.stack_sizes.last().copied().unwrap_or(0)
+    }
+
     fn change_instr(&mut self, index: usize, instr: Instr) {
         let block = &mut self.blocks[self.current_block];
         block.code[index] = instr;
diff --git a/src/vm.rs b/src/vm.rs
index 4bbdd68..02dd450 100644
--- a/src/vm.rs
+++ b/src/vm.rs
@@ -63,9 +63,8 @@ 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());
             self.pc += 1;
-            // debug stack size assertion
-            // todo!()
         }
     }
 
@@ -163,6 +162,12 @@ impl<'bc> Vm<'bc, '_> {
                 }
             }
             Instr::Jmp(pos) => self.pc = (self.pc as isize + pos) as usize,
+            Instr::ShrinkStack(size) => {
+                assert!(self.stack.len() > size);
+                let new_len = self.stack.len() - size;
+                // SAFETY: We only ever shrink the vec, and we don't overflow. Value is copy so no leaks as a bonus
+                unsafe { self.stack.set_len(new_len) }
+            }
         }
 
         Ok(())
diff --git a/test.dil b/test.dil
index 2504cee..791f9b0 100644
--- a/test.dil
+++ b/test.dil
@@ -1,5 +1,9 @@
-loop {
-break;
+let i = 0;
+
+while i < 1000 {
+    print "lol";
+    i = i + 1;
 }
 
+
 print "hi";
\ No newline at end of file