From f19736cb8a9fb69c85857856feec91b0b729f93c Mon Sep 17 00:00:00 2001
From: Ross Andrews <ross.andrews@gmail.com>
Date: Sat, 4 May 2024 20:37:47 -0500
Subject: [PATCH] Continue keyword

---
 forge_core/src/ast.rs                         |  2 +
 forge_core/src/compiler/ast_nodes/block.rs    | 39 +++++++++++-
 .../src/compiler/ast_nodes/repeat_loop.rs     |  6 +-
 .../src/compiler/ast_nodes/while_loop.rs      |  2 +-
 forge_core/src/compiler/compiled_fn.rs        | 16 ++---
 forge_core/src/parser/ast_nodes/statement.rs  |  1 +
 forge_core/src/parser/forge.pest              |  3 +-
 vasm_core/src/ast.rs                          |  1 +
 vasm_core/src/vasm.pest                       |  2 +-
 vasm_core/src/vasm_parser.rs                  |  1 +
 vasm_core/src/vasm_preprocessor.rs            | 61 ++++++++++++++++---
 vtest/src/forge_tests.rs                      | 34 +++++++++++
 vweb/src/forge_highlighter.js                 |  2 +-
 13 files changed, 147 insertions(+), 23 deletions(-)

diff --git a/forge_core/src/ast.rs b/forge_core/src/ast.rs
index 49c8396..4acde6e 100644
--- a/forge_core/src/ast.rs
+++ b/forge_core/src/ast.rs
@@ -59,6 +59,7 @@ pub struct FunctionPrototype {
 #[derive(PartialEq, Clone, Debug)]
 pub enum Statement {
     Break,
+    Continue,
     Return(Return),
     Assignment(Assignment),
     Expr(Expr),
@@ -233,6 +234,7 @@ impl Statement {
         String::from(
             match self {
                 Statement::Break => "break",
+                Statement::Continue => "continue",
                 Statement::Return(_) => "return",
                 Statement::Assignment(_) => "assignment",
                 Statement::Expr(_) => "expr",
diff --git a/forge_core/src/compiler/ast_nodes/block.rs b/forge_core/src/compiler/ast_nodes/block.rs
index b609f4e..0b84f4f 100644
--- a/forge_core/src/compiler/ast_nodes/block.rs
+++ b/forge_core/src/compiler/ast_nodes/block.rs
@@ -22,6 +22,21 @@ impl Compilable for Block {
                         return Err(CompileError(loc.line, loc.col, "Break outside loop".to_string()))
                     }
                 }
+                Statement::Continue => {
+                    if sig.in_loop() {
+                        // This is subtle. To "continue" a while or until loop, we simply need to #continue
+                        // and the assembler takes care of the macro for us. Repeat loops are different though
+                        // because the counter inc / dec takes place at the end of the loop, so we can't use
+                        // the macro, we just want to generate a jmpr instead
+                        if let Some(label) = sig.continue_point() {
+                            sig.emit_arg("jmpr", format!("@{}", label))
+                        } else {
+                            sig.emit("#continue")
+                        }
+                    } else {
+                        return Err(CompileError(loc.line, loc.col, "Continue outside loop".to_string()))
+                    }
+                }
                 Statement::Return(ret) => {
                     ret.process(state, Some(sig), loc)?;
                 }
@@ -89,7 +104,7 @@ mod test {
             vec![
                 "push 10", "#while", "dup", "agt 0", "#do", // standard repeat loop header
                 "#break", // the break
-                "sub 1", "#end", "pop" // repeat loop footer
+                "_forge_gensym_3:", "sub 1", "#end", "pop" // repeat loop footer
             ].join("\n")
         );
 
@@ -104,11 +119,28 @@ mod test {
     }
 
     #[test]
-    fn test_malformed_break_statements() {
+    fn test_continue_statements() {
+        assert_eq!(
+            test_body(state_for("fn test() { repeat(10) { continue; } }")),
+            vec![
+                "push 10", "#while", "dup", "agt 0", "#do", // standard repeat loop header
+                "jmpr @_forge_gensym_3", // the continue, jumping to the label
+                "_forge_gensym_3:", // The label demnoting the start of the counter update
+                "sub 1", "#end", "pop" // repeat loop footer
+            ].join("\n")
+        );
+    }
+
+    #[test]
+    fn test_malformed_break_continue_statements() {
         assert_eq!(
             test_err("fn test() { break; }"),
             Some(CompileError(1, 13, "Break outside loop".to_string()))
-        )
+        );
+        assert_eq!(
+            test_err("fn test() { continue; }"),
+            Some(CompileError(1, 13, "Continue outside loop".to_string()))
+        );
     }
 
     #[test]
@@ -129,6 +161,7 @@ mod test {
                 "#do",
                 "push 2", // Pointless loop body
                 "pop",
+                "_forge_gensym_3:",
                 "peekr", // Load c as an lvalue
                 "dup", // Dup it, load it, add 1
                 "loadw",
diff --git a/forge_core/src/compiler/ast_nodes/repeat_loop.rs b/forge_core/src/compiler/ast_nodes/repeat_loop.rs
index f21e286..e0788eb 100644
--- a/forge_core/src/compiler/ast_nodes/repeat_loop.rs
+++ b/forge_core/src/compiler/ast_nodes/repeat_loop.rs
@@ -41,9 +41,11 @@ impl Compilable for RepeatLoop {
 
         // Loop body:
         sig.emit("#do");
-        sig.enter_loop();
+        let cp = state.gensym();
+        sig.enter_loop(Some(cp.clone()));
         body.process(state, Some(sig), loc)?;
         sig.exit_loop();
+        sig.emit(format!("{}:", cp).as_str());
 
         // After the body we need to increment the counter if there is one
         if named_counter {
@@ -115,6 +117,7 @@ mod test {
                 "peekr", // Load x as an lvalue
                 "add 3",
                 "storew", // Store c + x into it
+                "_forge_gensym_3:",
                 "peekr", // Load c as an lvalue
                 "add 6",
                 "dup", // Dup it, load it, add 1
@@ -157,6 +160,7 @@ mod test {
                 "peekr", // Load x as an lvalue
                 "add 3",
                 "storew", // Store 2x into it
+                "_forge_gensym_3:", // Label for the continue point, before the counter change
                 "sub 1", // decrement the counter
                 "#end", // End of the loop body!
                 "pop", // Drop the limit off the top
diff --git a/forge_core/src/compiler/ast_nodes/while_loop.rs b/forge_core/src/compiler/ast_nodes/while_loop.rs
index bb2ab35..65c6749 100644
--- a/forge_core/src/compiler/ast_nodes/while_loop.rs
+++ b/forge_core/src/compiler/ast_nodes/while_loop.rs
@@ -11,7 +11,7 @@ impl Compilable for WhileLoop {
         sig.emit("#while");
         condition.process(state, Some(sig), loc)?;
         sig.emit("#do");
-        sig.enter_loop();
+        sig.enter_loop(None);
         body.process(state, Some(sig), loc)?;
         sig.exit_loop();
         sig.emit("#end");
diff --git a/forge_core/src/compiler/compiled_fn.rs b/forge_core/src/compiler/compiled_fn.rs
index 03a70ff..e6658cf 100644
--- a/forge_core/src/compiler/compiled_fn.rs
+++ b/forge_core/src/compiler/compiled_fn.rs
@@ -1,7 +1,6 @@
 use std::cmp::max;
 use std::collections::btree_map::Entry::Vacant;
 use std::fmt::Display;
-use crate::ast::Location;
 use crate::compiler::compile_error::CompileError;
 use crate::compiler::text::Text;
 use crate::compiler::utils::{Label, Scope, Variable};
@@ -57,7 +56,7 @@ pub struct CompiledFn {
     pub max_frame_size: usize,
     pub local_scope: Scope,
     pub arity: usize,
-    pub loop_level: u32,
+    pub continue_points: Vec<Option<Label>>,
 
     /// The preamble is the area of the function capturing the args, etc
     pub preamble: Text,
@@ -85,12 +84,12 @@ impl CompiledFn {
         }
     }
 
-    pub fn enter_loop(&mut self) {
-        self.loop_level += 1
+    pub fn enter_loop(&mut self, continue_point: Option<Label>) {
+        self.continue_points.push(continue_point)
     }
     pub fn exit_loop(&mut self) {
-        if self.loop_level > 0 {
-            self.loop_level -= 1;
+        if self.continue_points.len() > 0 {
+            self.continue_points.pop();
         } else {
             // I doubt it's possible for this to happen, since it would be
             // caught at the parse stage
@@ -98,7 +97,10 @@ impl CompiledFn {
         }
     }
     pub fn in_loop(&self) -> bool {
-        self.loop_level > 0
+        self.continue_points.len() > 0
+    }
+    pub fn continue_point(&self) -> &Option<Label> {
+        self.continue_points.last().unwrap_or(&None)
     }
 
     /// Return an iterator over all the
diff --git a/forge_core/src/parser/ast_nodes/statement.rs b/forge_core/src/parser/ast_nodes/statement.rs
index f04ec59..b379dfa 100644
--- a/forge_core/src/parser/ast_nodes/statement.rs
+++ b/forge_core/src/parser/ast_nodes/statement.rs
@@ -7,6 +7,7 @@ impl AstNode for Statement {
         let pair = pair.first();
         match pair.as_rule() {
             PestRule::break_stmt => Self::Break,
+            PestRule::continue_stmt => Self::Continue,
             PestRule::return_stmt => Self::Return(Return::from_pair(pair)),
             PestRule::assignment => Self::Assignment(Assignment::from_pair(pair)),
             PestRule::expr => Self::Expr(Expr::from_pair(pair)),
diff --git a/forge_core/src/parser/forge.pest b/forge_core/src/parser/forge.pest
index 5ae7698..19f0d74 100644
--- a/forge_core/src/parser/forge.pest
+++ b/forge_core/src/parser/forge.pest
@@ -73,7 +73,7 @@ suffix    = _{ subscript | arglist }
 subscript = { "[" ~ expr ~ "]" }
 arglist   = { "(" ~ (expr ~ ("," ~ expr)*)? ~ ")" }
 
-statement = { asm | ((break_stmt | return_stmt | var_decl | assignment | expr) ~ ";") | conditional | while_loop | repeat_loop }
+statement = { asm | ((break_stmt | continue_stmt | return_stmt | var_decl | assignment | expr) ~ ";") | conditional | while_loop | repeat_loop }
 
 asm      = { "asm" ~ asm_args? ~ "{" ~ asm_body ~ "}" }
 asm_args = _{ "(" ~ expr ~ ("," ~ expr)* ~ ")" }
@@ -90,6 +90,7 @@ declaration = { function | function_prototype | global | const_decl }
 program     = { (COMMENT | WHITESPACE)? ~ declaration* ~ EOI }
 
 break_stmt = { "break" }
+continue_stmt = { "continue" }
 return_stmt = { "return" ~ expr? }
 conditional = { "if" ~ "(" ~ expr ~ ")" ~ block ~ ("else" ~ block)? }
 while_loop  = { "while" ~ "(" ~ expr ~ ")" ~ block }
diff --git a/vasm_core/src/ast.rs b/vasm_core/src/ast.rs
index fed9657..42c7d9a 100644
--- a/vasm_core/src/ast.rs
+++ b/vasm_core/src/ast.rs
@@ -63,6 +63,7 @@ pub enum Macro {
     Do,
     End,
     Break,
+    Continue,
 }
 
 #[derive(Debug, PartialEq, Clone)]
diff --git a/vasm_core/src/vasm.pest b/vasm_core/src/vasm.pest
index f02e4e9..add5a62 100644
--- a/vasm_core/src/vasm.pest
+++ b/vasm_core/src/vasm.pest
@@ -69,7 +69,7 @@ org_directive = { label_def? ~ ".org" ~ expr }
 equ_directive = { label_def ~ ".equ" ~ expr }
 
 include = { "include" ~ string }
-control = { "if" | "unless" | "while" | "until" | "do" | "else" | "end" | "break" }
+control = { "if" | "unless" | "while" | "until" | "do" | "else" | "end" | "break" | "continue" }
 preprocessor = {"#" ~ (control | include) }
 
 blank = { WHITESPACE? ~ COMMENT? }
diff --git a/vasm_core/src/vasm_parser.rs b/vasm_core/src/vasm_parser.rs
index 37bec12..dcd351c 100644
--- a/vasm_core/src/vasm_parser.rs
+++ b/vasm_core/src/vasm_parser.rs
@@ -209,6 +209,7 @@ fn parse_macro(line: Pair) -> Macro {
             "do" => Macro::Do,
             "end" => Macro::End,
             "break" => Macro::Break,
+            "continue" => Macro::Continue,
             _ => unreachable!(),
         },
         Rule::include => Macro::Include(create_string(pre.only())),
diff --git a/vasm_core/src/vasm_preprocessor.rs b/vasm_core/src/vasm_preprocessor.rs
index 0aa97bc..522a55c 100644
--- a/vasm_core/src/vasm_preprocessor.rs
+++ b/vasm_core/src/vasm_preprocessor.rs
@@ -217,23 +217,32 @@ impl<T: IntoIterator<Item = String>, F: Fn(String) -> Result<T, AssembleError>>
                 }
                 _ => return Some(AssembleError::MacroError(self.current_location())),
             },
-
             Macro::Break => {
-                let innermost_loop = self.control_stack.iter().rev().find(|cs| {
-                    match cs {
-                        ControlStructure::LoopDo(_, _) => true,
-                        _ => false
-                    }
-                });
-                if let Some(ControlStructure::LoopDo(_, after)) = innermost_loop {
+                if let Some(ControlStructure::LoopDo(_, after)) = self.innermost_loop() {
                     self.emit(format!("jmpr @{}", after));
                 } else {
                     return Some(AssembleError::MacroError(self.current_location()))
                 }
             }
+            Macro::Continue => {
+                if let Some(ControlStructure::LoopDo(test, _)) = self.innermost_loop() {
+                    self.emit(format!("jmpr @{}", test));
+                } else {
+                    return Some(AssembleError::MacroError(self.current_location()))
+                }
+            }
         }
         None
     }
+
+    pub fn innermost_loop(&self) -> Option<&ControlStructure> {
+        self.control_stack.iter().rev().find(|cs| {
+            match cs {
+                ControlStructure::LoopDo(_, _) => true,
+                _ => false
+            }
+        })
+    }
 }
 
 #[cfg(test)]
@@ -386,4 +395,40 @@ mod tests {
             Some(AssembleError::MacroError(Location::from(2)))
         )
     }
+
+    #[test]
+    fn test_preprocess_continue() {
+        // First a working one
+        assert_eq!(
+            lines_for(stringify(vec!["#while", "#do", "#continue", "#end"])),
+            vec![
+                VASMLine::LabelDef(Label("__gensym_1".to_string())),
+                VASMLine::Instruction(
+                    None,
+                    Brz,
+                    Some(Node::RelativeLabel("__gensym_2".to_string()))
+                ),
+                VASMLine::Instruction(
+                    None,
+                    Jmpr,
+                    Some(Node::RelativeLabel("__gensym_1".to_string()))
+                ),
+                VASMLine::Instruction(
+                    None,
+                    Jmpr,
+                    Some(Node::RelativeLabel("__gensym_1".to_string()))
+                ),
+                VASMLine::LabelDef(Label::from("__gensym_2"))
+            ]
+        );
+    }
+
+    #[test]
+    fn test_malformed_continue() {
+        // Malformed one
+        assert_eq!(
+            error_for(stringify(vec!["#while", "#continue"])),
+            Some(AssembleError::MacroError(Location::from(2)))
+        )
+    }
 }
diff --git a/vtest/src/forge_tests.rs b/vtest/src/forge_tests.rs
index a1c63a4..51424e2 100644
--- a/vtest/src/forge_tests.rs
+++ b/vtest/src/forge_tests.rs
@@ -159,6 +159,40 @@ fn peekpoke_test() {
     )
 }
 
+#[test]
+fn break_test() {
+    // Loop 0..10? Not quite!
+    assert_eq!(
+        main_return(
+            "fn main() {
+                    var n = 0;
+                    while(n < 10) {
+                        n = n + 1;
+                        if (n == 4) { break; }
+                    }
+                    return n;
+                }"),
+        4
+    )
+}
+
+#[test]
+fn continue_test() {
+    // Loop 0..4, but skip 3
+    assert_eq!(
+        main_return(
+            "fn main() {
+                    var n = 0;
+                    repeat(5) i {
+                        if (i == 3) { continue; }
+                        n = n + i;
+                    }
+                    return n;
+                }"),
+        7
+    )
+}
+
 #[test]
 fn static_test() {
     // Because the static(1) returns the same address each time, a[0] is the same variable
diff --git a/vweb/src/forge_highlighter.js b/vweb/src/forge_highlighter.js
index 99e671e..3b307c9 100644
--- a/vweb/src/forge_highlighter.js
+++ b/vweb/src/forge_highlighter.js
@@ -16,7 +16,7 @@ export default {
             if (word = stream.match(/^[a-zA-Z_$][a-zA-Z0-9_$]*/)) {
                 // This is either a keyword or an identifier.
                 // All the keywords are valid names, so, check that:
-                if (word[0].match(/^fn|var|new|static|asm|return|if|while|repeat|global|const|peek|poke|break$/)) {
+                if (word[0].match(/^fn|var|new|static|asm|return|if|while|repeat|global|const|peek|poke|break|continue$/)) {
                     if (word[0] === 'asm') {
                         state.asm = true
                     } // This is the start of an asm control structure