Fucky wucky and some testy westies

6 files changed, 256 insertions, 77 deletions

diff --git a/crates/sloth/src/lexer.rs b/crates/sloth/src/lexer.rs
index 62b1a8d..cc17b3d 100644
--- a/crates/sloth/src/lexer.rs
+++ b/crates/sloth/src/lexer.rs
@@ -86,6 +86,7 @@ pub enum TokenType {
     Var,
 
     Fn,
+    Return,
 
     If,
     Else,
@@ -382,6 +383,7 @@ impl<'a> Iterator for Lexer<'a> {
                     "val" => TokenType::Val,
                     "var" => TokenType::Var,
                     "fn" => TokenType::Fn,
+                    "return" => TokenType::Return,
                     "if" => TokenType::If,
                     "else" => TokenType::Else,
                     "while" => TokenType::While,
diff --git a/crates/sloth/src/main.rs b/crates/sloth/src/main.rs
index 7821097..fd4a1eb 100644
--- a/crates/sloth/src/main.rs
+++ b/crates/sloth/src/main.rs
@@ -13,6 +13,7 @@ use std::{env, fs};
 
 use itertools::Itertools;
 use lexer::Lexer;
+use parser::AstParser;
 
 fn main() {
     let args = env::args().collect_vec();
@@ -29,7 +30,6 @@ fn main() {
         return;
     };
 
-    let _tokens = Lexer::new(&source).collect_vec();
-
-    // TODO: Write a parser
+    let tokens = Lexer::new(&source).collect_vec();
+    let _parser = AstParser::new(tokens);
 }
diff --git a/crates/sloth/src/parser/ast.rs b/crates/sloth/src/parser/ast.rs
index c008905..67bc9da 100644
--- a/crates/sloth/src/parser/ast.rs
+++ b/crates/sloth/src/parser/ast.rs
@@ -1,4 +1,3 @@
-use crate::lexer::{Token, TokenType};
 #[derive(Debug, PartialEq)]
 pub enum BinaryOp {
     Add,
@@ -61,17 +60,17 @@ pub enum Expr {
     Literal(Literal),
     Lambda, // TODO: Lambda bitch
 }
-
+#[derive(PartialEq, Debug)]
 pub struct FuncArgs {
     pub name: String,
     pub typ: Option<String>,
 }
-
+#[derive(PartialEq, Debug)]
 pub enum Stmt {
     ExprStmt(Expr),
     DefineFunction {
         ident: String,
-        args: Vec<FuncArgs>,
+        args: Option<Vec<FuncArgs>>,
         body: Vec<Stmt>,
         return_type: Option<String>,
     },
@@ -100,56 +99,7 @@ pub enum Stmt {
         condition: Expr,
         body: Vec<Stmt>,
     },
-}
-
-pub struct AstParser<'a> {
-    tokens: Vec<Token<'a>>,
-    index: usize,
-}
-
-/// Implementation containing utilities used by the parsers internal components
-impl<'a> AstParser<'a> {
-    pub fn new(tokens: Vec<Token<'a>>) -> Self {
-        Self { tokens, index: 0 }
-    }
-    pub fn peek(&self) -> &Token {
-        &self.tokens[self.index]
-    }
-
-    pub fn advance(&mut self) -> Option<&Token> {
-        if self.eof() {
-            return None;
-        }
-
-        self.index += 1;
-        Some(&self.tokens[self.index - 1])
-    }
-
-    pub fn advance_if(&mut self, next: impl FnOnce(&Token) -> bool) -> bool {
-        if self.eof() {
-            return false;
-        }
-
-        if next(self.peek()) {
-            self.advance();
-            return true;
-        }
-
-        false
-    }
-
-    pub fn advance_if_eq(&mut self, next: &TokenType) -> bool {
-        self.advance_if(|it| it.tt == *next)
-    }
-
-    pub fn consume(&mut self, next: TokenType, error: &str) {
-        if std::mem::discriminant(&self.peek().tt) != std::mem::discriminant(&next) {
-            panic!("{error}");
-        }
-        self.advance();
-    }
-
-    pub fn eof(&self) -> bool {
-        self.index >= self.tokens.len()
-    }
+    Return {
+        value: Expr,
+    },
 }
diff --git a/crates/sloth/src/parser/expr.rs b/crates/sloth/src/parser/expr.rs
index 7cce919..a23359d 100644
--- a/crates/sloth/src/parser/expr.rs
+++ b/crates/sloth/src/parser/expr.rs
@@ -1,4 +1,5 @@
-use super::ast::{AstParser, BinaryOp, Expr, Literal, UnaryOp};
+use super::ast::{BinaryOp, Expr, Literal, UnaryOp};
+use super::AstParser;
 use crate::lexer::TokenType;
 
 /// Implementation containing parsers internal components related to expressions
@@ -20,7 +21,7 @@ impl<'a> AstParser<'a> {
                 TokenType::Bang => UnaryOp::Not,
                 TokenType::Tilde => UnaryOp::BWComp,
                 TokenType::Minus => UnaryOp::Neg,
-                _ => UnaryOp::Neg, // TODO: Idk how to not have this shit
+                _ => panic!(), // TODO: Idk how to not have this shit
             };
 
             let rhs = self.unary();
@@ -72,7 +73,9 @@ impl<'a> AstParser<'a> {
             TokenType::Character(literal) => Expr::Literal(Literal::Char(literal)),
             TokenType::String(literal) => Expr::Literal(Literal::String(literal)),
             TokenType::Regex(literal) => Expr::Literal(Literal::Regex(literal)),
-            TokenType::Identifier(ident) => Expr::Variable(ident),
+            TokenType::Identifier(ident) => Expr::Variable(ident), /* TODO: make this shit check
+                                                                     * if it is calling a
+                                                                     * function dumbass */
             TokenType::OpeningParen => {
                 let expr = self.expression();
                 self.consume(TokenType::ClosingParen, "Must end expression with ')'");
@@ -114,7 +117,7 @@ macro_rules! binary_expr {
                     TokenType::BangEq => BinaryOp::NotEq,
                     TokenType::AmpAmp => BinaryOp::LogAnd,
                     TokenType::PipePipe => BinaryOp::LogOr,
-                    _ => BinaryOp::Add, // TODO: Idk how to not have this shit
+                    _ => panic!("fuck"), // TODO: Idk how to not have this shit
                 };
 
                 let rhs = self.$parent();
diff --git a/crates/sloth/src/parser/mod.rs b/crates/sloth/src/parser/mod.rs
index f74a3d7..c900168 100644
--- a/crates/sloth/src/parser/mod.rs
+++ b/crates/sloth/src/parser/mod.rs
@@ -1,3 +1,57 @@
 pub mod ast;
 pub mod expr;
 pub mod stmt;
+
+use crate::lexer::{Token, TokenType};
+
+pub struct AstParser<'a> {
+    tokens: Vec<Token<'a>>,
+    index: usize,
+}
+
+/// Implementation containing utilities used by the parsers internal components
+impl<'a> AstParser<'a> {
+    pub fn new(tokens: Vec<Token<'a>>) -> Self {
+        Self { tokens, index: 0 }
+    }
+    pub fn peek(&self) -> &Token {
+        &self.tokens[self.index]
+    }
+
+    pub fn advance(&mut self) -> Option<&Token> {
+        if self.eof() {
+            return None;
+        }
+
+        self.index += 1;
+        Some(&self.tokens[self.index - 1])
+    }
+
+    pub fn advance_if(&mut self, next: impl FnOnce(&Token) -> bool) -> bool {
+        if self.eof() {
+            return false;
+        }
+
+        if next(self.peek()) {
+            self.advance();
+            return true;
+        }
+
+        false
+    }
+
+    pub fn advance_if_eq(&mut self, next: &TokenType) -> bool {
+        self.advance_if(|it| it.tt == *next)
+    }
+
+    pub fn consume(&mut self, next: TokenType, error: &str) {
+        if std::mem::discriminant(&self.peek().tt) != std::mem::discriminant(&next) {
+            panic!("{error}");
+        }
+        self.advance();
+    }
+
+    pub fn eof(&self) -> bool {
+        self.index >= self.tokens.len()
+    }
+}
diff --git a/crates/sloth/src/parser/stmt.rs b/crates/sloth/src/parser/stmt.rs
index c5c0c6a..bd7b4b2 100644
--- a/crates/sloth/src/parser/stmt.rs
+++ b/crates/sloth/src/parser/stmt.rs
@@ -1,4 +1,5 @@
-use super::ast::{AstParser, Expr, Stmt};
+use super::ast::{Expr, FuncArgs, Stmt};
+use super::AstParser;
 use crate::lexer::TokenType;
 
 impl<'a> AstParser<'a> {
@@ -33,6 +34,14 @@ impl<'a> AstParser<'a> {
             return self.while_statement();
         }
 
+        if self.advance_if_eq(&TokenType::Fn) {
+            return self.function_statement();
+        }
+
+        if self.advance_if_eq(&TokenType::Return) {
+            return self.return_statement();
+        }
+
         // If we couldn't parse a statement return an expression statement
         self.expression_statement()
     }
@@ -119,7 +128,7 @@ impl<'a> AstParser<'a> {
     //     }
 
     //     Stmt::For { name: (binding), iter: (), body: (body) }
-    // } TODO: Fix this garbage
+    // } // TODO: Fix this garbage
 
     fn while_statement(&mut self) -> Stmt {
         let condition = self.expression();
@@ -140,20 +149,181 @@ impl<'a> AstParser<'a> {
         let expr = self.expression();
 
         // FIXME: Move assignment handling
-        if self.advance_if_eq(&TokenType::Eq) {
-            if let Expr::Variable(ident) = &expr {
-                let value = self.expression();
-
-                self.consume(TokenType::SemiColon, "Expected ';' at end of statement");
-                return Stmt::DefineVariable {
-                    name: (ident.clone()),
-                    value: (value),
-                    typ: (None),
-                };
-            }
-        }
+        // if self.advance_if_eq(&TokenType::Eq) {
+        //     if let Expr::Literal(_ident) = &expr {
+        //         let value = self.expression();
+
+        //         self.consume(TokenType::SemiColon, "Expected ';' at end of
+        // statement");         // return Stmt::DefineVariable {
+        //         //     name: (ident.clone()),
+        //         //     value: (value),
+        //         //     typ: (None),
+        //         // };
+        //         return Stmt::ExprStmt(expr);
+        //     }
+        // }
 
         self.consume(TokenType::SemiColon, "Expected ';' at end of statement");
         Stmt::ExprStmt(expr)
     }
+
+    fn function_statement(&mut self) -> Stmt {
+        let TokenType::Identifier(ident) = self.advance().unwrap().tt.clone() else {
+            panic!("Identifier expected after 'fn'");
+        };
+
+        self.consume(TokenType::OpeningParen, "Expected '(' after identifier");
+        let mut args: Vec<FuncArgs> = Vec::new();
+        while !self.eof() && self.peek().tt != TokenType::ClosingParen {
+            let TokenType::Identifier(name) = self.advance().unwrap().tt.clone() else {
+                panic!("Identifier expected after '('");
+            };
+
+            self.advance_if_eq(&TokenType::Comma);
+
+            let arg = FuncArgs {
+                name: (name),
+                typ: (None),
+            };
+            args.push(arg);
+        }
+        self.advance();
+        self.consume(TokenType::OpeningBrace, "Expected '{' after parameters");
+        let mut body = Vec::new();
+        while !self.eof() && self.peek().tt != TokenType::ClosingBrace {
+            body.push(self.statement());
+        }
+
+        Stmt::DefineFunction {
+            ident: (ident),
+            args: (Some(args)),
+            body: (body),
+            return_type: (None),
+        }
+    }
+
+    fn return_statement(&mut self) -> Stmt {
+        let expr = self.expression();
+        Stmt::Return { value: (expr) }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use itertools::Itertools;
+
+    use super::{AstParser, Expr, Stmt};
+    use crate::lexer::Lexer;
+    use crate::parser::ast::{BinaryOp, FuncArgs, Literal};
+
+    #[test]
+    fn basic_statement_a() {
+        let lexer = Lexer::new("var test_a = 5 + 3;");
+        let tokens = lexer.collect_vec();
+
+        let expected_ast = Stmt::DefineVariable {
+            name: ("test_a".to_string()),
+            value: (Expr::BinaryOp {
+                op: (BinaryOp::Add),
+                lhs: (Box::new(Expr::Literal(Literal::Integer(5)))),
+                rhs: (Box::new(Expr::Literal(Literal::Integer(3)))),
+            }),
+            typ: (None),
+        };
+
+        let mut parser = AstParser::new(tokens);
+        let generated_ast = parser.statement();
+
+        println!("Expected AST:\n{expected_ast:#?}\n\n");
+        println!("Generated AST:\n{generated_ast:#?}\n\n");
+
+        assert_eq!(expected_ast, generated_ast);
+    }
+
+    #[test]
+    fn basic_statement_b() {
+        let lexer = Lexer::new("val test_b = \"Hello World\";");
+        let tokens = lexer.collect_vec();
+
+        let expected_ast = Stmt::DefineValue {
+            name: ("test_b".to_string()),
+            value: (Expr::Literal(Literal::String("Hello World".to_string()))),
+            typ: (None),
+        };
+
+        let mut parser = AstParser::new(tokens);
+        let generated_ast = parser.statement();
+
+        println!("Expected AST:\n{expected_ast:#?}\n\n");
+        println!("Generated AST:\n{generated_ast:#?}\n\n");
+
+        assert_eq!(expected_ast, generated_ast);
+    }
+
+    #[test]
+    fn basic_statement_c() {
+        let lexer = Lexer::new("fn test_c (a, b, c) {\nreturn (b + a * c);\n}");
+        let tokens = lexer.collect_vec();
+        println!("{tokens:?}");
+
+        let expected_ast = Stmt::DefineFunction {
+            ident: ("test_c".to_string()),
+            args: Some(vec![
+                FuncArgs {
+                    name: ("a".to_string()),
+                    typ: None,
+                },
+                FuncArgs {
+                    name: ("b".to_string()),
+                    typ: None,
+                },
+                FuncArgs {
+                    name: ("c".to_string()),
+                    typ: None,
+                },
+            ]),
+            body: (vec![Stmt::Return {
+                value: (Expr::BinaryOp {
+                    op: BinaryOp::Add,
+                    lhs: Box::new(Expr::Variable("a".to_string())),
+                    rhs: Box::new(Expr::BinaryOp {
+                        op: BinaryOp::Mul,
+                        lhs: Box::new(Expr::Variable("b".to_string())),
+                        rhs: Box::new(Expr::Variable("c".to_string())),
+                    }),
+                }),
+            }]),
+            return_type: (None),
+        };
+
+        let mut parser = AstParser::new(tokens);
+        let generated_ast = parser.statement();
+
+        println!("Expected AST:\n{expected_ast:#?}\n\n");
+        println!("Generated AST:\n{generated_ast:#?}\n\n");
+
+        assert_eq!(expected_ast, generated_ast);
+    }
+    #[test]
+    fn basic_statement_d() {
+        let lexer = Lexer::new("while true {\nprint(\"Hello World\");}");
+        let tokens = lexer.collect_vec();
+        println!("{tokens:?}");
+
+        let expected_ast = Stmt::While {
+            condition: (Expr::Literal(Literal::Bool(true))),
+            body: (vec![Stmt::ExprStmt(Expr::Call {
+                ident: (Box::new(Expr::Literal(Literal::String("print".to_string())))),
+                args: (vec![Expr::Literal(Literal::String("Hello World".to_string()))]),
+            })]),
+        };
+
+        let mut parser = AstParser::new(tokens);
+        let generated_ast = parser.statement();
+
+        println!("Expected AST:\n{expected_ast:#?}\n\n");
+        println!("Generated AST:\n{generated_ast:#?}\n\n");
+
+        assert_eq!(expected_ast, generated_ast);
+    }
 }