From f2bde7ddcbae4b199b379d8b4895c906334158f6 Mon Sep 17 00:00:00 2001
From: Nick Krichevsky <nick@ollien.com>
Date: Fri, 24 May 2024 11:09:41 -0400
Subject: [PATCH] Add whlie loop support

---
 build.rs     |  1 +
 src/eval.rs  |  8 ++++++++
 src/parse.rs | 29 ++++++++++++++++++++++++++++-
 3 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/build.rs b/build.rs
index beaf7ac..7f9a7ce 100644
--- a/build.rs
+++ b/build.rs
@@ -68,6 +68,7 @@ fn do_ast_codegen<W: Write>(mut output: W) {
             "If",
             "condition: Expr, then_branch: Box<Stmt>, else_branch: Option<Box<Stmt>>",
         ),
+        ("While", "condition: Expr, body: Box<Stmt>"),
     ]);
 
     define_ast(&mut output, "Stmt", &statement_types).expect("failed to generate ast values");
diff --git a/src/eval.rs b/src/eval.rs
index a39bfec..d50ecf5 100644
--- a/src/eval.rs
+++ b/src/eval.rs
@@ -269,6 +269,14 @@ impl StmtVisitor<Result<(), ScriptError>> for InterpreterRunner<'_> {
             Ok(())
         }
     }
+
+    fn visit_while(&mut self, condition: &Expr, body: &Stmt) -> Result<(), ScriptError> {
+        while self.visit_expr(condition)?.is_truthy() {
+            self.visit_stmt(body)?;
+        }
+
+        Ok(())
+    }
 }
 
 fn convert_arithmetic_operands(
diff --git a/src/parse.rs b/src/parse.rs
index b6b443c..5e6dc3d 100644
--- a/src/parse.rs
+++ b/src/parse.rs
@@ -175,6 +175,8 @@ fn parse_statement<I: Iterator<Item = Token>>(
 ) -> Result<ParsedStatement, ParseError> {
     if let Ok(token) = match_token_kind!(iter, TokenKind::If) {
         parse_if_statement(iter, &token)
+    } else if let Ok(token) = match_token_kind!(iter, TokenKind::While) {
+        parse_while_statement(iter, &token)
     } else if match_token_kind!(iter, TokenKind::Print).is_ok() {
         parse_statement_containing_expression(
             iter,
@@ -205,7 +207,7 @@ fn parse_if_statement<I: Iterator<Item = Token>>(
     let condition = parse_expression(iter)?;
 
     match_token_kind!(iter, TokenKind::RightParen).map_err(|_| ParseError {
-        message: "Expected ')' after 'if'".to_string(),
+        message: "Expected ')' after 'if' condition".to_string(),
         line: if_token.line().into(),
     })?;
 
@@ -224,6 +226,31 @@ fn parse_if_statement<I: Iterator<Item = Token>>(
     Ok(ParsedStatement::Stmt { stmt: if_stmt })
 }
 
+fn parse_while_statement<I: Iterator<Item = Token>>(
+    iter: &mut Peekable<I>,
+    while_token: &Token,
+) -> Result<ParsedStatement, ParseError> {
+    match_token_kind!(iter, TokenKind::LeftParen).map_err(|_| ParseError {
+        message: "Expected '(' after 'while'".to_string(),
+        line: while_token.line().into(),
+    })?;
+
+    let condition = parse_expression(iter)?;
+
+    match_token_kind!(iter, TokenKind::RightParen).map_err(|_| ParseError {
+        message: "Expected ')' after 'while' condition".to_string(),
+        line: while_token.line().into(),
+    })?;
+
+    let body = parse_statement(iter).and_then(ensure_explicit_statement)?;
+    let while_stmt = Stmt::While {
+        condition,
+        body: Box::new(body),
+    };
+
+    Ok(ParsedStatement::Stmt { stmt: while_stmt })
+}
+
 fn parse_statement_containing_expression<
     I: Iterator<Item = Token>,
     F: Fn(Expr) -> Stmt,