path: root/src/parse.zig

const std = @import("std");
const Allocator = std.mem.Allocator;

const root = @import("root");
const Lexer = root.Lexer;
const Token = Lexer.Token;
const Location = Lexer.Location;

fn Fmt(T: type) type {
    return std.fmt.Formatter(struct {
        fn format(
            data: struct { T, []const u8, usize },
            comptime _: []const u8,
            _: std.fmt.FormatOptions,
            writer: anytype,
        ) !void {
            const self, const file_source, const indent = data;
            return self.format(writer, file_source, indent);
        }
    }.format);
}

pub fn fmt(tree: anytype, source: []const u8, indent: usize) Fmt(@TypeOf(tree)) {
    return .{ .data = .{ tree, source, indent } };
}

pub const File = struct {
    decls: []Decl,

    const Decl = struct {
        loc: Location,
        inner: Stmt.Type.AssignVar,
    };

    fn format(self: File, writer: anytype, source: []const u8, indent: usize) !void {
        for (0.., self.decls) |i, decl| {
            try writer.print("{s}{s} {s} {}\n", .{
                if (i == 0) "" else "\n",
                decl.inner.ident.getIdent(source),
                ":=",
                fmt(decl.inner.value, source, indent),
            });
        }
    }
};

pub const Block = struct {
    loc: Location,
    stmts: []Stmt,

    fn format(self: Block, writer: anytype, source: []const u8, indent: usize) !void {
        try writer.writeAll("{\n");
        for (self.stmts) |stmt| {
            try writer.print("{}\n", .{fmt(stmt, source, indent + 4)});
        }
        try writer.writeByteNTimes(' ', indent);
        try writer.writeAll("}");
    }
};

pub const Stmt = struct {
    loc: Location,
    type: Type,

    pub const Type = union(enum) {
        expr: *const Expr,
        assign_var: AssignVar,
        block: Block,
        @"while": While,

        pub const AssignVar = struct {
            ident: Location,
            is_decl: bool,
            value: *const Expr,
        };

        pub const While = struct {
            cond: *const Expr,
            do: Block,
        };
    };

    fn format(self: Stmt, writer: anytype, source: []const u8, indent: usize) !void {
        try writer.writeByteNTimes(' ', indent);
        return switch (self.type) {
            .expr => |expr| writer.print("{}", .{fmt(expr, source, indent)}),
            .block => |b| writer.print("{}", .{fmt(b, source, indent)}),
            .assign_var => |assign_var| writer.print("{s} {s} {}", .{
                assign_var.ident.getIdent(source),
                if (assign_var.is_decl) ":=" else "=",
                fmt(assign_var.value, source, indent),
            }),
            .@"while" => |@"while"| try writer.print("while {} {}", .{
                fmt(@"while".cond, source, indent),
                fmt(@"while".do, source, indent),
            }),
        };
    }
};

pub const Expr = struct {
    loc: Location,
    type: Type,

    pub const Type = union(enum) {
        integer_literal,
        bin_op: BinOp,
        call: Call,
        identifier,
        @"if": If,
        proc: Proc,

        pub const BinOp = struct {
            lhs: *const Expr,
            rhs: *const Expr,
            op: Op,

            const Op = enum {
                plus,
                minus,
                left_angle,
                right_angle,
                left_angle_equal,
                right_angle_equal,

                pub fn format(self: Op, comptime _: []const u8, _: std.fmt.FormatOptions, writer: anytype) !void {
                    try writer.writeAll(switch (self) {
                        .plus => "+",
                        .minus => "-",
                        .left_angle => "<",
                        .right_angle => ">",
                        .left_angle_equal => "<=",
                        .right_angle_equal => ">=",
                    });
                }
            };
        };

        pub const Call = struct {
            proc: *const Expr,
            arg: *const Expr,
        };

        pub const If = struct {
            cond: *const Expr,
            then: Block,
            @"else": ?Block,
        };

        pub const Proc = struct {
            body: Block,
            param: ?Location,
        };
    };

    fn format(self: Expr, writer: anytype, source: []const u8, indent: usize) !void {
        switch (self.type) {
            .integer_literal => try writer.print("{}", .{self.loc.getInt(source)}),
            .bin_op => |bin_op| {
                try writer.print("{} {} {}", .{ fmt(bin_op.lhs, source, indent), bin_op.op, fmt(bin_op.rhs, source, indent) });
            },
            .call => |call| {
                try writer.print("{}({})", .{ fmt(call.proc, source, indent), fmt(call.arg, source, indent) });
            },
            .identifier => try writer.print("{s}", .{self.loc.getIdent(source)}),
            .@"if" => |@"if"| {
                try writer.print("if {} {}", .{
                    fmt(@"if".cond, source, indent),
                    fmt(@"if".then, source, indent),
                });
                if (@"if".@"else") |@"else"| {
                    try writer.print(" else {}", .{fmt(@"else", source, indent)});
                }
            },
            .proc => |proc| {
                try writer.print("proc() {}", .{fmt(proc.body, source, indent)});
            },
        }
    }
};

const ParseError = error{
    OutOfMemory,
    ExpectedRightParen,
    UnexpectedToken,
    InvalidAssignTarget,
    ExprStatementMustBeCall,
};

pub fn file(allocator: Allocator, lexer: *Lexer) !File {
    var decls: std.ArrayList(File.Decl) = .init(allocator);
    while (lexer.peek().type != .eof) {
        const stmt = try parseStatement(allocator, lexer);
        if (stmt.type != .assign_var or !stmt.type.assign_var.is_decl) {
            return error.ExpectedProcedureDeclaration;
        }
        try decls.append(.{
            .loc = stmt.loc,
            .inner = stmt.type.assign_var,
        });
    }
    return .{
        .decls = try decls.toOwnedSlice(),
    };
}

fn parseBlock(allocator: Allocator, lexer: *Lexer) !Block {
    const left_curly = try mustEat(lexer, .left_curly);
    var stmts: std.ArrayList(Stmt) = .init(allocator);
    while (lexer.peek().type != .right_curly) {
        try stmts.append(try parseStatement(allocator, lexer));
    }
    const right_curly = try mustEat(lexer, .right_curly);
    return .{
        .loc = left_curly.loc.combine(right_curly.loc),
        .stmts = try stmts.toOwnedSlice(),
    };
}

fn parseStatement(allocator: Allocator, lexer: *Lexer) ParseError!Stmt {
    switch (lexer.peek().type) {
        .left_curly => {
            const b = try parseBlock(allocator, lexer);
            return .{
                .loc = b.loc,
                .type = .{ .block = b },
            };
        },
        .@"while" => {
            const @"while" = lexer.next();
            const cond = try expression(allocator, lexer);
            const do = try parseBlock(allocator, lexer);
            return .{
                .loc = @"while".loc.combine(do.loc),
                .type = .{ .@"while" = .{
                    .cond = cond,
                    .do = do,
                } },
            };
        },
        else => {
            const lhs = try expression(allocator, lexer);

            const colon = if (lexer.peek().type == .colon)
                try mustEat(lexer, .colon)
            else
                null;
            if (colon != null or lexer.peek().type == .equal) {
                _ = try mustEat(lexer, .equal);
                const value = try expression(allocator, lexer);
                if (lhs.type != .identifier) {
                    std.debug.print("Invalid assign target. Found '{s}', expected an identifier.", .{@tagName(lhs.type)});
                    return error.InvalidAssignTarget;
                }
                return .{
                    .loc = lhs.loc.combine(value.loc),
                    .type = .{ .assign_var = .{ .ident = lhs.loc, .is_decl = colon != null, .value = value } },
                };
            } else if (lhs.type != .call) {
                return error.ExprStatementMustBeCall;
            }
            return .{
                .loc = lhs.loc,
                .type = .{ .expr = lhs },
            };
        },
    }
}

fn expression(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
    return parseProc(allocator, lexer);
}

fn parseProc(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
    if (lexer.peek().type != .proc) return parseComparisons(allocator, lexer);
    const proc = try mustEat(lexer, .proc);
    _ = try mustEat(lexer, .left_paren);
    const param_tok = lexer.next();
    const param: ?Location = switch (param_tok.type) {
        .identifier => param_tok.loc,
        .right_paren => null,
        else => |ty| return expected(&.{ .identifier, .right_paren }, ty),
    };
    if (param != null) _ = try mustEat(lexer, .right_paren);
    const body = try parseBlock(allocator, lexer);

    return allocate(Expr, allocator, .{
        .loc = proc.loc.combine(body.loc),
        .type = .{ .proc = .{ .body = body, .param = param } },
    });
}

fn parseComparisons(allocator: Allocator, lexer: *Lexer) ParseError!*Expr {
    const lhs = try parseTerms(allocator, lexer);

    const op: Expr.Type.BinOp.Op = switch (lexer.peek().type) {
        .left_angle => .left_angle,
        .right_angle => .right_angle,
        .left_angle_equal => .left_angle_equal,
        .right_angle_equal => .right_angle_equal,
        else => return lhs,
    };
    _ = lexer.next();

    const rhs = try parseTerms(allocator, lexer);
    return try allocate(Expr, allocator, .{
        .loc = lhs.loc.combine(rhs.loc),
        .type = .{ .bin_op = .{ .lhs = lhs, .op = op, .rhs = rhs } },
    });
}

fn parseTerms(allocator: Allocator, lexer: *Lexer) !*Expr {
    var lhs = try parseIf(allocator, lexer);
    while (true) {
        const op: Expr.Type.BinOp.Op = switch (lexer.peek().type) {
            .plus => .plus,
            .minus => .minus,
            else => break,
        };
        _ = lexer.next();

        const rhs = try parseIf(allocator, lexer);
        lhs = try allocate(Expr, allocator, .{
            .loc = lhs.loc.combine(rhs.loc),
            .type = .{ .bin_op = .{ .lhs = lhs, .op = op, .rhs = rhs } },
        });
    }
    return lhs;
}

fn parseIf(allocator: Allocator, lexer: *Lexer) !*Expr {
    switch (lexer.peek().type) {
        .@"if" => {
            const @"if" = lexer.next();
            const cond = try expression(allocator, lexer);
            const then = try parseBlock(allocator, lexer);
            const @"else" = if (lexer.peek().type == .@"else") blk: {
                _ = lexer.next();
                break :blk try parseBlock(allocator, lexer);
            } else null;
            return try allocate(Expr, allocator, .{
                .loc = @"if".loc.combine((@"else" orelse then).loc),
                .type = .{ .@"if" = .{
                    .cond = cond,
                    .then = then,
                    .@"else" = @"else",
                } },
            });
        },
        else => return try parseInvocations(allocator, lexer),
    }
}

fn parseInvocations(allocator: Allocator, lexer: *Lexer) !*Expr {
    var proc = try parsePrimaryExpr(allocator, lexer);
    while (true) {
        if (lexer.peek().type != .left_paren) break;
        _ = lexer.next();
        const arg = try expression(allocator, lexer);
        _ = try mustEat(lexer, .right_paren);
        proc = try allocate(Expr, allocator, .{
            .loc = proc.loc.combine(arg.loc),
            .type = .{ .call = .{ .proc = proc, .arg = arg } },
        });
    }
    return proc;
}

fn parsePrimaryExpr(allocator: Allocator, lexer: *Lexer) !*Expr {
    const token = lexer.next();
    return allocate(Expr, allocator, switch (token.type) {
        .left_paren => {
            const res = expression(allocator, lexer);
            const right_paren = lexer.next();
            if (right_paren.type != .right_paren)
                return error.ExpectedRightParen;
            return res;
        },
        .integer_literal => .{ .loc = token.loc, .type = .integer_literal },
        .identifier => .{ .loc = token.loc, .type = .identifier },
        else => |t| {
            std.debug.print("Expected '(', integer literal, or identifier. Got '{s}'\n", .{@tagName(t)});
            return error.UnexpectedToken;
        },
    });
}

fn expected(one_of: []const Token.Type, got: Token.Type) error{UnexpectedToken} {
    for (one_of) |ty|
        std.debug.assert(ty != got);
    std.debug.print("Expected ", .{});
    for (0.., one_of) |i, ty| {
        std.debug.print("{s}{s}", .{ if (i == 0) "" else " or ", @tagName(ty) });
    }
    std.debug.print(". Got {s}.\n", .{@tagName(got)});
    return error.UnexpectedToken;
}

fn mustEat(lexer: *Lexer, ty: Token.Type) !Token {
    const token = lexer.next();
    if (token.type != ty) {
        std.debug.print("Expected {s}. Got {s} at {}.\n", .{ @tagName(ty), @tagName(token.type), token.loc });
        return error.UnexpectedToken;
    }
    return token;
}

fn allocate(T: type, allocator: Allocator, t: T) !*T {
    const res = try allocator.create(T);
    res.* = t;
    return res;
}