From ccd5de67202ae40fb7f32f8540bb695f100bb57a Mon Sep 17 00:00:00 2001
From: Nick Krichevsky <nick@ollien.com>
Date: Sat, 25 May 2024 10:49:24 -0400
Subject: [PATCH] Add for loop support

---
 src/parse.rs | 115 ++++++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 104 insertions(+), 11 deletions(-)

diff --git a/src/parse.rs b/src/parse.rs
index 5e6dc3d..ee9aec1 100644
--- a/src/parse.rs
+++ b/src/parse.rs
@@ -175,26 +175,40 @@ 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::For) {
+        parse_for_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,
-            |expression| Stmt::Print { expression },
-            |expression| Stmt::Print { expression },
-        )
+        parse_print_statement(iter)
     } else if match_token_kind!(iter, TokenKind::LeftBrace).is_ok() {
         parse_block_statement(iter).map(|stmt| ParsedStatement::Stmt { stmt })
     } else {
-        parse_statement_containing_expression(
-            iter,
-            |expression| Stmt::Expression { expression },
-            // We should print implicit statements
-            |expression| Stmt::Print { expression },
-        )
+        parse_expression_statement(iter)
     }
 }
 
+fn parse_expression_statement<I: Iterator<Item = Token>>(
+    iter: &mut Peekable<I>,
+) -> Result<ParsedStatement, ParseError> {
+    parse_statement_containing_expression(
+        iter,
+        |expression| Stmt::Expression { expression },
+        // We should print implicit statements
+        |expression| Stmt::Print { expression },
+    )
+}
+
+fn parse_print_statement<I: Iterator<Item = Token>>(
+    iter: &mut Peekable<I>,
+) -> Result<ParsedStatement, ParseError> {
+    parse_statement_containing_expression(
+        iter,
+        |expression| Stmt::Print { expression },
+        |expression| Stmt::Print { expression },
+    )
+}
+
 fn parse_if_statement<I: Iterator<Item = Token>>(
     iter: &mut Peekable<I>,
     if_token: &Token,
@@ -251,6 +265,85 @@ fn parse_while_statement<I: Iterator<Item = Token>>(
     Ok(ParsedStatement::Stmt { stmt: while_stmt })
 }
 
+fn parse_for_statement<I: Iterator<Item = Token>>(
+    iter: &mut Peekable<I>,
+    for_token: &Token,
+) -> Result<ParsedStatement, ParseError> {
+    // TODO: this function is kinda messy
+    match_token_kind!(iter, TokenKind::LeftParen).map_err(|_| ParseError {
+        message: "Expected '(' after 'for'".to_string(),
+        line: for_token.line().into(),
+    })?;
+
+    let initializer = if match_token_kind!(iter, TokenKind::SemiColon).is_ok() {
+        None
+    } else if iter
+        .peek()
+        .is_some_and(|token| token.kind() == &TokenKind::Var)
+    {
+        Some(ensure_explicit_statement(parse_declaration(iter)?)?)
+    } else {
+        Some(ensure_explicit_statement(parse_expression_statement(
+            iter,
+        )?)?)
+    };
+
+    let condition = if match_token_kind!(iter, TokenKind::SemiColon).is_ok() {
+        None
+    } else {
+        let expr = parse_expression(iter)?;
+        match_token_kind!(iter, TokenKind::SemiColon).map_err(|_| ParseError {
+            message: "Expected ';' after for loop condition".to_string(),
+            line: for_token.line().into(),
+        })?;
+        Some(expr)
+    };
+
+    let increment = if match_token_kind!(iter, TokenKind::RightParen).is_ok() {
+        None
+    } else {
+        let expr = parse_expression(iter)?;
+        match_token_kind!(iter, TokenKind::RightParen).map_err(|_| ParseError {
+            message: "Expected ')' after for conditions".to_string(),
+            line: for_token.line().into(),
+        })?;
+        Some(expr)
+    };
+
+    let body = ensure_explicit_statement(parse_statement(iter)?)?;
+    let while_condition = condition.unwrap_or(Expr::Literal {
+        value: LiteralValue::True,
+        // Kinda wrong, but not completely off either
+        token: for_token.clone(),
+    });
+    let body_statements = [
+        Some(body),
+        increment.map(|expression| Stmt::Expression { expression }),
+    ]
+    .into_iter()
+    .flatten()
+    .collect::<Vec<_>>();
+
+    let wrapping_body_statements = [
+        initializer,
+        Some(Stmt::While {
+            condition: while_condition,
+            body: Box::new(Stmt::Block {
+                statements: body_statements,
+            }),
+        }),
+    ]
+    .into_iter()
+    .flatten()
+    .collect::<Vec<_>>();
+
+    Ok(ParsedStatement::Stmt {
+        stmt: Stmt::Block {
+            statements: wrapping_body_statements,
+        },
+    })
+}
+
 fn parse_statement_containing_expression<
     I: Iterator<Item = Token>,
     F: Fn(Expr) -> Stmt,