For loops and lists

4 files changed, 128 insertions, 27 deletions

diff --git a/crates/sloth/src/parser/ast.rs b/crates/sloth/src/parser/ast.rs
index fb03634..99bbf52 100644
--- a/crates/sloth/src/parser/ast.rs
+++ b/crates/sloth/src/parser/ast.rs
@@ -22,6 +22,7 @@ pub enum BinaryOp {
     NotEq,
     LogAnd,
     LogOr,
+    Range,
 }
 #[derive(Debug, PartialEq)]
 pub enum UnaryOp {
diff --git a/crates/sloth/src/parser/expr.rs b/crates/sloth/src/parser/expr.rs
index 866663a..0cc3a93 100644
--- a/crates/sloth/src/parser/expr.rs
+++ b/crates/sloth/src/parser/expr.rs
@@ -85,6 +85,18 @@ impl<'a> AstParser<'a> {
                 self.consume(TokenType::ClosingParen, "Must end expression with ')'");
                 Expr::Grouping(Box::new(expr))
             }
+            TokenType::OpeningBracket => {
+                let mut expr: Vec<Expr> = Vec::new();
+
+                while !self.eof() && self.peek().tt != TokenType::ClosingBracket {
+                    let exp = self.expression();
+                    expr.push(exp);
+
+                    self.advance_if_eq(&TokenType::Comma);
+                }
+                self.consume(TokenType::ClosingBracket, "Expected ']' at end of list");
+                Expr::Literal(Literal::List(expr))
+            }
             _ => unimplemented!("{:?}", self.peek()),
         }
     }
@@ -106,6 +118,7 @@ macro_rules! binary_expr {
                     TokenType::StarStar => BinaryOp::Pow,
                     TokenType::Slash => BinaryOp::Div,
                     TokenType::Perc => BinaryOp::Mod,
+                    TokenType::DotDot => BinaryOp::Range,
 
                     TokenType::LtLt => BinaryOp::BWSftRight,
                     TokenType::GtGt => BinaryOp::BWSftLeft,
@@ -121,7 +134,7 @@ macro_rules! binary_expr {
                     TokenType::BangEq => BinaryOp::NotEq,
                     TokenType::AmpAmp => BinaryOp::LogAnd,
                     TokenType::PipePipe => BinaryOp::LogOr,
-                    _ => panic!("fuck"), // TODO: Idk how to not have this shit
+                    _ => panic!("uh oh spagghetio"), // TODO: Idk how to not have this shit
                 };
 
                 let rhs = self.$parent();
@@ -142,7 +155,8 @@ macro_rules! binary_expr {
 impl<'a> AstParser<'a> {
     // Binary expressions in order of precedence from lowest to highest.
     binary_expr!(logical_or      , logical_and     , (TokenType::PipePipe));
-    binary_expr!(logical_and     , equality        , (TokenType::AmpAmp));
+    binary_expr!(logical_and     , range           , (TokenType::AmpAmp));
+    binary_expr!(range           , equality         , (TokenType::DotDot));
     binary_expr!(equality        , comparison      , (TokenType::BangEq | TokenType::EqEq));
     binary_expr!(comparison      , bitwise_shifting, (TokenType::Lt     | TokenType::Gt    | TokenType::LtEq | TokenType::GtEq));
     binary_expr!(bitwise_shifting, additive        , (TokenType::LtLt   | TokenType::GtGt));
@@ -212,4 +226,42 @@ mod tests {
 
         assert_eq!(expected_ast, generated_ast);
     }
+    #[test]
+    fn basic_expression_c() {
+        let lexer = Lexer::new("[1, 2, 3]");
+        let tokens = lexer.collect_vec();
+
+        let expected_ast = Expr::Literal(Literal::List(vec![
+            Expr::Literal(Literal::Integer(1)),
+            Expr::Literal(Literal::Integer(2)),
+            Expr::Literal(Literal::Integer(3)),
+        ]));
+
+        let mut parser = AstParser::new(tokens);
+        let generated_ast = parser.expression();
+
+        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_expression_d() {
+        let lexer = Lexer::new("1 .. 17");
+        let tokens = lexer.collect_vec();
+
+        let expected_ast = Expr::BinaryOp {
+            op: (BinaryOp::Range),
+            lhs: (Box::new(Expr::Literal(Literal::Integer(1)))),
+            rhs: (Box::new(Expr::Literal(Literal::Integer(17)))),
+        };
+
+        let mut parser = AstParser::new(tokens);
+        let generated_ast = parser.expression();
+
+        println!("Expected AST:\n{expected_ast:#?}\n\n");
+        println!("Generated AST:\n{generated_ast:#?}\n\n");
+
+        assert_eq!(expected_ast, generated_ast);
+    }
 }
diff --git a/crates/sloth/src/parser/stmt.rs b/crates/sloth/src/parser/stmt.rs
index 6c02f35..0d28131 100644
--- a/crates/sloth/src/parser/stmt.rs
+++ b/crates/sloth/src/parser/stmt.rs
@@ -27,9 +27,9 @@ impl<'a> AstParser<'a> {
             return self.if_statement();
         }
 
-        // if self.advance_if_eq(&TokenType::For) {
-        //     return self.for_statement();
-        // }
+        if self.advance_if_eq(&TokenType::For) {
+            return self.for_statement();
+        }
 
         if self.advance_if_eq(&TokenType::While) {
             return self.while_statement();
@@ -169,31 +169,39 @@ impl<'a> AstParser<'a> {
         } // TODO: implement else if and else
     }
 
-    // fn for_statement(&mut self) -> Stmt {
-    //     let binding = self.expression();
-    //     let Expr::Variable(binding) = binding else {
-    //         panic!("Left side of for statement must be identifier");
-    //     };
+    fn for_statement(&mut self) -> Stmt {
+        let binding = self.expression();
+        let Expr::Variable(binding) = binding else {
+            panic!("Left side of for statement must be identifier");
+        };
 
-    //     self.consume(
-    //         TokenType::In,
-    //         "Expected 'in' in between identifier and range",
-    //     );
+        self.consume(
+            TokenType::In,
+            "Expected 'in' in between identifier and range",
+        );
 
-    //     let range_start = self.expression();
-    //     self.consume(
-    //         TokenType::DotDot,
-    //         "Expected '..' denoting min and max of range",
-    //     );
-    //     let range_end = self.expression();
+        // let range_start = self.expression();
+        // self.consume(
+        //     TokenType::DotDot,
+        //     "Expected '..' denoting min and max of range",
+        // );
+        // let range_end = self.expression();
 
-    //     let mut body = Vec::new();
-    //     while !self.eof() && self.peek().tt != TokenType::ClosingBrace {
-    //         body.push(self.statement());
-    //     }
+        let expr = self.expression();
+
+        self.consume(TokenType::OpeningBrace, "Expected '{' after iterator");
+
+        let mut body = Vec::new();
+        while !self.eof() && self.peek().tt != TokenType::ClosingBrace {
+            body.push(self.statement());
+        }
 
-    //     Stmt::For { name: (binding), iter: (), body: (body) }
-    // } // TODO: Fix this garbage
+        Stmt::For {
+            name: (binding),
+            iter: (expr),
+            body: (body),
+        }
+    } // TODO: Fix this garbage
 
     fn while_statement(&mut self) -> Stmt {
         let condition = self.expression();
@@ -241,7 +249,7 @@ impl<'a> AstParser<'a> {
         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 '('");
+                panic!("parameter expected after '('");
             };
 
             let mut typ: Option<String> = None;
@@ -580,4 +588,43 @@ mod tests {
 
         assert_eq!(expected_ast, generated_ast);
     }
+
+    #[test]
+    fn basic_statement_h() {
+        let lexer = Lexer::new("for i in 1 .. 3 {\nfor j in [1, 2, 3] {\nprint(j*i);}}");
+        let tokens = lexer.collect_vec();
+
+        let expected_ast = Stmt::For {
+            name: ("i".to_string()),
+            iter: (Expr::BinaryOp {
+                op: (BinaryOp::Range),
+                lhs: (Box::new(Expr::Literal(Literal::Integer(1)))),
+                rhs: (Box::new(Expr::Literal(Literal::Integer(3)))),
+            }),
+            body: (vec![Stmt::For {
+                name: ("j".to_string()),
+                iter: (Expr::Literal(Literal::List(vec![
+                    Expr::Literal(Literal::Integer(1)),
+                    Expr::Literal(Literal::Integer(2)),
+                    Expr::Literal(Literal::Integer(3)),
+                ]))),
+                body: (vec![Stmt::ExprStmt(Expr::Call {
+                    ident: Box::new(Expr::Literal(Literal::String("print".to_string()))),
+                    args: (vec![Expr::BinaryOp {
+                        op: (BinaryOp::Mul),
+                        lhs: (Box::new(Expr::Variable("j".to_string()))),
+                        rhs: (Box::new(Expr::Variable("i".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);
+    }
 }
diff --git a/documentation/order.txt b/documentation/order.txt
index ba22100..3c06ce0 100644
--- a/documentation/order.txt
+++ b/documentation/order.txt
@@ -13,6 +13,7 @@
 | bitwise or     | |         | Left       |
 | comparison     | < > <= >= | Left       |
 | equality       | == !=     | Left       |
+| range          | ..        | Left       |
 | logical and    | &&        | Left       |
 | logical or     | ||        | Left       |