try to fix stack things

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),
 }

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;

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(())