feat: add while loops, break and continue statements

src/frontend/ast.rs

@@ -98,6 +98,10 @@ pub enum StmtKind<P: Phase = Untyped> {
         then: Box<Stmt<P>>,
         elze: Option<Box<Stmt<P>>>,
     },
+    While {
+        condition: Expr<P>,
+        body: Box<Stmt<P>>,
+    },
     Return {
         value: Option<Expr<P>>,
     },
@@ -110,6 +114,8 @@ pub enum StmtKind<P: Phase = Untyped> {
     Expression {
         expr: Expr<P>,
     },
+    Break,
+    Continue,
 }
 
 #[derive(Debug, PartialEq, Eq)]

src/frontend/lexer.rs

@@ -68,6 +68,9 @@ impl<'src> Lexer<'src> {
             "else" => TokenKind::Else,
             "return" => TokenKind::Return,
             "let" => TokenKind::Let,
+            "while" => TokenKind::While,
+            "break" => TokenKind::Break,
+            "continue" => TokenKind::Continue,
 
             "i8" => TokenKind::I8,
             "i16" => TokenKind::I16,
@@ -206,7 +209,7 @@ mod test {
     #[test]
     fn identifiers() {
         assert_eq!(
-            tokenize("HELLO _hello _0@ fn if else return let"),
+            tokenize("HELLO _hello _0@ fn if else return let while break continue"),
             vec![
                 Token::new(TokenKind::Identifier, "HELLO", Span::new(0, 5)),
                 Token::new(TokenKind::Identifier, "_hello", Span::new(6, 12)),
@@ -217,6 +220,9 @@ mod test {
                 Token::new(TokenKind::Else, "else", Span::new(23, 27)),
                 Token::new(TokenKind::Return, "return", Span::new(28, 34)),
                 Token::new(TokenKind::Let, "let", Span::new(35, 38)),
+                Token::new(TokenKind::While, "while", Span::new(39, 44)),
+                Token::new(TokenKind::Break, "break", Span::new(45, 50)),
+                Token::new(TokenKind::Continue, "continue", Span::new(51, 59)),
             ]
         )
     }

src/frontend/parser.rs

@@ -287,6 +287,10 @@ impl<'src> Parser<'src> {
     ///      | if_stmt
     ///      | return_stmt
     ///      | let_stmt
+    ///      | while_stmt
+    ///      | break_stmt
+    ///      | continue_stmt
+    ///      | expr_stmt
     ///      ;
     /// ```
     pub fn parse_stmt(&mut self) -> ParseResult<Stmt> {
@@ -295,8 +299,11 @@ impl<'src> Parser<'src> {
         match peek_token.kind {
             TokenKind::LBrace => self.parse_compound_stmt(),
             TokenKind::If => self.parse_if_stmt(),
+            TokenKind::While => self.parse_while_stmt(),
             TokenKind::Return => self.parse_return_stmt(),
             TokenKind::Let => self.parse_let_stmt(),
+            TokenKind::Break => self.parse_break_stmt(),
+            TokenKind::Continue => self.parse_continue_stmt(),
             _ => self.parse_expr_stmt(),
         }
     }
@@ -374,6 +381,26 @@ impl<'src> Parser<'src> {
         })
     }
 
+    /// Parses a while statement.
+    ///
+    /// ```ebnf
+    /// while_stmt = "while" expr compound_stmt ;
+    /// ```
+    fn parse_while_stmt(&mut self) -> ParseResult<Stmt> {
+        let while_token = self.expect(TokenKind::While)?;
+        let condition = self.parse_expr()?;
+        let body = self.parse_compound_stmt()?;
+        let span = while_token.span.join(body.span);
+
+        Ok(Stmt {
+            kind: StmtKind::While {
+                condition,
+                body: Box::new(body),
+            },
+            span,
+        })
+    }
+
     /// Parses an expression statement.
     ///
     /// ```ebnf
@@ -454,6 +481,38 @@ impl<'src> Parser<'src> {
         })
     }
 
+    /// Parses a break statement.
+    ///
+    /// ```ebnf
+    /// break_stmt = "break" ";" ;
+    /// ```
+    fn parse_break_stmt(&mut self) -> ParseResult<Stmt> {
+        let break_token = self.expect(TokenKind::Break)?;
+        let semi_token = self.expect(TokenKind::Semicolon)?;
+        let span = break_token.span.join(semi_token.span);
+
+        Ok(Stmt {
+            kind: StmtKind::Break,
+            span,
+        })
+    }
+
+    /// Parses a continue statement.
+    ///
+    /// ```ebnf
+    /// continue_stmt = "continue" ";" ;
+    /// ```
+    fn parse_continue_stmt(&mut self) -> ParseResult<Stmt> {
+        let continue_token = self.expect(TokenKind::Continue)?;
+        let semi_token = self.expect(TokenKind::Semicolon)?;
+        let span = continue_token.span.join(semi_token.span);
+
+        Ok(Stmt {
+            kind: StmtKind::Continue,
+            span,
+        })
+    }
+
     // ====== Pratt Parsing Implementation ======
 
     /// Parses an expression.
@@ -910,6 +969,54 @@ mod test {
         );
     }
 
+    #[test]
+    fn while_stmt() {
+        assert_eq!(
+            parse("while true { break; }", Parser::parse_stmt),
+            Success(Stmt {
+                kind: StmtKind::While {
+                    condition: Expr {
+                        kind: ExprKind::Boolean { value: true },
+                        ty: (),
+                        span: Span::new(6, 10)
+                    },
+                    body: Box::new(Stmt {
+                        kind: StmtKind::Compound {
+                            inner: vec![Stmt {
+                                kind: StmtKind::Break,
+                                span: Span::new(13, 19)
+                            }]
+                        },
+                        span: Span::new(11, 21)
+                    }),
+                },
+                span: Span::new(0, 21)
+            })
+        )
+    }
+
+    #[test]
+    fn break_stmt() {
+        assert_eq!(
+            parse("break;", Parser::parse_stmt),
+            Success(Stmt {
+                kind: StmtKind::Break,
+                span: Span::new(0, 6)
+            })
+        );
+    }
+
+    #[test]
+    fn continue_stmt() {
+        assert_eq!(
+            parse("continue;", Parser::parse_stmt),
+            Success(Stmt {
+                kind: StmtKind::Continue,
+                span: Span::new(0, 9)
+            })
+        );
+    }
+
     #[test]
     fn assign_expr_stmt() {
         assert_eq!(

src/frontend/sema.rs

@@ -78,6 +78,7 @@ pub struct Sema {
     deferred_unary_neg: Vec<(Span, Ty, Ty, Option<u64>)>,
     deferred_binary: Vec<(Span, Ty)>,
     deferred_literals: Vec<(Span, Ty)>,
+    loop_depth: usize,
 }
 
 impl Sema {
@@ -91,6 +92,7 @@ impl Sema {
             deferred_unary_neg: Vec::new(),
             deferred_binary: Vec::new(),
             deferred_literals: Vec::new(),
+            loop_depth: 0,
         }
     }
 
@@ -323,6 +325,47 @@ impl Sema {
                     span: stmt.span,
                 }
             }
+            StmtKind::While { condition, body } => {
+                let typed_condition = self.analyze_expr(condition);
+
+                if let Err(err) = self.unify(&typed_condition.ty, &Ty::Bool) {
+                    self.errors.push(SemanticError::new(err, condition.span));
+                }
+
+                self.loop_depth += 1;
+                let typed_body = self.analyze_stmt(body, expected_ret_ty);
+                self.loop_depth -= 1;
+
+                TypedStmt {
+                    kind: TypedStmtKind::While {
+                        condition: typed_condition,
+                        body: Box::new(typed_body),
+                    },
+                    span: stmt.span,
+                }
+            }
+            StmtKind::Break => {
+                if self.loop_depth == 0 {
+                    self.errors
+                        .push(SemanticError::new("`break` outside of a loop", stmt.span));
+                }
+                TypedStmt {
+                    kind: TypedStmtKind::Break,
+                    span: stmt.span,
+                }
+            }
+            StmtKind::Continue => {
+                if self.loop_depth == 0 {
+                    self.errors.push(SemanticError::new(
+                        "`continue` outside of a loop",
+                        stmt.span,
+                    ));
+                }
+                TypedStmt {
+                    kind: TypedStmtKind::Continue,
+                    span: stmt.span,
+                }
+            }
             StmtKind::Return { value } => {
                 if let Some(expr) = value {
                     let typed_expr = self.analyze_expr(expr);
@@ -604,6 +647,14 @@ impl Sema {
                 elze: elze.map(|s| Box::new(self.apply_subst_stmt(*s))),
             },
 
+            TypedStmtKind::While { condition, body } => TypedStmtKind::While {
+                condition: self.apply_subst_expr(condition),
+                body: Box::new(self.apply_subst_stmt(*body)),
+            },
+
+            TypedStmtKind::Break => TypedStmtKind::Break,
+            TypedStmtKind::Continue => TypedStmtKind::Continue,
+
             TypedStmtKind::Return { value } => TypedStmtKind::Return {
                 value: value.map(|e| self.apply_subst_expr(e)),
             },
@@ -933,4 +984,21 @@ mod test {
                 .any(|e| e.message.contains("invalid left-hand side of assignment"))
         );
     }
+
+    #[test]
+    fn valid_while() {
+        let src = "fn test() { while true { break; continue; } }";
+        assert!(analyze(src).is_ok());
+    }
+
+    #[test]
+    fn invalid_break() {
+        let src = "fn test() { break; }";
+        let errors = analyze(src).unwrap_err();
+        assert!(
+            errors
+                .iter()
+                .any(|e| e.message.contains("`break` outside of a loop"))
+        );
+    }
 }

src/frontend/token.rs

@@ -60,6 +60,9 @@ pub enum TokenKind {
     Else,
     Return,
     Let,
+    While,
+    Break,
+    Continue,
 
     // Types
     I8,
@@ -117,6 +120,9 @@ impl Display for TokenKind {
             TokenKind::Else => "`else`",
             TokenKind::Return => "`return`",
             TokenKind::Let => "`let`",
+            TokenKind::While => "`while`",
+            TokenKind::Break => "`break`",
+            TokenKind::Continue => "`continue`",
             TokenKind::I8 => "`i8`",
             TokenKind::I16 => "`i16`",
             TokenKind::I32 => "`i32`",