pog pog pogpogpogpogpogpogpogpog pog pog pogpogpogpogpogpog

11 files changed, 789 insertions, 1070 deletions

diff --git a/examples/hello.sloth b/examples/hello.sloth
index 8201dae..0da2fd0 100644
--- a/examples/hello.sloth
+++ b/examples/hello.sloth
@@ -1,6 +1,7 @@
-print("Hello World!");
-
-# Comment
-for greeting in ["Hello", "Hola", "你好"] {
-    print(greeting + " World!");
+fn main() {
+    var i: Int = 10;
+    while i > 0 {
+        i = i - 1;
+    }
 }
+
diff --git a/sloth/src/codegen.rs b/sloth/src/codegen.rs
index 217be71..d4a5442 100644
--- a/sloth/src/codegen.rs
+++ b/sloth/src/codegen.rs
@@ -5,7 +5,7 @@ use inkwell::types::{AnyTypeEnum, BasicMetadataTypeEnum, BasicType, BasicTypeEnu
 use inkwell::values::FunctionValue;
 use itertools::Itertools;
 
-use crate::parser::ast::{FuncArgs, Stmt};
+use crate::parser::ast::{FunctionInput, Stmt};
 use crate::symbol::SymbolTableStack;
 
 /// A module codegen is a struct designated to compiling a single module
@@ -38,7 +38,7 @@ impl<'ctx, 'a> ModuleCodegen<'ctx, 'a> {
     pub fn compile_function(
         &mut self,
         ident: String,
-        args: Vec<FuncArgs>,
+        args: Vec<FunctionInput>,
         return_type: Option<String>,
         body: &[Stmt],
     ) {
@@ -52,15 +52,12 @@ impl<'ctx, 'a> ModuleCodegen<'ctx, 'a> {
 
     fn compile_function_type(
         &self,
-        args: &[FuncArgs],
+        args: &[FunctionInput],
         return_type: Option<&str>,
     ) -> FunctionType<'ctx> {
         let args = args
             .iter()
-            .map(|it| {
-                self.compile_basic_metadata_type(it.typ.as_ref().unwrap())
-                    .unwrap()
-            })
+            .map(|it| self.compile_basic_metadata_type(&it.typ).unwrap())
             .collect_vec();
 
         match return_type {
diff --git a/sloth/src/compiler.rs b/sloth/src/compiler.rs
index 2a02ce9..e52e3e2 100644
--- a/sloth/src/compiler.rs
+++ b/sloth/src/compiler.rs
@@ -41,13 +41,13 @@ impl Compiler {
 
     fn resolve_globals(&mut self, code: &[Stmt]) {
         for stmt in code.iter() {
-            if let Stmt::DefineFunction { ident, .. } = stmt {
-                let symbol = Symbol {
-                    typ: Some(SymbolType::Function),
-                };
-
-                self.symbol_table.insert(ident, symbol);
-            }
+            // if let Stmt::DefineFunction { ident, .. } = stmt {
+            //     let symbol = Symbol {
+            //         typ: Some(SymbolType::Function),
+            //     };
+            //
+            //     self.symbol_table.insert(ident, symbol);
+            // }
         }
     }
 }
diff --git a/sloth/src/compiler/mod.rs.disabled b/sloth/src/compiler/mod.rs.disabled
deleted file mode 100644
index 87c0618..0000000
--- a/sloth/src/compiler/mod.rs.disabled
+++ /dev/null
@@ -1,131 +0,0 @@
-#![allow(unused)]
-
-use std::collections::HashMap;
-use std::path::Path;
-use std::vec;
-
-use inkwell::builder::Builder;
-use inkwell::context::Context;
-use inkwell::module::Module;
-use inkwell::targets::{
-    CodeModel, FileType, InitializationConfig, RelocMode, Target, TargetMachine,
-};
-use inkwell::values::IntValue;
-use inkwell::OptimizationLevel;
-
-use crate::parser::ast::{BinaryOp, Expr, FuncArgs, Literal, Stmt, UnaryOp};
-
-pub struct Compiler<'ctx> {
-    context: &'ctx Context,
-    builder: Builder<'ctx>,
-    module: Module<'ctx>,
-}
-
-impl<'ctx> Compiler<'ctx> {
-    pub fn new(context: &'ctx Context) -> Self {
-        let builder = context.create_builder();
-        let module = context.create_module("sloth");
-
-        Self {
-            context,
-            builder,
-            module,
-        }
-    }
-
-    pub fn compile(&self, src: Vec<Stmt>) {
-        for stmt in src {
-            match stmt {
-                Stmt::DefineFunction {
-                    ident,
-                    args,
-                    body,
-                    return_type,
-                } => {
-                    self.compile_function(&ident, &args, return_type.is_some(), body);
-                }
-                _ => panic!("You may only define a function top level"),
-            }
-        }
-
-        Target::initialize_native(&InitializationConfig::default()).unwrap();
-
-        let triple = TargetMachine::get_default_triple();
-        let target = Target::from_triple(&triple).unwrap();
-        let machine = target
-            .create_target_machine(
-                &triple,
-                "x86-64",
-                "",
-                OptimizationLevel::None,
-                RelocMode::Default,
-                CodeModel::Default,
-            )
-            .unwrap();
-
-        self.module.set_triple(&triple);
-        machine
-            .write_to_file(&self.module, FileType::Object, Path::new("output.o"))
-            .unwrap();
-    }
-
-    fn compile_function(&self, identifier: &str, args: &[FuncArgs], returns: bool, src: Vec<Stmt>) {
-        let void_type = self.context.void_type();
-        let i64_type = self.context.i64_type();
-
-        let function_type = if returns {
-            i64_type.fn_type(&vec![i64_type.into(); args.len()], false)
-        } else {
-            void_type.fn_type(&vec![i64_type.into(); args.len()], false)
-        };
-        let function = self.module.add_function(identifier, function_type, None);
-
-        let basic_block = self.context.append_basic_block(function, "body");
-
-        self.builder.position_at_end(basic_block);
-
-        let mut arg_values = HashMap::<String, IntValue>::new();
-        for (i, arg) in args.iter().enumerate() {
-            arg_values.insert(
-                arg.name.clone(),
-                function.get_nth_param(i as u32).unwrap().into_int_value(),
-            );
-        }
-
-        for stmt in src {
-            match stmt {
-                Stmt::Return { value } => match value {
-                    Expr::BinaryOp { op, lhs, rhs } => {
-                        let lhs = match *lhs {
-                            Expr::Variable(a) => arg_values[&a],
-                            _ => unimplemented!(),
-                        };
-
-                        let rhs = match *rhs {
-                            Expr::Variable(a) => arg_values[&a],
-                            _ => unimplemented!(),
-                        };
-
-                        let res = match op {
-                            BinaryOp::Add => self.builder.build_int_add(lhs, rhs, "addop"),
-                            BinaryOp::Sub => self.builder.build_int_sub(lhs, rhs, "subop"),
-                            _ => unimplemented!(),
-                        };
-
-                        self.builder.build_return(Some(&res));
-                        return;
-                    }
-                    Expr::Variable(name) => {
-                        let var = arg_values[&name];
-                        self.builder.build_return(Some(&var));
-                        return;
-                    }
-                    _ => unimplemented!(),
-                },
-                _ => unimplemented!(),
-            }
-        }
-
-        self.builder.build_return(None);
-    }
-}
diff --git a/sloth/src/lexer.rs b/sloth/src/lexer.rs
index 0afaf1c..7ce95b5 100644
--- a/sloth/src/lexer.rs
+++ b/sloth/src/lexer.rs
@@ -82,6 +82,7 @@ pub enum TokenType {
     FatArrow, // =>
 
     // Keywords
+    Const,
     Val,
     Var,
 
@@ -101,18 +102,27 @@ pub enum TokenType {
 
     As,
 
-    // Literals
-    Integer(i128),
+    // Other
+    Literal(Literal),
+    Identifier(String),
+
+    // Utility
+    Error(LexerError),
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum Literal {
+    Integer(i64),
     Float(f64),
     Boolean(bool),
     Character(char),
     String(String),
-    Regex(String),
-
-    Identifier(String),
+}
 
-    // Utility
-    Error(LexerError),
+impl From<Literal> for TokenType {
+    fn from(value: Literal) -> Self {
+        Self::Literal(value)
+    }
 }
 
 #[derive(Debug, Default, Clone, Copy)]
@@ -237,9 +247,9 @@ impl<'a> Lexer<'a> {
                 value.push(self.advance());
             }
 
-            TokenType::Float(value.parse::<f64>().expect("Expected float"))
+            Literal::Float(value.parse::<f64>().expect("Expected float")).into()
         } else {
-            TokenType::Integer(value.parse::<i128>().expect("Expected integer"))
+            Literal::Integer(value.parse::<i64>().expect("Expected integer")).into()
         }
     }
 
@@ -272,7 +282,7 @@ impl<'a> Lexer<'a> {
             }
         }
 
-        TokenType::String(value)
+        Literal::String(value).into()
     }
 }
 
@@ -371,7 +381,7 @@ impl<'a> Iterator for Lexer<'a> {
             [':', ..] => self.advance_with(TokenType::Colon),
 
             // Literals
-            ['\'', c, '\''] => self.advance_by_with(3, TokenType::Character(c)),
+            ['\'', c, '\''] => self.advance_by_with(3, Literal::Character(c).into()),
             ['0'..='9', ..] => self.lex_number(),
             ['"', ..] => self.lex_string(),
 
@@ -382,6 +392,7 @@ impl<'a> Iterator for Lexer<'a> {
                 }
 
                 match value.as_str() {
+                    "const" => TokenType::Const,
                     "val" => TokenType::Val,
                     "var" => TokenType::Var,
                     "fn" => TokenType::Fn,
@@ -395,8 +406,8 @@ impl<'a> Iterator for Lexer<'a> {
                     "break" => TokenType::Break,
                     "continue" => TokenType::Continue,
                     "as" => TokenType::As,
-                    "true" => TokenType::Boolean(true),
-                    "false" => TokenType::Boolean(false),
+                    "true" => Literal::Boolean(true).into(),
+                    "false" => Literal::Boolean(false).into(),
                     _ => TokenType::Identifier(value),
                 }
             }
@@ -428,7 +439,7 @@ impl<'a> Iterator for Lexer<'a> {
 mod tests {
     use itertools::Itertools;
 
-    use super::{Lexer, TokenType};
+    use super::{Lexer, Literal, TokenType};
     use crate::lexer::LexerError;
 
     #[test]
@@ -504,13 +515,13 @@ mod tests {
             TokenType::Break,
             TokenType::Continue,
             TokenType::As,
-            TokenType::Boolean(true),
-            TokenType::Boolean(false),
+            Literal::Boolean(true).into(),
+            Literal::Boolean(false).into(),
         ]);
     }
 
     #[test]
-    fn lex_literals_a() {
+    fn lex_literals() {
         let source = "foo bar _foo __bar $0 $$1 \"foo\" \"bar\" \"baz\" \"\\\"\" \"\\n\" \"\\t\" \
                       'a' 'b' '\"' 93 3252 238 -382 -832 83 -25 52.9 83.7 12.4 35.2 3.3";
         let tokens = Lexer::new(source).map(|it| it.tt).collect_vec();
@@ -522,30 +533,30 @@ mod tests {
             TokenType::Identifier("__bar".to_owned()),
             TokenType::Identifier("$0".to_owned()),
             TokenType::Identifier("$$1".to_owned()),
-            TokenType::String("foo".to_owned()),
-            TokenType::String("bar".to_owned()),
-            TokenType::String("baz".to_owned()),
-            TokenType::String("\"".to_owned()),
-            TokenType::String("\n".to_owned()),
-            TokenType::String("\t".to_owned()),
-            TokenType::Character('a'),
-            TokenType::Character('b'),
-            TokenType::Character('"'),
-            TokenType::Integer(93),
-            TokenType::Integer(3252),
-            TokenType::Integer(238),
+            Literal::String("foo".to_owned()).into(),
+            Literal::String("bar".to_owned()).into(),
+            Literal::String("baz".to_owned()).into(),
+            Literal::String("\"".to_owned()).into(),
+            Literal::String("\n".to_owned()).into(),
+            Literal::String("\t".to_owned()).into(),
+            Literal::Character('a').into(),
+            Literal::Character('b').into(),
+            Literal::Character('"').into(),
+            Literal::Integer(93).into(),
+            Literal::Integer(3252).into(),
+            Literal::Integer(238).into(),
             TokenType::Minus,
-            TokenType::Integer(382),
+            Literal::Integer(382).into(),
             TokenType::Minus,
-            TokenType::Integer(832),
-            TokenType::Integer(83),
+            Literal::Integer(832).into(),
+            Literal::Integer(83).into(),
             TokenType::Minus,
-            TokenType::Integer(25),
-            TokenType::Float(52.9),
-            TokenType::Float(83.7),
-            TokenType::Float(12.4),
-            TokenType::Float(35.2),
-            TokenType::Float(3.3),
+            Literal::Integer(25).into(),
+            Literal::Float(52.9).into(),
+            Literal::Float(83.7).into(),
+            Literal::Float(12.4).into(),
+            Literal::Float(35.2).into(),
+            Literal::Float(3.3).into(),
         ]);
     }
 
diff --git a/sloth/src/main.rs b/sloth/src/main.rs
index 0e429b5..67f8f97 100644
--- a/sloth/src/main.rs
+++ b/sloth/src/main.rs
@@ -17,6 +17,7 @@ use std::{env, fs};
 use compiler::Compiler;
 use itertools::Itertools;
 use lexer::Lexer;
+use parser::graph::GraphBuilder;
 use parser::AstParser;
 
 fn main() {
@@ -35,9 +36,12 @@ fn main() {
     };
 
     let tokens = Lexer::new(&source).collect_vec();
-    let ast = AstParser::new(tokens).parse();
+    let ast = AstParser::parse(tokens).unwrap();
 
-    Compiler::compile(ast).unwrap();
+    let graph = GraphBuilder::generate(&ast).unwrap();
+    println!("{graph}");
+
+    // Compiler::compile(ast).unwrap();
 
     // let context = Context::create();
     // let compiler = Compiler::new(&context);
diff --git a/sloth/src/parser/ast.rs b/sloth/src/parser/ast.rs
index 543ea3a..2385443 100644
--- a/sloth/src/parser/ast.rs
+++ b/sloth/src/parser/ast.rs
@@ -1,52 +1,49 @@
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum BinaryOp {
-    Add,
-    Con,
-    Sub,
-    Mul,
-    Pow,
-    Div,
-    Mod,
+use crate::lexer::{self, TokenType};
+use crate::parser::ParsingError;
 
-    BWSftRight,
-    BWSftLeft,
-    BWAnd,
-    BWOr,
-    BWXor,
+// #[derive(PartialEq, Clone, Debug)]
+// pub struct Node {
+//     id: i32,
+//     kind: NodeKind,
+// }
+//
+// impl Node {
+//     pub fn new(id: i32, kind: NodeKind) -> Self {
+//         Self { id, kind }
+//     }
+// }
+//
+// #[derive(PartialEq, Clone, Debug)]
+// pub enum NodeKind {
+//     Expr(Expr),
+//     Stmt(Stmt),
+// }
+//
+// impl From<Expr> for NodeKind {
+//     fn from(value: Expr) -> Self {
+//         todo!()
+//     }
+// }
+//
+// impl From<Stmt> for NodeKind {}
 
-    Lt,
-    Gt,
-    LtEq,
-    GtEq,
-    EqEq,
-    NotEq,
-    LogAnd,
-    LogOr,
-    Range,
-}
-
-#[derive(Debug, Clone, Copy, PartialEq)]
-pub enum UnaryOp {
-    Not,
-    Neg,
-
-    BWComp,
+#[derive(PartialEq, Clone, Debug)]
+pub struct Expr {
+    pub id: i32,
+    pub kind: ExprKind,
 }
 
-#[derive(Debug, Clone, PartialEq)]
-pub enum Literal {
-    Integer(i128),
-    Float(f64),
-    Bool(bool),
-    Char(char),
-    String(String),
-    Regex(String),
-    List(Vec<Expr>),
+impl Expr {
+    pub fn new(id: i32, kind: ExprKind) -> Self {
+        Self { id, kind }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq)]
-pub enum Expr {
+pub enum ExprKind {
     Grouping(Box<Expr>),
+    Literal(Literal),
+    Identifier(String),
     BinaryOp {
         op: BinaryOp,
         lhs: Box<Expr>,
@@ -57,59 +54,153 @@ pub enum Expr {
         value: Box<Expr>,
     },
     Call {
-        ident: Box<Expr>,
+        callee: Box<Expr>,
         args: Vec<Expr>,
     },
-    Variable(String),
-    Literal(Literal),
-    Lambda, // TODO: Lambda
 }
 
 #[derive(PartialEq, Clone, Debug)]
-pub struct FuncArgs {
-    pub name: String,
-    pub typ: Option<String>,
+pub struct FunctionInput {
+    pub identifier: String,
+    pub typ: String,
 }
 
+// TODO: For loops
+// TODO: Values & Constants
 #[derive(PartialEq, Clone, Debug)]
 pub enum Stmt {
+    Block(Vec<Stmt>),
     ExprStmt(Expr),
-    DefineFunction {
-        ident: String,
-        args: Vec<FuncArgs>,
-        body: Vec<Stmt>,
-        return_type: Option<String>,
+    IfStmt {
+        condition: Expr,
+        if_then: Box<Stmt>,
+        else_then: Option<Box<Stmt>>,
     },
-    DefineVariable {
-        name: String,
-        value: Expr,
-        typ: Option<String>,
+    WhileStmt {
+        condition: Expr,
+        body: Box<Stmt>,
     },
-    DefineValue {
-        name: String,
+    DefineVariable {
+        identifier: String,
         value: Expr,
-        typ: Option<String>,
+        typ: String,
     },
     AssignVariable {
-        name: String,
+        identifier: String,
         value: Expr,
     },
-    If {
-        expr: Expr,
-        body: Vec<Stmt>,
-        else_if: Vec<(Expr, Stmt)>,
-        els: Option<Box<Stmt>>,
-    },
-    For {
-        name: String,
-        iter: Expr,
-        body: Vec<Stmt>,
-    },
-    While {
-        condition: Expr,
-        body: Vec<Stmt>,
-    },
-    Return {
-        value: Expr,
+    /// A function definition. Output is None when the function returns nothing
+    /// meaning void, otherwise it is the name of the type the function
+    /// returns.
+    DefineFunction {
+        identifier: String,
+        inputs: Vec<FunctionInput>,
+        output: Option<String>,
+        body: Box<Stmt>,
     },
+    Return(Expr),
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum Literal {
+    Integer(i64),
+    Float(f64),
+    Boolean(bool),
+    Character(char),
+    String(String),
+    Array(Vec<ExprKind>),
+}
+
+impl From<lexer::Literal> for Literal {
+    fn from(value: lexer::Literal) -> Self {
+        use lexer::Literal;
+
+        match value {
+            Literal::Integer(value) => Self::Integer(value),
+            Literal::Float(value) => Self::Float(value),
+            Literal::Boolean(value) => Self::Boolean(value),
+            Literal::Character(value) => Self::Character(value),
+            Literal::String(value) => Self::String(value),
+        }
+    }
+}
+
+impl From<Literal> for ExprKind {
+    fn from(value: Literal) -> Self {
+        Self::Literal(value)
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum BinaryOp {
+    Add,
+    Con,
+    Sub,
+    Mul,
+    Div,
+    Mod,
+
+    Lt,
+    Gt,
+    LtEq,
+    GtEq,
+    EqEq,
+    NotEq,
+
+    LogicalAnd,
+    LogicalOr,
+
+    Range,
+}
+
+impl TryFrom<TokenType> for BinaryOp {
+    type Error = ParsingError;
+
+    fn try_from(value: TokenType) -> Result<Self, Self::Error> {
+        let operation = match value {
+            TokenType::Plus => Self::Add,
+            TokenType::PlusPlus => Self::Con,
+            TokenType::Minus => Self::Sub,
+            TokenType::Star => Self::Mul,
+            TokenType::Slash => Self::Div,
+            TokenType::Perc => Self::Mod,
+
+            TokenType::Lt => Self::Lt,
+            TokenType::Gt => Self::Gt,
+            TokenType::LtEq => Self::LtEq,
+            TokenType::GtEq => Self::GtEq,
+            TokenType::EqEq => Self::EqEq,
+            TokenType::BangEq => Self::NotEq,
+
+            TokenType::AmpAmp => Self::LogicalAnd,
+            TokenType::PipePipe => Self::LogicalOr,
+
+            TokenType::DotDot => Self::Range,
+
+            _ => return Err(ParsingError::InvalidOp),
+        };
+
+        Ok(operation)
+    }
+}
+
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum UnaryOp {
+    Not,
+    Neg,
+}
+
+impl TryFrom<TokenType> for UnaryOp {
+    type Error = ParsingError;
+
+    fn try_from(value: TokenType) -> Result<Self, Self::Error> {
+        let operation = match value {
+            TokenType::Bang => Self::Not,
+            TokenType::Minus => Self::Neg,
+
+            _ => return Err(ParsingError::InvalidOp),
+        };
+
+        Ok(operation)
+    }
 }
diff --git a/sloth/src/parser/expr.rs b/sloth/src/parser/expr.rs
index 9e81f7f..aaf76a9 100644
--- a/sloth/src/parser/expr.rs
+++ b/sloth/src/parser/expr.rs
@@ -1,98 +1,47 @@
-use super::ast::{BinaryOp, Expr, Literal, UnaryOp};
+use super::ast::{Expr, Literal, UnaryOp};
 use super::AstParser;
 use crate::lexer::TokenType;
+use crate::parser::ast::{BinaryOp, ExprKind};
+use crate::parser::ParsingError;
 
-/// Implementation containing parsers internal components related to expressions
 impl<'a> AstParser<'a> {
-    // FIXME: Should probably avoid cloning token types
-
-    pub fn expression(&mut self) -> Expr {
+    pub(super) fn expression(&mut self) -> Result<Expr, ParsingError> {
         self.logical_or()
     }
 
-    fn unary(&mut self) -> Expr {
-        if !self.eof()
-            && matches!(
-                self.peek().tt,
-                TokenType::Bang | TokenType::Plus | TokenType::Minus
-            )
-        {
-            let operator = match self.advance().unwrap().tt.clone() {
-                TokenType::Bang => UnaryOp::Not,
-                TokenType::Tilde => UnaryOp::BWComp,
-                TokenType::Minus => UnaryOp::Neg,
-                _ => panic!(),
-            };
-
-            let rhs = self.unary();
-            return Expr::UnaryOp {
-                op: (operator),
-                value: (Box::new(rhs)),
-            };
-        }
-
-        self.call()
-    }
-
-    fn call(&mut self) -> Expr {
-        let mut expr = self.primary();
-
-        if self.advance_if_eq(&TokenType::OpeningParen) {
-            let mut arguments = Vec::<Expr>::new();
-
-            if self.peek().tt != TokenType::ClosingParen {
-                loop {
-                    arguments.push(self.expression());
-                    if !self.advance_if_eq(&TokenType::Comma) {
-                        break;
-                    }
-                }
-            }
+    fn unary(&mut self) -> Result<Expr, ParsingError> {
+        if matches!(self.peek().tt, TokenType::Bang | TokenType::Minus) {
+            let oeprator_tt = self.advance().unwrap().tt.clone();
+            let operator = UnaryOp::try_from(oeprator_tt)?;
 
-            self.consume(
-                TokenType::ClosingParen,
-                "Expected ')' to close off function call",
-            );
+            let value = self.unary()?;
 
-            // let Expr::Variable(_ident) = expr else { panic!("uh oh spaghettio"); };
+            let kind = ExprKind::UnaryOp {
+                op: operator,
+                value: Box::new(value),
+            };
 
-            expr = Expr::Call {
-                ident: (Box::new(expr)),
-                args: (arguments),
-            }
+            return Ok(Expr::new(self.reserve_id(), kind));
         }
 
-        expr
+        self.primary()
     }
 
-    fn primary(&mut self) -> Expr {
-        match self.advance().unwrap().tt.clone() {
-            TokenType::Integer(literal) => Expr::Literal(Literal::Integer(literal)),
-            TokenType::Float(literal) => Expr::Literal(Literal::Float(literal)),
-            TokenType::Boolean(literal) => Expr::Literal(Literal::Bool(literal)),
-            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),
+    fn primary(&mut self) -> Result<Expr, ParsingError> {
+        let kind = match self.advance().unwrap().tt.clone() {
+            TokenType::Literal(literal) => ExprKind::Literal(literal.into()),
+            TokenType::Identifier(identifier) => ExprKind::Identifier(identifier),
+
             TokenType::OpeningParen => {
-                let expr = self.expression();
-                self.consume(TokenType::ClosingParen, "Must end expression with ')'");
-                Expr::Grouping(Box::new(expr))
+                let expr = self.expression()?;
+                self.consume(TokenType::ClosingParen, "Must end grouping with ')'")?;
+                ExprKind::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);
+            _ => return Err(ParsingError::UnexpectedToken),
+        };
 
-                    self.advance_if_eq(&TokenType::Comma);
-                }
-                self.consume(TokenType::ClosingBracket, "Expected ']' at end of list");
-                Expr::Literal(Literal::List(expr))
-            }
-            _ => unimplemented!("{:?}", self.peek()),
-        }
+        Ok(Expr::new(self.reserve_id(), kind))
     }
 }
 
@@ -100,162 +49,78 @@ impl<'a> AstParser<'a> {
 // multiplication, exc.
 macro_rules! binary_expr {
     ($name:ident, $parent:ident, $pattern:pat) => {
-        fn $name(&mut self) -> Expr {
-            let mut expr = self.$parent();
+        fn $name(&mut self) -> Result<Expr, ParsingError> {
+            let mut expr = self.$parent()?;
 
             while !self.eof() && matches!(self.peek().tt, $pattern) {
-                let operator = match self.advance().unwrap().tt.clone() {
-                    TokenType::Plus => BinaryOp::Add,
-                    TokenType::PlusPlus => BinaryOp::Con,
-                    TokenType::Minus => BinaryOp::Sub,
-                    TokenType::Star => BinaryOp::Mul,
-                    TokenType::StarStar => BinaryOp::Pow,
-                    TokenType::Slash => BinaryOp::Div,
-                    TokenType::Perc => BinaryOp::Mod,
-                    TokenType::DotDot => BinaryOp::Range,
-
-                    TokenType::LtLt => BinaryOp::BWSftRight,
-                    TokenType::GtGt => BinaryOp::BWSftLeft,
-                    TokenType::Amp => BinaryOp::BWAnd,
-                    TokenType::Pipe => BinaryOp::BWOr,
-                    TokenType::Caret => BinaryOp::BWXor,
-
-                    TokenType::Lt => BinaryOp::Lt,
-                    TokenType::Gt => BinaryOp::Gt,
-                    TokenType::LtEq => BinaryOp::LtEq,
-                    TokenType::GtEq => BinaryOp::GtEq,
-                    TokenType::EqEq => BinaryOp::EqEq,
-                    TokenType::BangEq => BinaryOp::NotEq,
-                    TokenType::AmpAmp => BinaryOp::LogAnd,
-                    TokenType::PipePipe => BinaryOp::LogOr,
-                    _ => panic!("uh oh spagghetio"),
+                let operator_tt = self.advance().unwrap().tt.clone();
+                let operator = BinaryOp::try_from(operator_tt)?;
+
+                let rhs = self.$parent()?;
+                let kind = ExprKind::BinaryOp {
+                    op: operator,
+                    lhs: Box::new(expr),
+                    rhs: Box::new(rhs),
                 };
 
-                let rhs = self.$parent();
-                expr = Expr::BinaryOp {
-                    op: (operator),
-                    lhs: (Box::new(expr)),
-                    rhs: (Box::new(rhs)),
-                }
+                expr = Expr::new(self.reserve_id(), kind);
             }
 
-            expr
+            Ok(expr)
         }
     };
 }
 
 #[rustfmt::skip]
 #[allow(unused_parens)]
-impl<'a> AstParser<'a> {
+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     , 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));
-    binary_expr!(additive        , multiplicative  , (TokenType::Plus   | TokenType::Minus));
-    binary_expr!(multiplicative  , unary           , (TokenType::Star   | TokenType::Slash | TokenType::Perc));
+    binary_expr!(logical_or      , logical_and   , (TokenType::PipePipe));
+    binary_expr!(logical_and     , range         , (TokenType::AmpAmp));
+    binary_expr!(range           , equality      , (TokenType::DotDot));
+    binary_expr!(equality        , comparison    , (TokenType::BangEq | TokenType::EqEq));
+    binary_expr!(comparison      , additive      , (TokenType::Lt     | TokenType::Gt    | TokenType::LtEq | TokenType::GtEq));
+    binary_expr!(additive        , multiplicative, (TokenType::Plus   | TokenType::Minus));
+    binary_expr!(multiplicative  , unary         , (TokenType::Star   | TokenType::Slash | TokenType::Perc));
 }
 
 #[cfg(test)]
 mod tests {
     use itertools::Itertools;
 
-    use super::{AstParser, BinaryOp, Expr, Literal};
+    use super::{AstParser, BinaryOp, ExprKind, Literal};
     use crate::lexer::Lexer;
-    use crate::parser::ast::UnaryOp;
-
-    #[test]
-    fn basic_expression_a() {
-        let lexer = Lexer::new("3 + 5 * 4");
-        let tokens = lexer.collect_vec();
-
-        let expected_ast = Expr::BinaryOp {
-            op: BinaryOp::Add,
-            lhs: Box::new(Expr::Literal(Literal::Integer(3))),
-            rhs: Box::new(Expr::BinaryOp {
-                op: BinaryOp::Mul,
-                lhs: Box::new(Expr::Literal(Literal::Integer(5))),
-                rhs: Box::new(Expr::Literal(Literal::Integer(4))),
-            }),
-        };
-
-        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);
-    }
+    use crate::parser::ast::Expr;
 
     #[test]
-    fn basic_expression_b() {
-        let lexer = Lexer::new("17 - (-5 + 5) / 6");
+    fn basic_expression() {
+        let lexer = Lexer::new("3 + 5 * 4 - 9 / 3");
         let tokens = lexer.collect_vec();
 
-        let expected_ast = Expr::BinaryOp {
-            op: BinaryOp::Sub,
-            lhs: Box::new(Expr::Literal(Literal::Integer(17))),
-            rhs: Box::new(Expr::BinaryOp {
-                op: BinaryOp::Div,
-                lhs: Box::new(Expr::Grouping(Box::new(Expr::BinaryOp {
-                    op: BinaryOp::Add,
-                    lhs: Box::new(Expr::UnaryOp {
-                        op: UnaryOp::Neg,
-                        value: Box::new(Expr::Literal(Literal::Integer(5))),
-                    }),
-                    rhs: Box::new(Expr::Literal(Literal::Integer(5))),
-                }))),
-                rhs: Box::new(Expr::Literal(Literal::Integer(6))),
-            }),
-        };
-
-        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_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 expected_ast = Ok(ExprKind::BinaryOp {
+        //     op: BinaryOp::Sub,
+        //     lhs: Box::new(ExprKind::BinaryOp {
+        //         op: BinaryOp::Add,
+        //         lhs: Box::new(ExprKind::Literal(Literal::Integer(3))),
+        //         rhs: Box::new(ExprKind::BinaryOp {
+        //             op: BinaryOp::Mul,
+        //             lhs: Box::new(ExprKind::Literal(Literal::Integer(5))),
+        //             rhs: Box::new(ExprKind::Literal(Literal::Integer(4))),
+        //         }),
+        //     }),
+        //     rhs: Box::new(ExprKind::BinaryOp {
+        //         op: BinaryOp::Div,
+        //         lhs: Box::new(ExprKind::Literal(Literal::Integer(9))),
+        //         rhs: Box::new(ExprKind::Literal(Literal::Integer(3))),
+        //     }),
+        // });
 
         let mut parser = AstParser::new(tokens);
         let generated_ast = parser.expression();
 
-        println!("Expected AST:\n{expected_ast:#?}\n\n");
+        // println!("Expected AST:\n{expected_ast:#?}\n\n");
         println!("Generated AST:\n{generated_ast:#?}\n\n");
 
-        assert_eq!(expected_ast, generated_ast);
+        assert_eq!(Ok(Expr::new(0, Literal::Integer(0).into())), generated_ast);
     }
 }
diff --git a/sloth/src/parser/graph.rs b/sloth/src/parser/graph.rs
new file mode 100644
index 0000000..b35be05
--- /dev/null
+++ b/sloth/src/parser/graph.rs
@@ -0,0 +1,160 @@
+use std::fmt::{Error, Write};
+
+use super::ast::{ExprKind, Stmt};
+
+pub struct GraphBuilder {
+    i: i32,
+    graph: String,
+}
+
+impl GraphBuilder {
+    pub fn generate(ast: &[Stmt]) -> Result<String, Error> {
+        let mut this = Self {
+            i: 0,
+            graph: String::new(),
+        };
+
+        this.graph.push_str("digraph {\n");
+        for stmt in ast.iter() {
+            this.traverse_stmt0(stmt)?;
+        }
+        this.i = 0;
+        for stmt in ast.iter() {
+            this.traverse_stmt(stmt)?;
+        }
+        this.graph.push('}');
+
+        Ok(this.graph)
+    }
+
+    fn traverse_stmt0(&mut self, stmt: &Stmt) -> Result<(), Error> {
+        self.i += 1;
+
+        match stmt {
+            Stmt::Block(body) => {
+                writeln!(&mut self.graph, "N{} [shape=box label=\"Block\"];", self.i)?;
+                for stmt in body.iter() {
+                    self.traverse_stmt0(stmt)?;
+                }
+            }
+            Stmt::ExprStmt(expr) => {
+                writeln!(
+                    &mut self.graph,
+                    "N{} [shape=box label=\"ExprStmt\"];",
+                    self.i
+                )?;
+                // self.traverse_expr0(expr);
+            }
+            Stmt::IfStmt {
+                condition,
+                if_then,
+                else_then,
+            } => {
+                writeln!(&mut self.graph, "N{} [shape=box label=\"IfStmt\"];", self.i)?;
+                // self.traverse_expr0(condition);
+                self.traverse_stmt0(if_then)?;
+                if let Some(else_then) = else_then {
+                    self.traverse_stmt0(else_then)?;
+                }
+            }
+            Stmt::WhileStmt { condition, body } => {
+                writeln!(
+                    &mut self.graph,
+                    "N{} [shape=box label=\"WhileStmt\"];",
+                    self.i
+                )?;
+                // self.traverse_expr0(condition);
+                self.traverse_stmt0(body)?;
+            }
+            Stmt::DefineVariable {
+                identifier,
+                value,
+                typ,
+            } => {
+                writeln!(
+                    &mut self.graph,
+                    "N{} [shape=box label=\"DefineVariable\\n\\nIdentifier={}\\nType={}\"];",
+                    self.i, identifier, typ
+                )?;
+                // self.traverse_expr0(value);
+            }
+            Stmt::AssignVariable { identifier, value } => {
+                writeln!(
+                    &mut self.graph,
+                    "N{} [shape=box label=\"AssignVariable\\n\\nIdentifier={}\"];",
+                    self.i, identifier
+                )?;
+                // self.traverse_expr0(value);
+            }
+            Stmt::DefineFunction {
+                identifier,
+                inputs,
+                output,
+                body,
+            } => {
+                writeln!(
+                    &mut self.graph,
+                    "N{} [shape=box \
+                     label=\"DefineFunction\\n\\nIdentifier={}\\nInputs={}\\nOutput={}\"];",
+                    self.i,
+                    identifier,
+                    inputs.len(),
+                    output.is_some()
+                )?;
+                self.traverse_stmt0(body)?;
+            }
+            Stmt::Return(expr) => {
+                writeln!(&mut self.graph, "N{} [shape=box label=\"Return\"];", self.i)?;
+                // self.traverse_expr0(expr);
+            }
+        }
+
+        Ok(())
+    }
+
+    fn traverse_expr0(&mut self, expr: &ExprKind) {
+        self.i += 1;
+        // match expr {
+        //     Expr::Grouping(_) => todo!(),
+        //     Expr::Literal(_) => todo!(),
+        //     Expr::Identifier(_) => todo!(),
+        //     Expr::BinaryOp { op, lhs, rhs } => todo!(),
+        //     Expr::UnaryOp { op, value } => todo!(),
+        //     Expr::Call { callee, args } => todo!(),
+        // }
+    }
+
+    fn traverse_stmt(&mut self, stmt: &Stmt) -> Result<(), Error> {
+        self.i += 1;
+
+        match stmt {
+            Stmt::Block(_) => todo!(),
+            Stmt::ExprStmt(_) => todo!(),
+            Stmt::IfStmt {
+                condition,
+                if_then,
+                else_then,
+            } => todo!(),
+            Stmt::WhileStmt { condition, body } => todo!(),
+            Stmt::DefineVariable {
+                identifier,
+                value,
+                typ,
+            } => todo!(),
+            Stmt::AssignVariable { identifier, value } => todo!(),
+            Stmt::DefineFunction {
+                identifier,
+                inputs,
+                output,
+                body,
+            } => todo!(),
+            Stmt::Return(_) => todo!(),
+        }
+
+        Ok(())
+    }
+
+    fn traverse_expr(&mut self, expr: &ExprKind) {
+        //
+    }
+}
diff --git a/sloth/src/parser/mod.rs b/sloth/src/parser/mod.rs
index 9d77acc..bc47ddd 100644
--- a/sloth/src/parser/mod.rs
+++ b/sloth/src/parser/mod.rs
@@ -1,23 +1,57 @@
 pub mod ast;
 pub mod expr;
+pub mod graph;
 pub mod stmt;
 
+use self::ast::{Literal, Stmt};
 use crate::lexer::{Token, TokenType};
+
+#[derive(thiserror::Error, Debug, PartialEq)]
+pub enum ParsingError {
+    #[error("Invalid operation")]
+    InvalidOp,
+    #[error("Unexpected token")]
+    UnexpectedToken,
+}
+
 #[derive(Debug)]
 pub struct AstParser<'a> {
     tokens: Vec<Token<'a>>,
     index: usize,
+    id: i32,
+}
+
+impl<'a> AstParser<'a> {
+    pub fn parse(tokens: Vec<Token<'a>>) -> Result<Vec<Stmt>, ParsingError> {
+        let mut parser = Self::new(tokens);
+
+        let mut statements = Vec::new();
+        while !parser.eof() {
+            statements.push(parser.statement()?);
+        }
+
+        Ok(statements)
+    }
 }
 
 /// 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 }
+        Self {
+            tokens,
+            index: 0,
+            id: 0,
+        }
     }
+
     pub fn peek(&self) -> &Token {
         &self.tokens[self.index]
     }
 
+    pub fn peek2(&self) -> &Token {
+        &self.tokens[self.index + 1]
+    }
+
     pub fn advance(&mut self) -> Option<&Token> {
         if self.eof() {
             return None;
@@ -44,11 +78,35 @@ impl<'a> AstParser<'a> {
         self.advance_if(|it| it.tt == *next)
     }
 
-    pub fn consume(&mut self, next: TokenType, error: &str) {
+    pub fn consume(&mut self, next: TokenType, error: &str) -> Result<&Token, ParsingError> {
         if std::mem::discriminant(&self.peek().tt) != std::mem::discriminant(&next) {
-            panic!("{error} at index {:?}", self.index);
+            println!("{error} at index {:?}", self.index);
+            return Err(ParsingError::UnexpectedToken);
         }
-        self.advance();
+
+        Ok(self.advance().unwrap())
+    }
+
+    pub fn consume_literal(&mut self) -> Result<Literal, ParsingError> {
+        let Some(TokenType::Literal(literal)) = self.advance().map(|it| it.tt.clone()) else {
+            return Err(ParsingError::UnexpectedToken);
+        };
+
+        Ok(literal.into())
+    }
+
+    pub fn consume_identifier(&mut self) -> Result<String, ParsingError> {
+        let Some(TokenType::Identifier(identifier)) = self.advance().map(|it| it.tt.clone()) else {
+            return Err(ParsingError::UnexpectedToken);
+        };
+
+        Ok(identifier)
+    }
+
+    pub fn reserve_id(&mut self) -> i32 {
+        let id = self.id;
+        self.id += 1;
+        id
     }
 
     pub fn eof(&self) -> bool {
diff --git a/sloth/src/parser/stmt.rs b/sloth/src/parser/stmt.rs
index 1a961b1..aaa65ff 100644
--- a/sloth/src/parser/stmt.rs
+++ b/sloth/src/parser/stmt.rs
@@ -1,317 +1,161 @@
-use super::ast::{Expr, FuncArgs, Stmt};
-use super::AstParser;
+use super::ast::{FunctionInput, Stmt};
+use super::{AstParser, ParsingError};
 use crate::lexer::TokenType;
 
 impl<'a> AstParser<'a> {
-    pub fn parse(&mut self) -> Vec<Stmt> {
-        let mut statements = Vec::new();
-
-        while !self.eof() {
-            statements.push(self.statement());
+    pub(super) fn statement(&mut self) -> Result<Stmt, ParsingError> {
+        match self.peek().tt {
+            TokenType::OpeningBrace => self.block(),
+
+            TokenType::If => self.if_stmt(),
+            TokenType::While => self.while_stmt(),
+            TokenType::Var => self.define_variable(),
+            TokenType::Fn => self.define_function(),
+            TokenType::Return => self.return_stmt(),
+
+            _ if self.peek2().tt == TokenType::Eq => self.assign_variable(),
+            _ => self.expression_stmt(),
         }
-
-        statements
     }
 
-    fn statement(&mut self) -> Stmt {
-        if self.advance_if_eq(&TokenType::Var) {
-            return self.var_statement();
-        }
-
-        if self.advance_if_eq(&TokenType::Val) {
-            return self.val_statement();
-        }
-
-        if self.advance_if_eq(&TokenType::If) {
-            return self.if_statement();
-        }
-
-        if self.advance_if_eq(&TokenType::For) {
-            return self.for_statement();
-        }
-
-        if self.advance_if_eq(&TokenType::While) {
-            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();
+    fn if_stmt(&mut self) -> Result<Stmt, ParsingError> {
+        // Consume the if token
+        self.consume(TokenType::If, "Expected if")?;
+
+        // Get the condition and if_then of the if statement
+        let condition = self.expression()?;
+        let if_then = self.block()?;
+
+        // Check if there is an else
+        let mut else_then = None;
+        if self.advance_if_eq(&TokenType::Else) {
+            if self.peek().tt == TokenType::If {
+                else_then = Some(self.if_stmt()?);
+            } else {
+                else_then = Some(self.block()?);
+            }
         }
 
-        self.mut_statement()
-
-        // If we couldn't parse a statement return an expression statement
-        // self.expression_statement()
+        Ok(Stmt::IfStmt {
+            condition,
+            if_then: if_then.into(),
+            else_then: else_then.map(|it| it.into()),
+        })
     }
 
-    fn mut_statement(&mut self) -> Stmt {
-        let TokenType::Identifier(ident) = self.peek().tt.clone() else {
-            panic!("Identifier error  {:?}", self.peek());
-        };
-
-        self.advance();
-        let next = self.advance().unwrap().tt.clone();
-        if next == TokenType::Eq {
-            let value = self.expression();
-            self.consume(TokenType::SemiColon, "No semi colon for me i guess");
-            return Stmt::AssignVariable {
-                name: (ident),
-                value: (value),
-            };
-        } else if next == TokenType::OpeningParen {
-            let mut arguments = Vec::<Expr>::new();
-
-            if self.peek().tt != TokenType::ClosingParen {
-                loop {
-                    arguments.push(self.expression());
-                    if !self.advance_if_eq(&TokenType::Comma) {
-                        break;
-                    }
-                }
-            }
+    fn while_stmt(&mut self) -> Result<Stmt, ParsingError> {
+        // Consume the while token
+        self.consume(TokenType::While, "Expected while")?;
 
-            self.consume(
-                TokenType::ClosingParen,
-                "Expected ')' to close off function call",
-            );
+        let condition = self.expression()?;
+        let body = self.block()?;
 
-            self.consume(TokenType::SemiColon, "No semi colon for me i guess");
-            return Stmt::ExprStmt(Expr::Call {
-                ident: Box::new(Expr::Variable(ident)),
-                args: (arguments),
-            });
-        }
-        self.expression_statement()
+        Ok(Stmt::WhileStmt {
+            condition,
+            body: body.into(),
+        })
     }
 
-    fn var_statement(&mut self) -> Stmt {
-        let TokenType::Identifier(ident) = self.peek().tt.clone() else {
-            panic!("Identifier expected after 'var', not {:?}", self.peek());
-        };
-
-        self.advance();
-
-        let mut typ: Option<String> = None;
-        if self.peek().tt.clone() == TokenType::Colon {
-            self.consume(TokenType::Colon, "How did you even get this error?");
-            let TokenType::Identifier(name) = self.peek().tt.clone() else {
-                panic!("Type expected after identifier, not {:?}", self.peek());
-            };
-            self.advance();
-            typ = Some(name);
-        }
+    // TODO: Make variable types optional
+    fn define_variable(&mut self) -> Result<Stmt, ParsingError> {
+        // Consume the var token
+        self.consume(TokenType::Var, "Expected var")?;
 
-        self.consume(TokenType::Eq, "Expected '=' after identifier at ");
+        // Get the identifier and type
+        let identifier = self.consume_identifier()?;
+        self.consume(TokenType::Colon, "Expected ':'")?;
+        let typ = self.consume_identifier()?;
 
-        let value = self.expression();
+        // Get the default value
+        self.consume(TokenType::Eq, "Expected '='")?;
+        let value = self.expression()?;
 
-        self.consume(TokenType::SemiColon, "Expected ';' at end of statement");
+        self.consume(TokenType::SemiColon, "Expected ';' at end of statement")?;
 
-        Stmt::DefineVariable {
-            name: (ident),
-            value: (value),
-            typ: (typ),
-        }
+        Ok(Stmt::DefineVariable {
+            identifier,
+            value,
+            typ,
+        })
     }
 
-    fn val_statement(&mut self) -> Stmt {
-        let TokenType::Identifier(ident) = self.peek().tt.clone() else {
-            panic!("Identifier expected after 'val'");
-        };
-
-        self.advance(); // Advancing from the identifier
+    // TODO: Make argument types optional
+    fn define_function(&mut self) -> Result<Stmt, ParsingError> {
+        // Consume the fn token
+        self.consume(TokenType::Fn, "Expected fn")?;
 
-        let mut typ: Option<String> = None;
-        if self.peek().tt.clone() == TokenType::Colon {
-            self.consume(TokenType::Colon, "How did you even get this error?");
-            let TokenType::Identifier(name) = self.peek().tt.clone() else {
-                panic!("Type expected after identifier, not {:?}", self.peek());
-            };
-            self.advance();
-            typ = Some(name);
-        }
+        let identifier = self.consume_identifier()?;
 
-        self.consume(TokenType::Eq, "Expected '=' after identifier");
+        // Get the function inputs
+        self.consume(TokenType::OpeningParen, "Expected '('")?;
 
-        let value = self.expression();
+        let mut inputs = Vec::new();
+        while matches!(self.peek().tt, TokenType::Identifier(_)) {
+            let input_identifier = self.consume_identifier()?;
+            self.consume(TokenType::Colon, "Expected ':'")?;
+            let input_type = self.consume_identifier()?;
 
-        self.consume(TokenType::SemiColon, "Expected ';' at end of statement");
-
-        Stmt::DefineValue {
-            name: (ident),
-            value: (value),
-            typ: (typ),
+            inputs.push(FunctionInput {
+                identifier: input_identifier,
+                typ: input_type,
+            });
         }
-    }
 
-    fn if_statement(&mut self) -> Stmt {
-        let condition = self.expression();
+        self.consume(TokenType::ClosingParen, "Expected ')'")?;
 
-        self.consume(
-            TokenType::OpeningBrace,
-            "Expected '{' at beggining of block",
-        );
-        let mut body = Vec::new();
-        while !self.eof() && self.peek().tt != TokenType::ClosingBrace {
-            body.push(self.statement());
-        }
-        self.advance();
-        Stmt::If {
-            expr: (condition),
-            body: (body),
-            else_if: (Vec::new()),
-            els: (None),
-        } // 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");
+        // Get the function output
+        let output = if matches!(self.peek().tt, TokenType::Identifier(_)) {
+            Some(self.consume_identifier()?)
+        } else {
+            None
         };
 
-        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 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());
-        }
-        self.advance();
-
-        Stmt::For {
-            name: (binding),
-            iter: (expr),
-            body: (body),
-        }
-    } // TODO: Fix this garbage
-
-    fn while_statement(&mut self) -> Stmt {
-        let condition = self.expression();
-
-        self.consume(
-            TokenType::OpeningBrace,
-            "Expected '{' at beggining of block",
-        );
-        let mut body = Vec::new();
-        while !self.eof() && self.peek().tt != TokenType::ClosingBrace {
-            println!("{:?}", self.peek().tt);
-            body.push(self.statement());
-        }
-        self.consume(
-            TokenType::ClosingBrace,
-            "Expected '}' after block on while loop",
-        );
+        // Get the function body
+        let body = self.block()?;
 
-        self.advance();
-        Stmt::While { condition, body }
+        Ok(Stmt::DefineFunction {
+            identifier,
+            inputs,
+            output,
+            body: body.into(),
+        })
     }
 
-    fn expression_statement(&mut self) -> Stmt {
-        let expr = self.expression();
-
-        // FIXME: Move assignment handling
-        // 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 expr statement",
-        );
-        Stmt::ExprStmt(expr)
+    fn return_stmt(&mut self) -> Result<Stmt, ParsingError> {
+        self.consume(TokenType::Return, "Expected return")?;
+        let value = self.expression()?;
+        self.consume(TokenType::SemiColon, "Expected ';' at end of statement")?;
+        Ok(Stmt::Return(value))
     }
 
-    fn function_statement(&mut self) -> Stmt {
-        let TokenType::Identifier(ident) = self.advance().unwrap().tt.clone() else {
-            panic!("Identifier expected after 'fn'");
-        };
+    fn assign_variable(&mut self) -> Result<Stmt, ParsingError> {
+        let identifier = self.consume_identifier()?;
+        self.consume(TokenType::Eq, "Expected '='")?;
+        let value = self.expression()?;
+        self.consume(TokenType::SemiColon, "Expected ';' at end of statement")?;
+        Ok(Stmt::AssignVariable { identifier, value })
+    }
 
-        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!("parameter expected after '('");
-            };
-
-            let mut typ: Option<String> = None;
-
-            if self.peek().tt.clone() == TokenType::Colon {
-                self.consume(TokenType::Colon, "How did you even get this error?");
-                let TokenType::Identifier(name) = self.peek().tt.clone() else {
-                    panic!("Type expected after ':', not {:?}", self.peek());
-                };
-                self.advance();
-                typ = Some(name);
-            }
+    fn expression_stmt(&mut self) -> Result<Stmt, ParsingError> {
+        let expr = self.expression()?;
+        self.consume(TokenType::SemiColon, "Expected ';' at end of statement")?;
+        Ok(Stmt::ExprStmt(expr))
+    }
 
-            self.advance_if_eq(&TokenType::Comma);
+    fn block(&mut self) -> Result<Stmt, ParsingError> {
+        // Consume the opening brace
+        self.consume(TokenType::OpeningBrace, "Expected '{'")?;
 
-            let arg = FuncArgs {
-                name: (name),
-                typ: (typ),
-            };
-            args.push(arg);
-        }
-        self.advance();
-        let mut typ: Option<String> = None;
-        if self.peek().tt.clone() == TokenType::Arrow {
-            self.advance();
-            let TokenType::Identifier(name) = self.peek().tt.clone() else {
-                panic!("Type expected after ':', not {:?}", self.peek());
-            };
-            typ = Some(name);
-            self.advance();
-        }
-        self.consume(TokenType::OpeningBrace, "Expected '{' after parameters");
+        // Get the body of the block
         let mut body = Vec::new();
         while !self.eof() && self.peek().tt != TokenType::ClosingBrace {
-            body.push(self.statement());
+            body.push(self.statement()?);
         }
-        self.consume(TokenType::ClosingBrace, "Expected '}' after body");
 
-        Stmt::DefineFunction {
-            ident: (ident),
-            args: (args),
-            body: (body),
-            return_type: (typ),
-        }
-    }
+        // Consume the closing brace
+        self.consume(TokenType::ClosingBrace, "Expected '}'")?;
 
-    fn return_statement(&mut self) -> Stmt {
-        let expr = self.expression();
-        self.consume(TokenType::SemiColon, "Expected ';' after return statement");
-        Stmt::Return { value: (expr) }
+        Ok(Stmt::Block(body))
     }
 }
 
@@ -321,326 +165,145 @@ mod tests {
 
     use super::{AstParser, Stmt};
     use crate::lexer::Lexer;
-    use crate::parser::ast::{BinaryOp, Expr, FuncArgs, Literal, UnaryOp};
-
-    #[test]
-    fn basic_statement_a() {
-        let lexer = Lexer::new("var test_a: int = 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: Some("int".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);
-    }
-
-    #[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 (a + b * c);\n}",
-        );
-        let tokens = lexer.collect_vec();
-        println!("{tokens:?}");
-
-        let expected_ast = Stmt::DefineFunction {
-            ident: ("test_c".to_string()),
-            args: (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::Grouping(Box::new(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\");\nprintln(5 + 7/-3);\n}",
-        );
-        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::Variable("print".to_string())),
-                    args: (vec![Expr::Literal(Literal::String("Hello World".to_string()))]),
-                }),
-                Stmt::ExprStmt(Expr::Call {
-                    ident: Box::new(Expr::Variable("println".to_string())),
-                    args: (vec![Expr::BinaryOp {
-                        op: (BinaryOp::Add),
-                        lhs: (Box::new(Expr::Literal(Literal::Integer(5)))),
-                        rhs: (Box::new(Expr::BinaryOp {
-                            op: (BinaryOp::Div),
-                            lhs: (Box::new(Expr::Literal(Literal::Integer(7)))),
-                            rhs: (Box::new(Expr::UnaryOp {
-                                op: (UnaryOp::Neg),
-                                value: (Box::new(Expr::Literal(Literal::Integer(3)))),
-                            })),
-                        })),
-                    }]),
-                }),
-            ]),
-        };
-
-        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_e() {
-        let lexer = Lexer::new(
-            "\
-            if a+5 > 10 {\nprint(a);\n}\nif a+5 < 10 {\nprintln(10);\n}\nif a+5 == 10 \
-             {\nprint(toString(10));\na = true;\n}",
-        );
-        let tokens = lexer.collect_vec();
-        // println!("{tokens:?}");
-
-        let expected_ast = vec![
-            Stmt::If {
-                expr: (Expr::BinaryOp {
-                    op: (BinaryOp::Gt),
-                    lhs: (Box::new(Expr::BinaryOp {
-                        op: (BinaryOp::Add),
-                        lhs: (Box::new(Expr::Variable("a".to_string()))),
-                        rhs: (Box::new(Expr::Literal(Literal::Integer(5)))),
-                    })),
-                    rhs: (Box::new(Expr::Literal(Literal::Integer(10)))),
-                }),
-                body: (vec![Stmt::ExprStmt(Expr::Call {
-                    ident: (Box::new(Expr::Variable("print".to_string()))),
-                    args: (vec![Expr::Variable("a".to_string())]),
-                })]),
-                else_if: (Vec::new()),
-                els: (None),
-            },
-            Stmt::If {
-                expr: (Expr::BinaryOp {
-                    op: (BinaryOp::Lt),
-                    lhs: (Box::new(Expr::BinaryOp {
-                        op: (BinaryOp::Add),
-                        lhs: (Box::new(Expr::Variable("a".to_string()))),
-                        rhs: (Box::new(Expr::Literal(Literal::Integer(5)))),
-                    })),
-                    rhs: (Box::new(Expr::Literal(Literal::Integer(10)))),
-                }),
-                body: (vec![Stmt::ExprStmt(Expr::Call {
-                    ident: (Box::new(Expr::Variable("println".to_string()))),
-                    args: (vec![Expr::Literal(Literal::Integer(10))]),
-                })]),
-                else_if: (Vec::new()),
-                els: (None),
-            },
-            Stmt::If {
-                expr: (Expr::BinaryOp {
-                    op: (BinaryOp::EqEq),
-                    lhs: (Box::new(Expr::BinaryOp {
-                        op: (BinaryOp::Add),
-                        lhs: (Box::new(Expr::Variable("a".to_string()))),
-                        rhs: (Box::new(Expr::Literal(Literal::Integer(5)))),
-                    })),
-                    rhs: (Box::new(Expr::Literal(Literal::Integer(10)))),
-                }),
-                body: (vec![
-                    Stmt::ExprStmt(Expr::Call {
-                        ident: (Box::new(Expr::Variable("print".to_string()))),
-                        // ident: (Box::new(Expr::Literal(Literal::String("print".to_string())))),
-                        args: (vec![Expr::Call {
-                            ident: (Box::new(Expr::Variable("toString".to_string()))),
-                            // ident: Box::new(Expr::Literal(Literal::String("toString".
-                            // to_string()))),
-                            args: vec![Expr::Literal(Literal::Integer(10))],
-                        }]),
-                    }),
-                    Stmt::AssignVariable {
-                        name: ("a".to_string()),
-                        value: (Expr::Literal(Literal::Bool(true))),
-                    },
-                ]),
-
-                else_if: (Vec::new()),
-                els: (None),
-            },
-        ];
-
-        let mut parser = AstParser::new(tokens);
-        let generated_ast = parser.parse();
-
-        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_f() {
-        let lexer = Lexer::new("test_a = 5 + 3;");
-        let tokens = lexer.collect_vec();
-
-        let expected_ast = Stmt::AssignVariable {
-            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)))),
-            }),
-        };
-
-        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_g() {
-        let lexer = Lexer::new(
-            "\
-        fn times_two(x: int) -> int {\nval y: int = x*2;\nreturn y;\n}",
-        );
-        let tokens = lexer.collect_vec();
-
-        let expected_ast = Stmt::DefineFunction {
-            ident: ("times_two".to_string()),
-            args: (vec![FuncArgs {
-                name: ("x".to_string()),
-                typ: (Some("int".to_string())),
-            }]),
-            body: (vec![
-                Stmt::DefineValue {
-                    name: "y".to_string(),
-                    value: (Expr::BinaryOp {
-                        op: (BinaryOp::Mul),
-                        lhs: (Box::new(Expr::Variable("x".to_string()))),
-                        rhs: (Box::new(Expr::Literal(Literal::Integer(2)))),
-                    }),
-                    typ: Some("int".to_string()),
-                },
-                Stmt::Return {
-                    value: (Expr::Variable("y".to_string())),
-                },
-            ]),
-
-            return_type: Some("int".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);
-    }
-
-    #[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::Variable("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);
-    }
+    use crate::parser::ast::{BinaryOp, ExprKind, FunctionInput, Literal};
+
+    // #[test]
+    // fn standalone_blocks() {
+    //     let tokens = Lexer::new("{{{ 0; }}}").collect_vec();
+    //
+    //     let expected_ast =
+    // Ok(Stmt::Block(vec![Stmt::Block(vec![Stmt::Block(vec![
+    //         Stmt::ExprStmt(Literal::Integer(0).into()),
+    //     ])])]));
+    //
+    //     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_variable_definition() {
+    //     let tokens = Lexer::new("var foo: Int = 5 + 3;").collect_vec();
+    //
+    //     let expected_ast = Ok(Stmt::DefineVariable {
+    //         identifier: "foo".to_string(),
+    //         value: ExprKind::BinaryOp {
+    //             op: BinaryOp::Add,
+    //             lhs: Box::new(ExprKind::Literal(Literal::Integer(5))),
+    //             rhs: Box::new(ExprKind::Literal(Literal::Integer(3))),
+    //         },
+    //         typ: "Int".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);
+    // }
+    //
+    // #[test]
+    // fn basic_function() {
+    //     let tokens = Lexer::new(
+    //         r#"
+    //         fn foo(bar: Int) Int {
+    //             var baz: Int = bar + 1;
+    //             baz = baz + 1;
+    //             return baz;
+    //         }
+    //     "#,
+    //     )
+    //     .collect_vec();
+    //
+    //     let expected_ast = Ok(Stmt::DefineFunction {
+    //         identifier: "foo".to_owned(),
+    //         inputs: vec![FunctionInput {
+    //             identifier: "bar".to_owned(),
+    //             typ: "Int".to_owned(),
+    //         }],
+    //         output: Some("Int".to_owned()),
+    //         body: Box::new(Stmt::Block(vec![
+    //             Stmt::DefineVariable {
+    //                 identifier: "baz".to_owned(),
+    //                 value: ExprKind::BinaryOp {
+    //                     op: BinaryOp::Add,
+    //                     lhs:
+    // Box::new(ExprKind::Identifier("bar".to_owned())),                    
+    // rhs: Box::new(Literal::Integer(1).into()),                 },
+    //                 typ: "Int".to_owned(),
+    //             },
+    //             Stmt::AssignVariable {
+    //                 identifier: "baz".to_owned(),
+    //                 value: ExprKind::BinaryOp {
+    //                     op: BinaryOp::Add,
+    //                     lhs:
+    // Box::new(ExprKind::Identifier("baz".to_owned())),                    
+    // rhs: Box::new(Literal::Integer(1).into()),                 },
+    //             },
+    //             Stmt::Return(ExprKind::Identifier("baz".to_owned())),
+    //         ])),
+    //     });
+    //
+    //     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_conditional() {
+    //     let tokens = Lexer::new(
+    //         r#"
+    //         if foo {
+    //             0;
+    //         } else if bar {
+    //             1;
+    //         } else if baz {
+    //             2;
+    //         } else {
+    //             3;
+    //         }
+    //     "#,
+    //     )
+    //     .collect_vec();
+    //
+    //     let expected_ast = Ok(Stmt::IfStmt {
+    //         condition: ExprKind::Identifier("foo".to_owned()),
+    //         if_then: Box::new(Stmt::Block(vec![Stmt::ExprStmt(
+    //             Literal::Integer(0).into(),
+    //         )])),
+    //         else_then: Some(Box::new(Stmt::IfStmt {
+    //             condition: ExprKind::Identifier("bar".to_owned()),
+    //             if_then: Box::new(Stmt::Block(vec![Stmt::ExprStmt(
+    //                 Literal::Integer(1).into(),
+    //             )])),
+    //             else_then: Some(Box::new(Stmt::IfStmt {
+    //                 condition: ExprKind::Identifier("baz".to_owned()),
+    //                 if_then: Box::new(Stmt::Block(vec![Stmt::ExprStmt(
+    //                     Literal::Integer(2).into(),
+    //                 )])),
+    //                 else_then: Some(Box::new(Stmt::Block(vec![Stmt::ExprStmt(
+    //                     Literal::Integer(3).into(),
+    //                 )]))),
+    //             })),
+    //         })),
+    //     });
+    //
+    //     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);
+    // }
 }