4 files changed, 40 insertions, 23 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index 5636dbd..59ef03b 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -223,7 +223,9 @@ test "Codegen exit" {
     var pop = symb.Populator.init(std.testing.allocator);
     var treeNode = parseTree.asNode();
     try pop.populateSymtable(&treeNode);
-    var gen = Generator.init(std.testing.allocator, parseTree);
+    var arena = std.heap.ArenaAllocator.init(std.testing.allocator);
+    defer arena.deinit();
+    var gen = Generator.init(arena.allocator(), parseTree);
     defer gen.deinit();
     const actual = try gen.generate();
     try expect(std.mem.eql(u8, actual, expected));
diff --git a/src/main.zig b/src/main.zig
index d748461..1e54dec 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -58,7 +58,9 @@ pub fn main() !void {
     try pop.populateSymtable(&treeNode);
 
     // Codegen
-    var generator = gen.Generator.init(allocator, tree);
+    var arena = std.heap.ArenaAllocator.init(allocator);
+    defer arena.deinit();
+    var generator = gen.Generator.init(arena.allocator(), tree);
     defer generator.deinit();
     const code = try generator.generate();
     try outWriter.writeAll(code);
@@ -69,7 +71,6 @@ pub fn main() !void {
     const ldproc = try std.process.Child.run(.{ .argv = &ldargv, .allocator = allocator });
     defer allocator.free(ldproc.stdout);
     defer allocator.free(ldproc.stderr);
-    std.debug.print("code: \n{s}", .{code});
 }
 
 /// Get file extension based on filename
diff --git a/src/parser.zig b/src/parser.zig
index 8cd539f..08f098c 100644
--- a/src/parser.zig
+++ b/src/parser.zig
@@ -113,16 +113,30 @@ pub const Parser = struct {
         };
     }
 
+    fn dinitHelper(self: *Parser, node: NodeStmt) !void {
+        switch (node.kind) {
+            .block => |blk| {
+                const children = try node.children(self.allocator);
+                defer self.allocator.free(children);
+                for (children) |child| try self.dinitHelper(child.Stmt);
+                self.allocator.free(blk);
+                node.symtable.deinit();
+            },
+            .function => |fun| {
+                const children = try node.children(self.allocator);
+                defer self.allocator.free(children);
+                for (children) |child| try self.dinitHelper(child.Stmt);
+                self.allocator.destroy(fun.block);
+            },
+            else => {},
+        }
+    }
+
     pub fn deinit(self: *Parser) void {
         for (self.nodes.items) |node| {
-            switch (node.kind) {
-                .block => |blk| self.allocator.free(blk),
-                .function => |fun| {
-                    self.allocator.free(fun.block.kind.block);
-                    self.allocator.destroy(fun.block);
-                },
-                else => {},
-            }
+            self.dinitHelper(node) catch |err| {
+                if (err == error.OutOfMemory) {}
+            };
         }
         self.nodes.deinit();
     }
diff --git a/src/symtable.zig b/src/symtable.zig
index f37ee12..60e871c 100644
--- a/src/symtable.zig
+++ b/src/symtable.zig
@@ -161,7 +161,6 @@ pub const Populator = struct {
                 const table: *SymbolTable = stmt.symtable;
                 switch (stmt.kind) {
                     .defVar => |variable| {
-                        std.debug.print("Populated {s}\n", .{variable.ident.ident});
                         const symbol: Symbol = try self.buildValueSymb(
                             table,
                             if (variable.expr.typ) |typ| typ else pars.TypeIdent{ .ident = "i32", .list = false },
@@ -175,15 +174,12 @@ pub const Populator = struct {
                             if (value.expr.typ) |typ| typ else pars.TypeIdent{ .ident = "i32", .list = false },
                             false,
                         );
-                        std.debug.print("Populated {s}\n", .{value.ident.ident});
                         if (!try table.insert(value.ident.ident, symbol)) return error.FailedToInsert;
                     },
                     .block => {
                         const children = try stmt.children(self.allocator);
                         defer self.allocator.free(children);
-                        std.debug.print("Populated Block\n", .{});
                         for (children) |child| {
-                            std.debug.print("Child: {d}\n", .{child.Stmt.id});
                             try self.populateSymtable(&child);
                         }
                     },
@@ -193,9 +189,9 @@ pub const Populator = struct {
                             fun.args,
                             fun.retType,
                         );
-                        std.debug.print("Populated Function {s}\n", .{fun.ident.ident});
                         if (!try table.insert(fun.ident.ident, symbol)) return error.FailedToInsert;
                         const children = try stmt.children(self.allocator);
+                        defer self.allocator.free(children);
                         for (children) |child| try self.populateSymtable(&child);
                     },
 
@@ -221,17 +217,21 @@ pub const Populator = struct {
 
         const output = try self.allocator.create(SymbType);
         output.* = if (retType) |typ| table.getType(typ).? else SymbType.Void;
+        const typ = SymbType{
+            .Function = .{
+                .output = output,
+                .input = input,
+            },
+        };
+        const id = self.reserveId();
+
+        _ = try table.insert("func_" ++ .{@as(u8, @truncate(id))}, typ.toSymb());
 
         return Symbol{
             .Value = SymbValue{
                 .mut = true,
-                .id = self.reserveId(),
-                .typ = SymbType{
-                    .Function = .{
-                        .input = input,
-                        .output = output,
-                    },
-                },
+                .id = id,
+                .typ = typ,
             },
         };
     }