diff --git a/src/parse.rs b/src/parse.rs index ee9aec1..57ae33f 100644 --- a/src/parse.rs +++ b/src/parse.rs @@ -19,6 +19,12 @@ enum ParsedStatement { ImplicitStmt { stmt: Stmt, line_number: usize }, } +struct ForLoopPreamble { + initializer: Option, + condition: Option, + increment: Option, +} + pub fn parse_program, F: FnMut(ScriptError)>( iter: I, mut on_error: F, @@ -269,63 +275,94 @@ fn parse_for_statement>( iter: &mut Peekable, for_token: &Token, ) -> Result { - // 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>( + iter: &mut Peekable, + for_token: &Token, +) -> Result { 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 initializer = optionally_parse_with_suffix(iter, TokenKind::SemiColon, |iter| { + let statement = if match_next_token(iter, &TokenKind::Var).is_some() { + parse_var_declaration(iter)? + } else { + parse_expression_statement(iter)? + }; - let condition = if match_token_kind!(iter, TokenKind::SemiColon).is_ok() { - None - } else { + ensure_explicit_statement(statement) + })?; + + 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() { - None - } else { + Ok(expr) + })?; + + 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)?)?; - let while_condition = condition.unwrap_or(Expr::Literal { + Ok(expr) + })?; + + Ok(ForLoopPreamble { + initializer, + condition, + increment, + }) +} + +fn optionally_parse_with_suffix< + I: Iterator, + R, + F: FnOnce(&mut Peekable) -> Result, +>( + iter: &mut Peekable, + suffix_kind: TokenKind, + parse: F, +) -> Result, 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::>(); 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>( .flatten() .collect::>(); - Ok(ParsedStatement::Stmt { - stmt: Stmt::Block { - statements: wrapping_body_statements, - }, - }) + Stmt::Block { + statements: wrapping_body_statements, + } } fn parse_statement_containing_expression<