From 3b7ce3eb2604c50e8a158bc27add9efc467b889b Mon Sep 17 00:00:00 2001
From: Ross Andrews <ross.andrews@gmail.com>
Date: Sat, 11 May 2024 01:24:41 -0500
Subject: [PATCH] Added once blocks

---
 forge_core/src/ast.rs                        | 11 ++++-
 forge_core/src/compiler/ast_nodes/block.rs   |  5 +-
 forge_core/src/compiler/ast_nodes/mod.rs     |  1 +
 forge_core/src/compiler/ast_nodes/once.rs    | 48 ++++++++++++++++++++
 forge_core/src/compiler/mod.rs               |  5 ++
 forge_core/src/compiler/state.rs             |  8 ++++
 forge_core/src/parser/ast_nodes/mod.rs       |  3 +-
 forge_core/src/parser/ast_nodes/once.rs      | 30 ++++++++++++
 forge_core/src/parser/ast_nodes/statement.rs |  1 +
 forge_core/src/parser/forge.pest             |  3 +-
 forge_core/src/parser/pairs_ext.rs           | 24 +---------
 vtest/src/forge_tests.rs                     | 21 +++++++++
 12 files changed, 132 insertions(+), 28 deletions(-)
 create mode 100644 forge_core/src/compiler/ast_nodes/once.rs
 create mode 100644 forge_core/src/parser/ast_nodes/once.rs

diff --git a/forge_core/src/ast.rs b/forge_core/src/ast.rs
index 4acde6e..277cccb 100644
--- a/forge_core/src/ast.rs
+++ b/forge_core/src/ast.rs
@@ -67,7 +67,8 @@ pub enum Statement {
     Conditional(Conditional),
     WhileLoop(WhileLoop),
     RepeatLoop(RepeatLoop),
-    Asm(Asm)
+    Asm(Asm),
+    Once(Once)
 }
 
 #[derive(PartialEq, Clone, Debug)]
@@ -117,6 +118,11 @@ pub struct Conditional {
     pub alternative: Option<Block>,
 }
 
+#[derive(PartialEq, Clone, Debug)]
+pub struct Once {
+    pub body: Block
+}
+
 #[derive(PartialEq, Clone, Debug)]
 pub struct WhileLoop {
     pub condition: Expr,
@@ -242,7 +248,8 @@ impl Statement {
                 Statement::Conditional(_) => "conditional",
                 Statement::WhileLoop(_) => "whileLoop",
                 Statement::RepeatLoop(_) => "repeatLoop",
-                Statement::Asm(_) => "asm"
+                Statement::Asm(_) => "asm",
+                Statement::Once(_) => "once"
             }
         )
     }
diff --git a/forge_core/src/compiler/ast_nodes/block.rs b/forge_core/src/compiler/ast_nodes/block.rs
index 0b84f4f..2093a83 100644
--- a/forge_core/src/compiler/ast_nodes/block.rs
+++ b/forge_core/src/compiler/ast_nodes/block.rs
@@ -5,7 +5,7 @@ use crate::compiler::CompileError;
 use crate::compiler::state::State;
 
 impl Compilable for Block {
-    fn process(self, state: &mut State, sig: Option<&mut CompiledFn>, loc: Location) -> Result<(), CompileError> {
+    fn process(self, state: &mut State, sig: Option<&mut CompiledFn>, _loc: Location) -> Result<(), CompileError> {
         let sig = sig.expect("Block outside function");
 
         let frame_size_before_block = sig.frame_size();
@@ -66,6 +66,9 @@ impl Compilable for Block {
                 Statement::RepeatLoop(repeat_loop) => {
                     repeat_loop.process(state, Some(sig), loc)?;
                 }
+                Statement::Once(once) => {
+                    once.process(state, Some(sig), loc)?;
+                }
             }
         }
 
diff --git a/forge_core/src/compiler/ast_nodes/mod.rs b/forge_core/src/compiler/ast_nodes/mod.rs
index e4b892b..4443a78 100644
--- a/forge_core/src/compiler/ast_nodes/mod.rs
+++ b/forge_core/src/compiler/ast_nodes/mod.rs
@@ -14,3 +14,4 @@ mod r#const;
 mod r#return;
 mod conditional;
 mod while_loop;
+mod once;
diff --git a/forge_core/src/compiler/ast_nodes/once.rs b/forge_core/src/compiler/ast_nodes/once.rs
new file mode 100644
index 0000000..2a4cb49
--- /dev/null
+++ b/forge_core/src/compiler/ast_nodes/once.rs
@@ -0,0 +1,48 @@
+use crate::ast::{Once, Location};
+use crate::compiler::compilable::Compilable;
+use crate::compiler::compiled_fn::CompiledFn;
+use crate::compiler::CompileError;
+use crate::compiler::state::State;
+
+impl Compilable for Once {
+    fn process(self, state: &mut State, sig: Option<&mut CompiledFn>, loc: Location) -> Result<(), CompileError> {
+        let sig = sig.expect("Once block outside function");
+
+        let Once { body } = self;
+        let sym = state.add_flag();
+
+        sig.emit_arg("loadw", sym.clone());
+        sig.emit("#if");
+        sig.emit_arg("push", 0);
+        sig.emit_arg("storew", sym);
+        body.process(state, Some(sig), loc)?;
+        sig.emit("#end");
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::compiler::test_utils::*;
+
+    #[test]
+    fn test_once() {
+        assert_eq!(
+            test_body(state_for("fn test() { var x = 0; once { x = 1; } }")),
+            vec![
+                "push 0",
+                "peekr",
+                "storew", // x = 0
+                "loadw _forge_gensym_3", // Load the once flag
+                "#if",
+                "push 0", // clear the once flag
+                "storew _forge_gensym_3",
+                "push 1", // x = 1
+                "peekr",
+                "storew",
+                "#end" // end once block
+            ]
+                .join("\n")
+        );
+    }
+}
\ No newline at end of file
diff --git a/forge_core/src/compiler/mod.rs b/forge_core/src/compiler/mod.rs
index d921c85..cdfa6cc 100644
--- a/forge_core/src/compiler/mod.rs
+++ b/forge_core/src/compiler/mod.rs
@@ -94,6 +94,11 @@ pub fn build_boot(src: &str, include_comments: bool) -> Result<Vec<String>, Comp
             asm.push(format!(".org {} + {}", label, size));
         }
 
+        // Once flags:
+        for label in state.flags {
+            asm.push(format!("{}: .db 1", label));
+        }
+
         // Final thing is to place a label for the stack:
         // (this is just a cell with the address of the following word)
         asm.push("stack: .db 0".into());
diff --git a/forge_core/src/compiler/state.rs b/forge_core/src/compiler/state.rs
index 855a0e5..73f3e82 100644
--- a/forge_core/src/compiler/state.rs
+++ b/forge_core/src/compiler/state.rs
@@ -20,6 +20,8 @@ pub(crate) struct State {
     pub strings: Vec<(Label, String)>,
     /// The statically-allocated buffers (size in bytes)
     pub buffers: Vec<(Label, usize)>,
+    /// The flags for once blocks
+    pub flags: Vec<Label>,
     /// The functions that have been prototyped but not yet defined
     pub prototypes: Scope,
     /// The initialization code for globals
@@ -72,6 +74,12 @@ impl State {
         sym
     }
 
+    pub(crate) fn add_flag(&mut self) -> Label {
+        let sym = self.gensym();
+        self.flags.push(sym.clone());
+        sym
+    }
+
     pub(crate) fn declare_function(&mut self, name: &str, _args: Vec<String>) -> Result<(), CompileError> {
         // If it's not already prototyped, gensym a label and put it in the list. If it is, just
         // ignore this (we don't check arity so it's not like the arglist being different matters)
diff --git a/forge_core/src/parser/ast_nodes/mod.rs b/forge_core/src/parser/ast_nodes/mod.rs
index 7ca4cf3..1c77662 100644
--- a/forge_core/src/parser/ast_nodes/mod.rs
+++ b/forge_core/src/parser/ast_nodes/mod.rs
@@ -14,4 +14,5 @@ mod while_loop;
 mod repeat_loop;
 mod program;
 mod operator;
-mod expr;
\ No newline at end of file
+mod expr;
+mod once;
\ No newline at end of file
diff --git a/forge_core/src/parser/ast_nodes/once.rs b/forge_core/src/parser/ast_nodes/once.rs
new file mode 100644
index 0000000..a7a7bea
--- /dev/null
+++ b/forge_core/src/parser/ast_nodes/once.rs
@@ -0,0 +1,30 @@
+use crate::ast::{Block, Once};
+use crate::parser::{AstNode, Pair, PestRule};
+
+impl AstNode for Once {
+    const RULE: PestRule = PestRule::once;
+    fn from_pair(pair: Pair) -> Self {
+        let mut inner = pair.into_inner().peekable();
+        let body = Block::from_pair(inner.next().unwrap());
+        Self { body }
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use crate::ast::Statement;
+    use crate::parser::test_utils::{dislocate, dislocated_block};
+    use crate::parser::Parseable;
+    use super::*;
+
+    #[test]
+    fn parse_once_blocks() {
+        assert!(match Statement::from_str("once { foo(x); }") {
+            Ok(Statement::Once(Once { body })) => {
+                assert_eq!(dislocate(body.0), dislocated_block("{ foo(x); }"));
+                true
+            },
+            _ => false
+        });
+    }
+}
\ No newline at end of file
diff --git a/forge_core/src/parser/ast_nodes/statement.rs b/forge_core/src/parser/ast_nodes/statement.rs
index b379dfa..9be561c 100644
--- a/forge_core/src/parser/ast_nodes/statement.rs
+++ b/forge_core/src/parser/ast_nodes/statement.rs
@@ -15,6 +15,7 @@ impl AstNode for Statement {
             PestRule::conditional => Self::Conditional(Conditional::from_pair(pair)),
             PestRule::while_loop => Self::WhileLoop(WhileLoop::from_pair(pair)),
             PestRule::repeat_loop => Self::RepeatLoop(RepeatLoop::from_pair(pair)),
+            PestRule::once => Self::Once(Once::from_pair(pair)),
             PestRule::asm => Self::Asm(Asm::from_pair(pair)),
             _ => unreachable!(),
         }
diff --git a/forge_core/src/parser/forge.pest b/forge_core/src/parser/forge.pest
index 19f0d74..30c488e 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 | continue_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 | once }
 
 asm      = { "asm" ~ asm_args? ~ "{" ~ asm_body ~ "}" }
 asm_args = _{ "(" ~ expr ~ ("," ~ expr)* ~ ")" }
@@ -95,6 +95,7 @@ return_stmt = { "return" ~ expr? }
 conditional = { "if" ~ "(" ~ expr ~ ")" ~ block ~ ("else" ~ block)? }
 while_loop  = { "while" ~ "(" ~ expr ~ ")" ~ block }
 repeat_loop = { "repeat" ~ "(" ~ expr ~ ")" ~ name? ~ block }
+once  = { "once" ~ block }
 var_decl    = { "var" ~ name ~ ("=" ~ expr)? }
 global      = { "global" ~ name ~ ("=" ~ expr)? ~ ";" }
 const_decl  = { "const" ~ name ~ "=" ~ (string | expr) ~ ";" }
diff --git a/forge_core/src/parser/pairs_ext.rs b/forge_core/src/parser/pairs_ext.rs
index 6d9bea6..9c0c611 100644
--- a/forge_core/src/parser/pairs_ext.rs
+++ b/forge_core/src/parser/pairs_ext.rs
@@ -7,12 +7,8 @@ pub trait PairExt {
     /// Go down one level in the AST, return the first node. Assumes the first node exists
     fn first(self) -> Self;
 
-    /// Returns the first child node as a String; useful for things like names and numbers
-    /// which are really just typed strings
-    fn first_as_string(self) -> String;
-
     /// Just like `first` but panics if there's more than one child node.
-    fn only(self) -> Self;
+    //fn only(self) -> Self;
 
     /// This should really be an AstNode or something, but it's used so often: turn a Pair into
     /// an i32 by parsing the Forge number format (0xwhatever, 0bwhatever, etc)
@@ -30,17 +26,6 @@ impl<'a> PairExt for Pair<'a> {
         self.into_inner().next().unwrap()
     }
 
-    fn first_as_string(self) -> String {
-        String::from(self.first().as_str())
-    }
-
-    fn only(self) -> Pair<'a> {
-        let mut iter = self.into_inner();
-        let child = iter.next().unwrap();
-        debug_assert_eq!(iter.next(), None);
-        child
-    }
-
     fn into_number(self) -> i32 {
         let first = self.into_inner().next().unwrap();
         match first.as_rule() {
@@ -75,17 +60,10 @@ pub trait PairsExt {
     /// Peek the next sibling but only if it matches a given rule: used for things
     /// with optional modifiers following
     fn next_if_rule(&mut self, rule: PestRule) -> Option<Pair>;
-
-    /// Grab the first thing from the list if you're sure it exists
-    fn first(&mut self) -> Pair;
 }
 
 impl PairsExt for Peekable<Pairs<'_>> {
     fn next_if_rule(&mut self, rule: PestRule) -> Option<Pair> {
         self.next_if(|p| p.as_rule() == rule)
     }
-
-    fn first(&mut self) -> Pair {
-        self.next().unwrap().first()
-    }
 }
diff --git a/vtest/src/forge_tests.rs b/vtest/src/forge_tests.rs
index 51424e2..7345868 100644
--- a/vtest/src/forge_tests.rs
+++ b/vtest/src/forge_tests.rs
@@ -231,6 +231,27 @@ fn static2_test() {
     );
 }
 
+#[test]
+fn once_test() {
+    // A once block.
+    // We pass a pointer in, it increments the var, but only the first time.
+    let src = "
+        fn foo(n) {
+            once {
+                *n = *n + 1;
+            }
+        }
+        fn main() {
+            var n = 0;
+            foo(&n); foo(&n); foo(&n);
+            return n;
+        }";
+    assert_eq!(
+        main_return(src),
+        1
+    );
+}
+
 //#[test]
 // This is no longer cursed, or a test. The revised calling convention with the pool pointer makes
 // it now perfectly sane. Left here for posterity.