ast: add traversable tree structure

Cargo.toml

@@ -23,3 +23,4 @@ rust-embed = "5.7.0"
 inkwell = { version = "0.1.0-beta.2", features = ["llvm10-0"], optional = true }
 regex = "1.5.4"
 lazy_static = "1.4.0"
+indextree = "4.3.1"

src/ast/mod.rs

@@ -1,5 +1,6 @@
 use crate::lexer::*;
 use core::convert::TryFrom;
+use indextree::Arena;
 use std::collections::HashMap;
 /**
  * Copyright 2021 Garrit Franke
@@ -21,6 +22,18 @@ use std::collections::HashSet;
 pub mod types;
 use types::Type;
 
+#[derive(Debug, Clone)]
+pub enum ASTNode {
+    ModuleNode(Module),
+    FunctionNode(Function),
+    StructDefNode(StructDef),
+    VariableNode(Variable),
+    StatementNode(Statement),
+    ExpressionNode(Expression),
+    MatchArmNode(MatchArm),
+    BinOpNode(BinOp),
+}
+
 /// Table that contains all symbol and its types
 pub type SymbolTable = HashMap<String, Option<Type>>;
 

src/main.rs

@@ -1,3 +1,4 @@
+extern crate indextree;
 extern crate lazy_static;
 extern crate regex;
 /**

src/parser/parser.rs

@@ -19,12 +19,14 @@ use crate::lexer::Position;
 use crate::lexer::{Token, TokenKind};
 use crate::parser::infer::infer;
 use crate::util::string_util::highlight_position_in_file;
+use indextree::Arena;
 use std::convert::TryFrom;
 use std::iter::Peekable;
 use std::vec::IntoIter;
 
 pub struct Parser {
     pub path: String,
+    pub ast: Arena<Box<ASTNode>>,
     tokens: Peekable<IntoIter<Token>>,
     peeked: Vec<Token>,
     current: Option<Token>,
@@ -45,13 +47,16 @@ impl Parser {
             current: None,
             prev: None,
             raw,
+            ast: Arena::new(),
         }
     }
 
     pub fn parse(&mut self) -> Result<Module, String> {
         let mut program = self.parse_module()?;
         // infer types
         infer(&mut program);
+        self.ast
+            .new_node(Box::new(ASTNode::ModuleNode(program.clone())));
 
         Ok(program)
     }

src/parser/rules.rs

@@ -3,6 +3,7 @@ use crate::ast::types::Type;
 use crate::ast::*;
 use crate::lexer::Keyword;
 use crate::lexer::{TokenKind, Value};
+use indextree::NodeId;
 use std::collections::HashMap;
 /**
  * Copyright 2020 Garrit Franke
@@ -22,39 +23,51 @@ use std::collections::HashMap;
 use std::collections::HashSet;
 use std::convert::TryFrom;
 
+impl Parser {
+    pub fn append_node(&mut self, parent: &mut NodeId, child: ASTNode) {
+        let child_node = self.ast.new_node(Box::new(child));
+        parent.append(child_node, &mut self.ast)
+    }
+}
+
 impl Parser {
     pub fn parse_module(&mut self) -> Result<Module, String> {
-        let mut functions = Vec::new();
-        let mut structs = Vec::new();
-        let mut imports = HashSet::new();
-        let globals = Vec::new();
+        let mut module = Module {
+            func: Vec::new(),
+            structs: Vec::new(),
+            globals: Vec::new(),
+            path: self.path.clone(),
+            imports: HashSet::new(),
+        };
+        let mut module_node = self
+            .ast
+            .new_node(Box::new(ASTNode::ModuleNode(module.clone())));
 
         while self.has_more() {
             let next = self.peek()?;
             match next.kind {
-                TokenKind::Keyword(Keyword::Function) => functions.push(self.parse_function()?),
-                TokenKind::Keyword(Keyword::Import) => {
-                    imports.insert(self.parse_import()?);
+                TokenKind::Keyword(Keyword::Function) => {
+                    let func = self.parse_function(&mut module_node)?;
+                    module.func.push(func);
                 }
-                TokenKind::Keyword(Keyword::Struct) => {
-                    structs.push(self.parse_struct_definition()?)
+                TokenKind::Keyword(Keyword::Import) => {
+                    module.imports.insert(self.parse_import()?);
                 }
+                TokenKind::Keyword(Keyword::Struct) => module
+                    .structs
+                    .push(self.parse_struct_definition(&mut module_node)?),
                 _ => return Err(format!("Unexpected token: {}", next.raw)),
             }
         }
 
         // TODO: Populate imports
 
-        Ok(Module {
-            func: functions,
-            structs,
-            globals,
-            path: self.path.clone(),
-            imports,
-        })
+        dbg!(&self.ast.count());
+
+        Ok(module)
     }
 
-    fn parse_struct_definition(&mut self) -> Result<StructDef, String> {
+    fn parse_struct_definition(&mut self, mut mod_node: &mut NodeId) -> Result<StructDef, String> {
         self.match_keyword(Keyword::Struct)?;
         let name = self.match_identifier()?;
 
@@ -65,7 +78,7 @@ impl Parser {
             let next = self.peek()?;
             match next.kind {
                 TokenKind::Keyword(Keyword::Function) => {
-                    methods.push(self.parse_function()?);
+                    methods.push(self.parse_function(&mut mod_node)?);
                 }
                 TokenKind::Identifier(_) => fields.push(self.parse_typed_variable()?),
                 _ => {
@@ -141,7 +154,7 @@ impl Parser {
     /// To reduce code duplication, this method can be either be used to parse a function or a method.
     /// If a function is parsed, the `fn` keyword is matched.
     /// If a method is parsed, `fn` will be omitted
-    fn parse_function(&mut self) -> Result<Function, String> {
+    fn parse_function(&mut self, mut mod_node: &mut NodeId) -> Result<Function, String> {
         self.match_keyword(Keyword::Function)?;
         let name = self.match_identifier()?;
 
@@ -161,12 +174,16 @@ impl Parser {
 
         let body = self.parse_block()?;
 
-        Ok(Function {
+        let func = Function {
             name,
             arguments,
             body,
             ret_type: ty,
-        })
+        };
+
+        self.append_node(&mut mod_node, ASTNode::FunctionNode(func.clone()));
+
+        Ok(func)
     }
 
     fn parse_import(&mut self) -> Result<String, String> {