Break out parsing for for loop components

src/parse.rs

@@ -19,6 +19,12 @@ enum ParsedStatement {
     ImplicitStmt { stmt: Stmt, line_number: usize },
 }
 
+struct ForLoopPreamble {
+    initializer: Option<Stmt>,
+    condition: Option<Expr>,
+    increment: Option<Expr>,
+}
+
 pub fn parse_program<I: Iterator<Item = Token>, F: FnMut(ScriptError)>(
     iter: I,
     mut on_error: F,
@@ -269,63 +275,94 @@ fn parse_for_statement<I: Iterator<Item = Token>>(
     iter: &mut Peekable<I>,
     for_token: &Token,
 ) -> Result<ParsedStatement, ParseError> {
-    // TODO: this function is kinda messy
+    let preamble = parse_for_loop_preamble(iter, for_token)?;
+    let body = ensure_explicit_statement(parse_statement(iter)?)?;
+    let loop_stmt = desugar_for_loop(preamble, body, for_token);
+
+    Ok(ParsedStatement::Stmt { stmt: loop_stmt })
+}
+
+fn parse_for_loop_preamble<I: Iterator<Item = Token>>(
+    iter: &mut Peekable<I>,
+    for_token: &Token,
+) -> Result<ForLoopPreamble, ParseError> {
     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() {
+    let initializer = optionally_parse_with_suffix(iter, TokenKind::SemiColon, |iter| {
-        None
+        let statement = if match_next_token(iter, &TokenKind::Var).is_some() {
-    } else if iter
+            parse_var_declaration(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(
+            parse_expression_statement(iter)?
-            iter,
-        )?)?)
         };
 
-    let condition = if match_token_kind!(iter, TokenKind::SemiColon).is_ok() {
+        ensure_explicit_statement(statement)
-        None
+    })?;
-    } else {
+
+    let condition = optionally_parse_with_suffix(iter, TokenKind::SemiColon, |iter| {
         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() {
+        Ok(expr)
-        None
+    })?;
-    } else {
+
+    let increment = optionally_parse_with_suffix(iter, TokenKind::RightParen, |iter| {
         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)?)?;
+        Ok(expr)
-    let while_condition = condition.unwrap_or(Expr::Literal {
+    })?;
+
+    Ok(ForLoopPreamble {
+        initializer,
+        condition,
+        increment,
+    })
+}
+
+fn optionally_parse_with_suffix<
+    I: Iterator<Item = Token>,
+    R,
+    F: FnOnce(&mut Peekable<I>) -> Result<R, ParseError>,
+>(
+    iter: &mut Peekable<I>,
+    suffix_kind: TokenKind,
+    parse: F,
+) -> Result<Option<R>, ParseError> {
+    if match_token_kind_eq(iter, &[suffix_kind]).is_some() {
+        Ok(None)
+    } else {
+        let res = parse(iter)?;
+        Ok(Some(res))
+    }
+}
+
+fn desugar_for_loop(preamble: ForLoopPreamble, body: Stmt, for_token: &Token) -> Stmt {
+    let while_condition = preamble.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 }),
+        preamble
+            .increment
+            .map(|expression| Stmt::Expression { expression }),
     ]
     .into_iter()
     .flatten()
     .collect::<Vec<_>>();
 
     let wrapping_body_statements = [
-        initializer,
+        preamble.initializer,
         Some(Stmt::While {
             condition: while_condition,
             body: Box::new(Stmt::Block {
@@ -337,11 +374,9 @@ fn parse_for_statement<I: Iterator<Item = Token>>(
     .flatten()
     .collect::<Vec<_>>();
 
-    Ok(ParsedStatement::Stmt {
+    Stmt::Block {
-        stmt: Stmt::Block {
         statements: wrapping_body_statements,
-        },
+    }
-    })
 }
 
 fn parse_statement_containing_expression<