bug(forth3): Fix the behavior of leave in forth3 (#276)

This PR fixes the behavior of `leave`, which didn't previously work
when there was code between the `leave` and the `loop`.

source/forth3/src/lib.rs

@@ -405,6 +405,23 @@ pub mod test {
         );
     }
 
+    #[test]
+    fn loop_leave() {
+        all_runtest(
+            r#"
+            > : test 10 0 do i . 43 emit i 5 = if leave then 10 0 do 42 emit i 5 = if leave then loop cr loop ;
+            < ok.
+            > test
+            < 0 +******
+            < 1 +******
+            < 2 +******
+            < 3 +******
+            < 4 +******
+            < 5 +ok.
+            "#,
+        )
+    }
+
     #[test]
     fn execute() {
         all_runtest(

source/forth3/src/vm/builtins.rs

@@ -182,12 +182,40 @@ impl<T: 'static> Forth<T> {
         builtin!("(jmp)", Self::jump),
         // NOTE: REQUIRED for `:` (if you want literals)
         builtin!("(literal)", Self::literal),
+        // NOTE: REQUIRED for `:` (if you want literals)
+        builtin!("(rliteral)", Self::rliteral),
         // NOTE: REQUIRED for `constant`
         builtin!("(constant)", Self::constant),
         // NOTE: REQUIRED for `variable` or `array`
         builtin!("(variable)", Self::variable),
+        builtin!("panic", Self::panic),
     ];
 
+    /// Dumps all stacks and ends the currently executing program
+    pub fn panic(&mut self) -> Result<(), Error> {
+        writeln!(&mut self.output, "cstack",)?;
+
+        while let Some(c) = self.call_stack.pop() {
+            writeln!(
+                &mut self.output,
+                "{} ({}/{})",
+                unsafe { (*c.eh.as_ptr()).name.as_str() },
+                c.idx,
+                c.len,
+            )?;
+        }
+
+        writeln!(&mut self.output, "\ndstack",)?;
+
+        while self.pop_print().is_ok() {}
+
+        writeln!(&mut self.output, "\nrstack",)?;
+
+        while self.return_to_data_stack().is_ok() && self.pop_print().is_ok() {}
+
+        Ok(())
+    }
+
     pub fn dict_free(&mut self) -> Result<(), Error> {
         let capa = self.dict.alloc.capacity();
         let used = self.dict.alloc.used();
@@ -577,10 +605,18 @@ impl<T: 'static> Forth<T> {
     }
 
     pub fn loop_leave(&mut self) -> Result<(), Error> {
+        // Pop the loop counter and limit
         let _ = self.return_stack.try_pop()?;
-        let a = self.return_stack.try_peek()?;
-        self.return_stack
-            .push(unsafe { Word::data(a.data.wrapping_sub(1)) })?;
+        let _ = self.return_stack.try_pop()?;
+        // Pop the "end of loop" value
+        let idx = self.return_stack.try_pop()?;
+        let idx = unsafe { idx.data };
+        let idx = u16::try_from(idx).map_err(|_| Error::BadCfaOffset)?;
+
+        // Move the parent's interpreter index forward to the end of loop index
+        let parent = self.call_stack.try_peek_back_n_mut(1)?;
+        parent.idx = idx;
+
         Ok(())
     }
 
@@ -593,6 +629,8 @@ impl<T: 'static> Forth<T> {
             self.return_stack.push(ctr)?;
             self.jump()
         } else {
+            self.return_stack.try_pop()?;
+            // also pop the loop len counter
             self.return_stack.try_pop()?;
             self.skip_literal()
         }
@@ -847,6 +885,16 @@ impl<T: 'static> Forth<T> {
         Ok(())
     }
 
+    /// `(rliteral)` is used mid-interpret to put the NEXT word of the parent's
+    /// CFA into the *return* stack as a value
+    pub fn rliteral(&mut self) -> Result<(), Error> {
+        let parent = self.call_stack.try_peek_back_n_mut(1)?;
+        let literal = parent.get_current_word()?;
+        parent.offset(1)?;
+        self.return_stack.push(literal)?;
+        Ok(())
+    }
+
     /// `(literal)` is used mid-interpret to put the NEXT word of the parent's
     /// CFA array into the stack as a value.
     pub fn literal(&mut self) -> Result<(), Error> {

source/forth3/src/vm/mod.rs

@@ -465,14 +465,32 @@ impl<T> Forth<T> {
     }
 
     fn munch_do(&mut self, len: &mut u16) -> Result<u16, Error> {
-        let start = *len;
+        let pre_start = *len;
+        // At the beginning of the loop, we want to place "the index of
+        // the end of the loop" on the rstack, UNDER the loop variables.
+        //
+        // If someone calls `leave`, the interpreter will be fast-forwarded
+        // to this index. We place the r-push, and leave a placeholder which
+        // we'll fill when we know where the loop ends
+        let rlit = self.find_word("(rliteral)").ok_or(Error::WordNotInDict)?;
+        self.dict.alloc.bump_write(Word::ptr(rlit.as_ptr()))?;
+        let rlit_offset: &mut i32 = {
+            let cj_offset_word = self.dict.alloc.bump::<Word>()?;
+            unsafe {
+                cj_offset_word.as_ptr().write(Word::data(0));
+                &mut (*cj_offset_word.as_ptr()).data
+            }
+        };
+        *len += 2;
 
-        // Write a conditional jump, followed by space for a literal
-        let literal_cj = self.find_word("2d>2r").ok_or(Error::WordNotInDict)?;
-        self.dict.alloc.bump_write(Word::ptr(literal_cj.as_ptr()))?;
+        // Take the loop start and end from the data stack to the return stack
+        let d2r2 = self.find_word("2d>2r").ok_or(Error::WordNotInDict)?;
+        self.dict.alloc.bump_write(Word::ptr(d2r2.as_ptr()))?;
         *len += 1;
 
+        // Start is where the LOOP starts, e.g. where we need to jump back to
         let do_start = *len;
+
         // Now work until we hit an else or then statement.
         loop {
             match self.munch_one(len) {
@@ -496,7 +514,9 @@ impl<T> Forth<T> {
         self.dict.alloc.bump_write(Word::data(offset))?;
         *len += 2;
 
-        Ok(*len - start)
+        *rlit_offset = (*len).into();
+
+        Ok(*len - pre_start)
     }
 
     fn munch_if(&mut self, len: &mut u16) -> Result<u16, Error> {