Added proper codegen and parsetree

4 files changed, 166 insertions, 45 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index c8f414e..a200e23 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -2,25 +2,37 @@ const std = @import("std");
 const parse = @import("parser.zig");
 
 pub const Generator = struct {
-    root: parse.NodeExit,
+    root: []const parse.NodeStmt,
     allocator: std.mem.Allocator,
     code: std.ArrayList(u8),
 
-    pub fn init(allocator: std.mem.Allocator, root: parse.NodeExit) Generator {
+    pub fn init(allocator: std.mem.Allocator, stmts: []const parse.NodeStmt) Generator {
         return .{
-            .root = root,
+            .root = stmts,
             .allocator = allocator,
             .code = std.ArrayList(u8).init(allocator),
         };
     }
 
-    fn genExit(self: *Generator) ![]const u8 {
+    pub fn deinit(self: *Generator) void {
+        self.code.deinit();
+    }
+
+    fn genExit(self: *Generator, expr: parse.NodeExpr) ![]const u8 {
         return try std.fmt.allocPrint(self.allocator,
             \\  mov rax, 60
-            \\  mov rdi, {}
+            \\  mov rdi, {d}
             \\  syscall
             \\
-        , .{self.root.expr.intLit.intLit});
+        , .{switch (expr) {
+            .intLit => expr.intLit.intlit.intLit,
+            else => return error.NotImplemented,
+        }});
+    }
+
+    fn genValue(self: *Generator) ![]const u8 {
+        _ = self;
+        return error.NotImplemented;
     }
 
     pub fn generate(self: *Generator) ![]const u8 {
@@ -28,9 +40,13 @@ pub const Generator = struct {
             \\global _start:
             \\
         );
-        const exitStmt = try self.genExit();
-        defer self.allocator.free(exitStmt);
-        try self.code.appendSlice(exitStmt);
-        return try self.code.toOwnedSlice();
+        for (self.root) |stmt| {
+            const code = switch (stmt) {
+                .exit => try self.genExit(stmt.exit.expr),
+                .value => try self.genValue(),
+            };
+            try self.code.appendSlice(code);
+        }
+        return self.code.items;
     }
 };
diff --git a/src/main.zig b/src/main.zig
index 472f1ad..cdf5a38 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -46,12 +46,14 @@ pub fn main() !void {
     const tokens = try tokenizer.tokenize();
 
     // Parse
-    var parser = parse.Parser.init(tokens);
+    var parser = parse.Parser.init(allocator, tokens);
+    defer parser.deinit();
     const tree = try parser.parse();
 
     // Codegen
     var generator = gen.Generator.init(allocator, tree);
     const code = try generator.generate();
+    std.debug.print("{s}", .{code});
     defer allocator.free(code);
     try outWriter.writeAll(code);
 
diff --git a/src/parser.zig b/src/parser.zig
index 422c9cb..18b2348 100644
--- a/src/parser.zig
+++ b/src/parser.zig
@@ -3,46 +3,112 @@ const tok = @import("tokenize.zig");
 const Iterator = tok.Iterator;
 const Token = tok.Token;
 
-const ParsingError = error{ InvalidExpression, ExpectedExit, ExpectedSemicolon };
+const ParsingError = error{
+    InvalidExpression,
+    ExpectedExit,
+    ExpectedSemicolon,
+    ExpectedEqual,
+    ExpectedIdentifier,
+    InvalidStatement,
+};
+
+pub const NodeExpr = union(enum) {
+    intLit: NodeIntlit,
+    ident: NodeIdent,
+};
 
-pub const NodeExpr = struct {
-    intLit: Token,
+pub const NodeStmt = union(enum) {
+    exit: NodeExit,
+    value: NodeValue,
+};
+
+pub const NodeValue = struct {
+    ident: Token,
+    value: NodeExpr,
+    isConst: bool,
 };
 
 pub const NodeExit = struct {
     expr: NodeExpr,
 };
 
+pub const NodeIntlit = struct {
+    intlit: Token,
+};
+
+pub const NodeIdent = struct {
+    ident: Token,
+};
+
 pub const Parser = struct {
     tokens: Iterator(Token),
+    allocator: std.mem.Allocator,
+    nodes: std.ArrayList(NodeStmt),
 
-    pub fn init(tokens: []Token) Parser {
+    pub fn init(allocator: std.mem.Allocator, tokens: []Token) Parser {
         return .{
+            .allocator = allocator,
             .tokens = Iterator(Token).init(tokens),
+            .nodes = std.ArrayList(NodeStmt).init(allocator),
         };
     }
 
-    fn parseExpr(self: *Parser) !NodeExpr {
-        if (tok.checkType(self.tokens.peek().?, tok.TokenType.intLit))
-            return NodeExpr{ .intLit = self.tokens.consume().? };
-        return ParsingError.InvalidExpression;
+    pub fn deinit(self: *Parser) void {
+        self.nodes.deinit();
     }
 
-    pub fn parse(self: *Parser) !NodeExit {
-        var root: NodeExit = undefined;
-        while (self.tokens.peek()) |token| {
-            switch (token) {
-                .exit => {
-                    self.tokens.skip();
-                    root.expr = try self.parseExpr();
-                    if (!tok.checkType(self.tokens.peek().?, tok.TokenType.semiCol))
-                        return ParsingError.ExpectedSemicolon;
-                    self.tokens.skip();
+    fn parseExpr(self: *Parser) !NodeExpr {
+        return switch (self.tokens.peek().?) {
+            .intLit => NodeExpr{
+                .intLit = NodeIntlit{
+                    .intlit = (try self.tokens.consume(.intLit)).?,
+                },
+            },
+            .ident => NodeExpr{
+                .ident = NodeIdent{
+                    .ident = (try self.tokens.consume(.ident)).?,
                 },
-                else => return ParsingError.ExpectedExit,
-            }
-        }
-        return root;
+            },
+            else => ParsingError.InvalidExpression,
+        };
+    }
+
+    fn parseStmt(self: *Parser) !NodeStmt {
+        return switch (self.tokens.peek().?) {
+            .exit => NodeStmt{ .exit = try self.parseExit() },
+            .constant => NodeStmt{ .value = try self.parseValue(true) },
+            .variable => NodeStmt{ .value = try self.parseValue(false) },
+            else => ParsingError.InvalidStatement,
+        };
+    }
+
+    fn parseExit(self: *Parser) !NodeExit {
+        _ = try self.tokens.consume(.exit);
+        const expr = try self.parseExpr();
+        _ = try self.tokens.consume(.semiCol);
+        return NodeExit{
+            .expr = expr,
+        };
+    }
+
+    fn parseValue(self: *Parser, isConst: bool) !NodeValue {
+        self.tokens.skip();
+        const ident = (try self.tokens.consume(.ident)).?;
+        _ = try self.tokens.consume(.equal);
+        const expr = try self.parseExpr();
+        _ = try self.tokens.consume(.semiCol);
+        return NodeValue{
+            .ident = ident,
+            .value = expr,
+            .isConst = isConst,
+        };
+    }
+
+    pub fn parse(self: *Parser) ![]const NodeStmt {
+        while (self.tokens.peek()) |_|
+            try self.nodes.append(try self.parseStmt());
+
+        return self.nodes.items;
     }
 };
 
@@ -52,8 +118,20 @@ test "Parser" {
     var tokenizer = tok.Tokenizer.init(std.testing.allocator, src);
     defer tokenizer.deinit();
     const toks = try tokenizer.tokenize();
-    var parser = Parser.init(toks);
+    var parser = Parser.init(std.testing.allocator, toks);
+    defer parser.deinit();
     const parseTree = try parser.parse();
-    const exp = NodeExit{ .expr = NodeExpr{ .intLit = Token{ .intLit = 120 } } };
-    try expect(std.meta.eql(parseTree, exp));
+    const exp: []const NodeStmt = &[_]NodeStmt{NodeStmt{
+        .exit = NodeExit{
+            .expr = NodeExpr{
+                .intLit = NodeIntlit{
+                    .intlit = Token{
+                        .intLit = 120,
+                    },
+                },
+            },
+        },
+    }};
+    for (parseTree, exp) |stmt, expStmt|
+        try expect(std.meta.eql(stmt, expStmt));
 }
diff --git a/src/tokenize.zig b/src/tokenize.zig
index f28cdcf..c78a59b 100644
--- a/src/tokenize.zig
+++ b/src/tokenize.zig
@@ -8,23 +8,29 @@ const TokenizeError = error{
 pub const TokenType = enum {
     ident,
     intLit,
+    constant,
+    variable,
     exit,
     plus,
     minus,
     star,
     slash,
     semiCol,
+    equal,
 };
 
 pub const Token = union(TokenType) {
     ident: []const u8,
     intLit: i32,
+    constant,
+    variable,
     exit,
     plus,
     minus,
     star,
     slash,
     semiCol,
+    equal,
 
     pub fn fromChar(char: u8) !Token {
         return switch (char) {
@@ -33,9 +39,18 @@ pub const Token = union(TokenType) {
             '*' => .star,
             '/' => .slash,
             ';' => .semiCol,
+            '=' => .equal,
             else => TokenizeError.UnknownToken,
         };
     }
+
+    pub fn fromStr(str: []const u8) Token {
+        const eql = std.mem.eql;
+        if (eql(u8, str, "exit")) return .exit;
+        if (eql(u8, str, "const")) return .constant;
+        if (eql(u8, str, "var")) return .variable;
+        return Token{ .ident = str };
+    }
 };
 
 pub fn checkType(tok: Token, comptime typ: TokenType) bool {
@@ -67,13 +82,18 @@ pub fn Iterator(comptime typ: type) type {
         }
 
         /// Get current item and iterate index
-        pub fn consume(self: *Iterator(typ)) ?typ {
+        pub fn next(self: *Iterator(typ)) ?typ {
             const ret = self.peek();
             self.index += 1;
             return ret;
         }
-        /// Get current item and iterate index
-        pub const next = consume;
+
+        pub fn consume(self: *Iterator(typ), comptime expected: TokenType) !?typ {
+            if (typ != Token) return error.TokenIteratorOnly;
+            if (!checkType(self.peek().?, expected))
+                return error.ExpectedToken;
+            return self.next();
+        }
 
         /// Skip over current item
         pub fn skip(self: *Iterator(typ)) void {
@@ -100,6 +120,10 @@ pub const Tokenizer = struct {
 
     /// Releases allocated memory
     pub fn deinit(self: *Tokenizer) void {
+        for (self.toks.items) |token| {
+            if (checkType(token, TokenType.ident))
+                self.allocator.free(token.ident);
+        }
         self.toks.deinit();
     }
 
@@ -113,7 +137,7 @@ pub const Tokenizer = struct {
                 ' ', '\n', '\t' => self.src.skip(),
                 '0'...'9' => {
                     while (std.ascii.isDigit(self.src.peek().?))
-                        try buff.append(self.src.consume().?);
+                        try buff.append(self.src.next().?);
 
                     const num: i32 = try std.fmt.parseInt(i32, buff.items, 10);
                     try self.toks.append(.{ .intLit = num });
@@ -121,13 +145,14 @@ pub const Tokenizer = struct {
                 },
                 'a'...'z', 'A'...'Z' => {
                     while (std.ascii.isAlphanumeric(self.src.peek().?))
-                        try buff.append(self.src.consume().?);
-                    if (std.mem.eql(u8, "exit", buff.items)) {
-                        try self.toks.append(.exit);
-                    } else return TokenizeError.UnknownToken;
+                        try buff.append(self.src.next().?);
+                    const str = try buff.toOwnedSlice();
+                    const token = Token.fromStr(str);
+                    try self.toks.append(token);
+                    if (!checkType(token, TokenType.ident)) self.allocator.free(str);
                     buff.clearAndFree();
                 },
-                else => self.toks.append(try Token.fromChar(self.src.consume().?)),
+                else => self.toks.append(try Token.fromChar(self.src.next().?)),
             };
         }
         return self.toks.items;