From d4a3468798c6b762747794875f371a7c02545db3 Mon Sep 17 00:00:00 2001
From: ghostlyzsh <pokeball.181818@gmail.com>
Date: Thu, 1 Aug 2024 23:27:59 -0500
Subject: [PATCH] return, break, continue

---
 kabel/grammar.ebnf  |  6 ++--
 kabel/src/parser.rs | 86 +++++++++++++++++++++++++++++++++++++++++----
 2 files changed, 83 insertions(+), 9 deletions(-)

diff --git a/kabel/grammar.ebnf b/kabel/grammar.ebnf
index bbe350b6fc4fd820eef32f24e09d276be368212b..5dfad1219b2a9754c03dcd9e9c7ce55d964ee67d 100644
--- a/kabel/grammar.ebnf
+++ b/kabel/grammar.ebnf
@@ -1,12 +1,14 @@
 program = { statement } ;
 
-statement = function | loop | while | if | expression_statement ;
+statement = function | return | loop | while | if | expression_statement ;
 
 function = "function" , identifier , "(" , { identifier , "," , } ")" , block
+return = "return" , expression , ";" ;
 
 loop = "loop" , block ;
-
 while = "while" , "(" , expression , ")" , block ;
+break = "break" , ";" ;
+continue = "continue" , ";" ;
 
 if = "if" , "(" , expression , ")" , block [ , "else" , ( if | block ) ] ;
 
diff --git a/kabel/src/parser.rs b/kabel/src/parser.rs
index 3ec8d8dba8f840b20e0e840a54f7b80feb095a90..777def36f407818addb41b8e953dfc2eb4b0de28 100644
--- a/kabel/src/parser.rs
+++ b/kabel/src/parser.rs
@@ -49,8 +49,11 @@ impl Parser {
         match self.peek()?.token_type {
             TokenType::Ident(ident) => match ident.as_str() {
                 "function" => self.function_statement(),
+                "return" => self.return_statement(),
                 "loop" => self.loop_statement(),
                 "while" => self.while_statement(),
+                "break" => self.break_statement(),
+                "continue" => self.continue_statement(),
                 "if" => self.if_statement(),
                 _ => self.expression_statement(),
             },
@@ -100,6 +103,37 @@ impl Parser {
         }
     }
 
+    pub fn return_statement(&mut self) -> Result<AST, KabelError> {
+        let return_ident = self.read_token()?;
+        if let TokenType::Semicolon = self.peek()?.token_type {
+            let semicolon = self.read_token()?;
+            return Ok(AST {
+                ast_type: ASTType::Return(None),
+                start: return_ident.start,
+                end: semicolon.end,
+                line: return_ident.line,
+                column: return_ident.column,
+            });
+        }
+        let expression = self.expression()?;
+        let semicolon = self.read_token()?;
+        if let TokenType::Semicolon = semicolon.token_type {
+            Ok(AST {
+                ast_type: ASTType::Return(Some(Box::new(expression))),
+                start: return_ident.start,
+                end: semicolon.end,
+                line: return_ident.line,
+                column: return_ident.column,
+            })
+        } else {
+            return Err(unexpected_token!(
+                self,
+                "Expected ; found {}",
+                semicolon
+            ));
+        }
+    }
+
     pub fn loop_statement(&mut self) -> Result<AST, KabelError> {
         let loop_ident = self.read_token()?;
         let block = self.block()?;
@@ -135,6 +169,38 @@ impl Parser {
         }
     }
 
+    pub fn break_statement(&mut self) -> Result<AST, KabelError> {
+        let break_ident = self.read_token()?;
+        let semicolon = self.read_token()?;
+        if let TokenType::Semicolon = semicolon.token_type {
+            Ok(AST {
+                ast_type: ASTType::Break,
+                start: break_ident.start,
+                end: semicolon.end,
+                line: break_ident.line, 
+                column: break_ident.column,
+            })
+        } else {
+            Err(unexpected_token!(self, "Expected ; found {}", semicolon))
+        }
+    }
+
+    pub fn continue_statement(&mut self) -> Result<AST, KabelError> {
+        let continue_ident = self.read_token()?;
+        let semicolon = self.read_token()?;
+        if let TokenType::Semicolon = semicolon.token_type {
+            Ok(AST {
+                ast_type: ASTType::Continue,
+                start: continue_ident.start,
+                end: semicolon.end,
+                line: continue_ident.line, 
+                column: continue_ident.column,
+            })
+        } else {
+            Err(unexpected_token!(self, "Expected ; found {}", semicolon))
+        }
+    }
+
     pub fn if_statement(&mut self) -> Result<AST, KabelError> {
         let if_ident = self.read_token()?;
         let left_paren = self.read_token()?;
@@ -756,16 +822,22 @@ pub struct AST {
 pub enum ASTType {
     Program(Vec<AST>),
 
-    Function(Box<AST>, Vec<AST>, Box<AST>), // name, args, block
-    Loop(Box<AST>),                         // block
-    While(Box<AST>, Box<AST>),              // condition, block
-    If(Box<AST>, Box<AST>, Option<Box<AST>>), // condition, block, else/else if
-    Block(Vec<AST>),                        // statements
-
-    Decl(Box<AST>, Box<AST>), // identifier, expression
+    // statements
+    Function(Box<AST>, Vec<AST>, Box<AST>),     // name, args, block
+    Return(Option<Box<AST>>),                   // expression
+    Loop(Box<AST>),                             // block
+    While(Box<AST>, Box<AST>),                  // condition, block
+    Break,
+    Continue,
+    If(Box<AST>, Box<AST>, Option<Box<AST>>),   // condition, block, else/else if
+    Block(Vec<AST>),                            // statements
+
+    // expressions
+    Decl(Box<AST>, Box<AST>),                   // identifier, expression
     Binary(Box<AST>, BinOp, Box<AST>),
     Unary(UnOp, Box<AST>),
 
+    // primary
     Group(Box<AST>),
     Lit(Lit),
     Call(Box<AST>, Vec<AST>),