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>( ) -> Result { 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>( + iter: &mut Peekable, +) -> Result { + parse_statement_containing_expression( + iter, + |expression| Stmt::Expression { expression }, + // We should print implicit statements + |expression| Stmt::Print { expression }, + ) +} + +fn parse_print_statement>( + iter: &mut Peekable, +) -> Result { + parse_statement_containing_expression( + iter, + |expression| Stmt::Print { expression }, + |expression| Stmt::Print { expression }, + ) +} + fn parse_if_statement>( iter: &mut Peekable, if_token: &Token, @@ -251,6 +265,85 @@ fn parse_while_statement>( Ok(ParsedStatement::Stmt { stmt: while_stmt }) } +fn parse_for_statement>( + iter: &mut Peekable, + for_token: &Token, +) -> Result { + // 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::>(); + + let wrapping_body_statements = [ + initializer, + Some(Stmt::While { + condition: while_condition, + body: Box::new(Stmt::Block { + statements: body_statements, + }), + }), + ] + .into_iter() + .flatten() + .collect::>(); + + Ok(ParsedStatement::Stmt { + stmt: Stmt::Block { + statements: wrapping_body_statements, + }, + }) +} + fn parse_statement_containing_expression< I: Iterator, F: Fn(Expr) -> Stmt,