put tail call substitution in an fn and call it inside if/then

1 files changed, 66 insertions, 18 deletions

diff --git a/src/forth/compiler.rs b/src/forth/compiler.rs
index 89964e1..b3f5f08 100644
--- a/src/forth/compiler.rs
+++ b/src/forth/compiler.rs
@@ -113,6 +113,30 @@ impl<'a> Compiler<'a> {
         self.anno_mut().push(anno);
     }
 
+    // substitute previous call/if with tail version
+    //
+    // only use where we might by pushing `OpCode::Ret`. honestly,
+    // should todo: put this in bc_push so it's always applied
+    // appropriately.
+    fn sub_tcall(&mut self, anno: Annotation) {
+        match self.bc().0.last() {
+            Some(OpCode::Call(i)) => {
+                let k = *i;
+                self.bc_mut().0.pop();
+                self.bc_mut().0.push(OpCode::TCall(k));
+            },
+            Some(OpCode::If(t, f)) => {
+                let t = *t;
+                let f = *f;
+                self.bc_mut().0.pop();
+                self.bc_mut().0.push(OpCode::TIf(t, f));
+                // technically only needed if ‘f’ is None, but whatever.
+                self.bc_push(OpCode::Ret, anno);
+            },
+            _ => self.bc_push(OpCode::Ret, anno),
+        }
+    }
+
     pub fn compile(&mut self) -> CompileResult<()> {
         while let Some((word, start, end)) = self.next_word() {
             let anno = Annotation { loc: (start, end) };
@@ -137,22 +161,7 @@ impl<'a> Compiler<'a> {
                         self.annotations.push(vec![]);
                     },
                     ";" => {
-                        match self.bc().0.last() {
-                            Some(OpCode::Call(i)) => {
-                                let k = *i;
-                                self.bc_mut().0.pop();
-                                self.bc_mut().0.push(OpCode::TCall(k));
-                            },
-                            Some(OpCode::If(t, f)) => {
-                                let t = *t;
-                                let f = *f;
-                                self.bc_mut().0.pop();
-                                self.bc_mut().0.push(OpCode::TIf(t, f));
-                                // technically only needed if ‘f’ is None, but whatever.
-                                self.bc_push(OpCode::Ret, anno);
-                            },
-                            _ => self.bc_push(OpCode::Ret, anno),
-                        }
+                        self.sub_tcall(anno);
                         self.defstack.pop().ok_or(CompileError::DefStackEmpty)?;
                     },
                     "if" => {
@@ -164,7 +173,7 @@ impl<'a> Compiler<'a> {
                         self.if_stack.push((IfClauses::True(i), anno));
                     },
                     "else" => {
-                        self.bc_push(OpCode::Ret, anno);
+                        self.sub_tcall(anno);
                         self.defstack.pop();
 
                         let i = self.wordlist.0.len();
@@ -180,7 +189,7 @@ impl<'a> Compiler<'a> {
                         self.if_stack.push((IfClauses::TrueFalse(true_clause, i), anno));
                     },
                     "then" => {
-                        self.bc_push(OpCode::Ret, anno);
+                        self.sub_tcall(anno);
                         self.defstack.pop();
 
                         match self.if_stack.pop() {
@@ -361,4 +370,43 @@ mod tests {
         let foo = &comp.wordlist.0[1];
         assert_eq!(foo, vec![OpCode::TCall(1)]);
     }
+
+    #[test]
+    fn tcall_if_then() {
+        let prog = r#"
+: foo
+  dup
+  1 > if
+    foo
+  then
+  ;
+"#;
+        let comp = compiler_for(prog);
+        eprintwordlist(&comp.wordlist);
+        let tbranch = &comp.wordlist.0[2];
+        eprintln!("true: {tbranch:?}");
+        assert_eq!(tbranch.0[0], OpCode::TCall(1));
+    }
+
+    #[test]
+    fn tcall_if_else_then() {
+        let prog = r#"
+: foo
+  dup
+  1 > if
+    foo
+  else
+    foo
+  then
+  ;
+"#;
+        let comp = compiler_for(prog);
+        eprintwordlist(&comp.wordlist);
+        let tbranch = &comp.wordlist.0[2];
+        eprintln!("true: {tbranch:?}");
+        assert_eq!(tbranch.0[0], OpCode::TCall(1));
+        let fbranch = &comp.wordlist.0[3];
+        eprintln!("false: {fbranch:?}");
+        assert_eq!(tbranch.0[0], OpCode::TCall(1));
+    }
 }