4 files changed, 96 insertions, 14 deletions

diff --git a/sloth/src/parser/ast.rs b/sloth/src/parser/ast.rs
index 281cd86..1624b0d 100644
--- a/sloth/src/parser/ast.rs
+++ b/sloth/src/parser/ast.rs
@@ -218,7 +218,7 @@ pub enum StmtKind {
     DefineVariable {
         identifier: String,
         value: Expr,
-        typ: String,
+        typ: TypeIdentifier,
     },
     AssignVariable {
         identifier: String,
@@ -235,7 +235,7 @@ pub enum StmtKind {
 pub struct Function {
     pub identifier: String,
     pub inputs: Vec<FunctionInput>,
-    pub output: Option<String>,
+    pub output: Option<TypeIdentifier>,
     pub kind: FunctionKind,
 }
 
@@ -248,7 +248,30 @@ pub enum FunctionKind {
 #[derive(PartialEq, Clone, Debug)]
 pub struct FunctionInput {
     pub identifier: String,
-    pub typ: String,
+    pub typ: TypeIdentifier,
+}
+
+#[derive(PartialEq, Clone, Debug)]
+pub struct TypeIdentifier {
+    pub name: String,
+    pub is_list: bool,
+    pub list_len: u32,
+}
+
+impl Display for TypeIdentifier {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        if self.is_list {
+            write!(f, "[")?;
+        }
+
+        write!(f, "{}", self.name)?;
+
+        if self.is_list {
+            write!(f, "]")?;
+        }
+
+        Ok(())
+    }
 }
 
 #[derive(Debug, Clone, PartialEq)]
@@ -258,7 +281,7 @@ pub enum Literal {
     Boolean(bool),
     Character(char),
     String(String),
-    Array(Vec<ExprKind>),
+    Array(Vec<Expr>),
 }
 
 impl From<lexer::Literal> for Literal {
diff --git a/sloth/src/parser/expr.rs b/sloth/src/parser/expr.rs
index 7ec283f..3decf5d 100644
--- a/sloth/src/parser/expr.rs
+++ b/sloth/src/parser/expr.rs
@@ -1,4 +1,4 @@
-use super::ast::{Expr, UnaryOp};
+use super::ast::{Expr, Literal, UnaryOp};
 use super::AstParser;
 use crate::lexer::TokenType;
 use crate::parser::ast::{BinaryOp, ExprKind};
@@ -67,6 +67,20 @@ impl<'a> AstParser<'a> {
             TokenType::Literal(literal) => ExprKind::Literal(literal.into()),
             TokenType::Identifier(identifier) => ExprKind::Identifier(identifier),
 
+            TokenType::OpeningBracket => {
+                let mut contents = Vec::new();
+                while !self.eof() && self.peek().tt != TokenType::ClosingBracket {
+                    contents.push(self.expression()?);
+                    if !self.advance_if_eq(&TokenType::Comma) {
+                        break;
+                    }
+                }
+
+                self.consume(TokenType::ClosingBracket, "Expected ']'")?;
+
+                ExprKind::Literal(Literal::Array(contents))
+            }
+
             TokenType::OpeningParen => {
                 let expr = self.expression()?;
                 self.consume(TokenType::ClosingParen, "Must end grouping with ')'")?;
diff --git a/sloth/src/parser/mod.rs b/sloth/src/parser/mod.rs
index 748a0da..6dd534a 100644
--- a/sloth/src/parser/mod.rs
+++ b/sloth/src/parser/mod.rs
@@ -3,7 +3,7 @@ pub mod expr;
 pub mod graph;
 pub mod stmt;
 
-use self::ast::{Literal, Stmt, StmtKind};
+use self::ast::{Literal, Stmt, StmtKind, TypeIdentifier};
 use crate::lexer::{Token, TokenType};
 use crate::symtable::SymbolTable;
 
@@ -117,6 +117,31 @@ impl<'a> AstParser<'a> {
         Ok(identifier)
     }
 
+    pub fn consume_type(&mut self) -> Result<TypeIdentifier, ParsingError> {
+        let is_list = self.peek().tt == TokenType::OpeningBracket;
+
+        if is_list {
+            self.consume(TokenType::OpeningBracket, "Expected '['")?;
+        }
+
+        let name = self.consume_identifier()?;
+
+        let mut list_len = 0;
+        if is_list {
+            if let Literal::Integer(i) = self.consume_literal()? {
+                list_len = i as u32;
+            }
+
+            self.consume(TokenType::ClosingBracket, "Expected ']'")?;
+        }
+
+        Ok(TypeIdentifier {
+            name,
+            is_list,
+            list_len,
+        })
+    }
+
     pub fn reserve_id(&mut self) -> i32 {
         let id = self.id;
         self.id += 1;
diff --git a/sloth/src/parser/stmt.rs b/sloth/src/parser/stmt.rs
index 973b9c3..d8c091d 100644
--- a/sloth/src/parser/stmt.rs
+++ b/sloth/src/parser/stmt.rs
@@ -91,7 +91,7 @@ impl<'a> AstParser<'a> {
         // Get the identifier and type
         let identifier = self.consume_identifier()?;
         self.consume(TokenType::Colon, "Expected ':'")?;
-        let typ = self.consume_identifier()?;
+        let typ = self.consume_type()?;
 
         // Get the default value
         self.consume(TokenType::Eq, "Expected '='")?;
@@ -127,7 +127,7 @@ impl<'a> AstParser<'a> {
         while matches!(self.peek().tt, TokenType::Identifier(_)) {
             let input_identifier = self.consume_identifier()?;
             self.consume(TokenType::Colon, "Expected ':'")?;
-            let input_type = self.consume_identifier()?;
+            let input_type = self.consume_type()?;
 
             inputs.push(FunctionInput {
                 identifier: input_identifier,
@@ -144,8 +144,11 @@ impl<'a> AstParser<'a> {
         self.consume(TokenType::ClosingParen, "Expected ')'")?;
 
         // Get the function output
-        let output = if matches!(self.peek().tt, TokenType::Identifier(_)) {
-            Some(self.consume_identifier()?)
+        let output = if matches!(
+            self.peek().tt,
+            TokenType::Identifier(_) | TokenType::OpeningBracket
+        ) {
+            Some(self.consume_type()?)
         } else {
             None
         };
@@ -257,6 +260,7 @@ mod tests {
     use crate::lexer::Lexer;
     use crate::parser::ast::{
         BinaryOp, Expr, ExprKind, Function, FunctionInput, FunctionKind, Literal, Stmt,
+        TypeIdentifier,
     };
     use crate::symtable::SymbolTable;
 
@@ -304,7 +308,11 @@ mod tests {
                     ExprKind::Literal(Literal::Integer(3)),
                 )),
             }),
-            typ: "Int".to_string(),
+            typ: TypeIdentifier {
+                name: "Int".to_string(),
+                is_list: false,
+                list_len: 0,
+            },
         }));
 
         let mut parser = AstParser::new(tokens, SymbolTable::new());
@@ -335,9 +343,17 @@ mod tests {
                 identifier: "foo".to_owned(),
                 inputs: vec![FunctionInput {
                     identifier: "bar".to_owned(),
-                    typ: "Int".to_owned(),
+                    typ: TypeIdentifier {
+                        name: "Int".to_owned(),
+                        is_list: false,
+                        list_len: 0,
+                    },
                 }],
-                output: Some("Int".to_owned()),
+                output: Some(TypeIdentifier {
+                    name: "Int".to_owned(),
+                    is_list: false,
+                    list_len: 0,
+                }),
                 kind: FunctionKind::Normal {
                     body: Box::new(Stmt::without_table(
                         10,
@@ -355,7 +371,11 @@ mod tests {
                                         Literal::Integer(1).into(),
                                     )),
                                 }),
-                                typ: "Int".to_owned(),
+                                typ: TypeIdentifier {
+                                    name: "Int".to_owned(),
+                                    is_list: false,
+                                    list_len: 0,
+                                },
                             }),
                             Stmt::without_table(7, StmtKind::AssignVariable {
                                 identifier: "baz".to_owned(),