parser.lua

local SyntaxError = require("error").Syntax
local lexer = require("lexer").lexer
local x = require("libraries.inspect").label
local util = require("util")

x = function() end

local function make(tag, line, props)
    props = props or {}
    props.tag = tag
    props.line = line
    return props
end

local Parser = {}
Parser.__index = Parser

function Parser.new(lex, path)
    return setmetatable({ lex = lex, prev = {}, curr = {}, path = path, imports = {} }, Parser)
end

function Parser:eof()
    return self.curr.tag == "EOF"
end

function Parser:check(tag)
    return self.curr.tag == tag
end

function Parser:check_any(...)
    local tags = {...}
    for _, tag in ipairs(tags) do
        if self.curr.tag == tag then
            return true
        end
    end
    return false
end

function Parser:match(tag)
    if self:check(tag) then
        self:advance()
        return true
    end
    return false
end

function Parser:match_any(...)
    if self:check_any(...) then
        self:advance()
        return true
    end
    return false
end

function Parser:advance()
    self.prev, self.curr = self.curr, self.lex()
    return self.prev
end

function Parser:consume(tag, msg)
    if self.curr.tag == tag then
        return self:advance()
    else
        SyntaxError(msg or "expected " .. tag, self.curr.line or "EOF", self.path)
    end
end

-- Parse the root of the ast.
function Parser:chunk()
    local stmts = {}
    while not self:eof() do
        table.insert(stmts, self:stmt())
    end
    return make("Chunk", stmts[1] and stmts[1].line or 1, { body = stmts })
end

-- Don't consume here.
function Parser:stmt()
    if self:eof() then error("stmt() EOF") end
    local tok = self.curr
    local t = tok.tag

    if     t == "let"       then return self:var_decl()
    elseif t == "mut"       then return self:var_decl(true)
    elseif t == "pub"       then return self:pub()
    elseif t == "fun"       then return self:fn_decl()
    elseif t == "pure"      then return self:fn_decl(true)
    elseif t == "if"        then return self:if_()
    elseif t == "while"     then return self:while_()
    elseif t == "for"       then return self:for_()
    elseif t == "return"    then return self:return_()
    elseif t == "repeat"    then return self:repeat_()
    elseif t == "break"     then return make("Break", self:advance().line)
    elseif t == "class"     then return self:class()
    elseif t == "struct"    then return self:struct()
    elseif t == "match"     then return self:match_()
    elseif t == "interface" then return self:interface()
    elseif t == "enum"      then return self:enum()
    elseif t == "do"        then return self:do_()
    elseif t == "use"       then return self:use()
    elseif t == "@"         then return self:comptime()
    else                    return self:expr_stmt()
    end
end

-- Call or assignment.
function Parser:expr_stmt()
    local tok = self:consume("id", "expect name for expression statement prefix")
    local expr = make("Name", tok.line, { name = tok.val })
    expr = self:complete_postfix(expr)

    if expr.tag == "Call" or expr.tag == "MethodCall" then
        return make("ExprStmt", expr.line, { expr = expr })
    elseif self:match("=") then
        local rhs = self:expr()
        if expr.tag == "Index" then
            expr.kind = "assign"  -- A hack.
            return make("IndexAssign", expr.line, { lhs = expr, rhs = rhs })
        elseif expr.tag == "Member" then
            expr.kind = "assign"  -- A hack.
            return make("SetProp", expr.line, { lhs = expr, rhs = rhs })
        else
            return make("Assign", expr.line, { lhs = expr, rhs = rhs })
        end
    end
    SyntaxError("invalid expr stmt", self.curr.line, self.path)
end

function Parser:complete_postfix(expr)
    while true do
        if self:match("(") then
            local args = {}
            if not self:check(")") then
                repeat
                    table.insert(args, self:expr())
                until not self:match(",")
            end
            self:consume(")", "expect ')' after args")
            expr = make("Call", expr.line, { fn = expr, args = args })
        elseif self:match(".") then
            local name = self:consume("id", "expect field name").val
            expr = make("Member", expr.line, { base = expr, name = name })
        elseif self:match(":") then
            local name = self:consume("id", "expect method name").val
            self:consume("(", "expect '(' after method name")
            local args = {}
            if not self:check(")") then
                repeat
                    table.insert(args, self:expr())
                until not self:match(",")
            end
            self:consume(")", "expect ')' after method args")
            expr = make("MethodCall", expr.line, { base = expr, name = name, args = args })
        elseif self:match("[") then
            local start = self:expr()
            local stop
            if self:match(":") then
                stop = self:expr()
            end
            self:consume("]", "expect ']' after indexing")
            expr = make("Index", expr.line, { base = expr, index = start, stop = stop })
        else
            break
        end
    end

    -- TODO: This is a hack. Think again later.
    if self:match("is") then
        local type = self:type()
        expr = make("Is", expr.line, { left = expr, right = type })
    elseif self:match("isnt") then
        local type = self:type()
        expr = make("Is", expr.line, { left = expr, right = type, negate = true })
    end

    return expr
end

function Parser:param_body()
    self:consume("(", "expect '(' after function name")
    local params = {}
    if not self:check(")") then
        repeat
            local pname = self:consume("id", "expect param name").val
            self:consume(":", "expect ':' before type")
            local ptype = self:type()
            table.insert(params, { name = pname, type = ptype })
        until not self:match(",")
    end
    self:consume(")", "expect ')' after param list")
    self:consume("->", "expect '->' before type")
    local ret_type = self:type()

    local body = {}
    while not self:check_any("end", "EOF") do
        table.insert(body, self:stmt())
    end
    self:consume("end", "expect 'end' after function body")
    return params, ret_type, body
end

function Parser:fn_decl(pure, pub)
    local line = self:advance().line
    local name = self:consume("id", "missing function name").val
    if self:match("<") then
        local tparams = {}
        repeat
            local var = self:consume("id", "expect type variable").val
            table.insert(tparams, var)
        until not self:match(",") or self:eof()
        self:consume(">", "expect ']' after type param list")
        local params, ret_type, body = self:param_body()
        return make("Generic", line, { name = name, tparams = tparams, params = params, ret_type = ret_type, body = body, pure = pure, pub = pub })
    else
        local params, ret_type, body = self:param_body()
        return make("FunctionDecl", line, { name = name, params = params, ret_type = ret_type, body = body, pure = pure, pub = pub })
    end
end

function Parser:pub()
    self:advance()
    if self:check("fun") then
        return self:fn_decl(nil, true)
    elseif self:check("pure") then
        return self:fn_decl(true, true)
    elseif self:check("let") then
        return self:var_decl(nil, true)
    elseif self:check("class") then
        return self:class(true)
    elseif self:check("interface") then
        return self:interface(true)
    else
        SyntaxError("invalid pub placement", self.curr.line, self.path)
    end
end

-- Can only declare one variable.
function Parser:var_decl(mut, pub)
    local line = self:advance().line
    if self:match("(") then
        local els = {self:expr()}
        while self:match(",") do
            table.insert(els, self:expr())
        end
        self:consume(")", "expect ')' after tuple")
        self:consume("=", "expect '=' when unpacking")
        local init = self:expr()
        local tuple = make("Tuple", line, {els = els})
        return make("Unpack", line, { tuple = tuple, init = init, mut = mut })
    else
        local name = self:consume("id", "expect var name").val
        local type
        if self:match(":") then
            type = self:type()
        end
        local init
        if self:match("=") then
            init = self:expr()
        end
        return make("NameDecl", line, { name = name, type = type, init = init, mut = mut, pub = pub })
    end
end

function Parser:struct()
    local line = self:advance().line
    local name = self:consume("id", "expect var name").val

    local fname = self:consume("id", "expect struct field name").val
    self:consume(":", "expect ':' before type")
    local ftype = self:type()
    local fields = { {name = fname, type = ftype} }  -- NOTE: Struct cannot be empty.
    while self:match(",") do
        if self:check("end") then break end  -- It means we got trailing comma.
        local name = self:consume("id", "expect struct field name").val
        self:consume(":", "expect ':' before type")
        local type = self:type()
        table.insert(fields, { name = name, type = type })
    end

    self:consume("end", "expect 'end' after struct declaration")
    return make("StructDecl", line, { name = name, fields = fields })
end

function Parser:class(pub)
    local line = self:advance().line
    local name = self:consume("id", "expect class name").val

    -- Parse data section.
    self:consume("data", "expect 'data' after class name")
    local data = {}
    repeat
        local name = self:consume("id", "expect field name").val
        self:consume(":", "expect ':' after field name")
        local type = self:type()
        table.insert(data, { name = name, type = type })
    until self:match("end")

    -- Parse constructor 'new'.
    local tok = self:consume("id", "expect method name")
    local params, ret_type, body = self:param_body()
    local new = make("New", tok.line, { name = tok.val, params = params, ret_type = ret_type, body = body })

    local methods = {}
    while not self:match("end") do
        local tok = self:consume("id", "expect method name")
        local params, ret_type, body = self:param_body()
        local method = make("Method", tok.line, { name = tok.val, params = params, ret_type = ret_type, body = body })
        table.insert(methods, method)
    end

    return make("ClassDecl", line, { name = name, data = data, new = new, methods = methods, pub = pub })
end

function Parser:interface(pub)
    local line = self:advance().line
    local name = self:consume("id", "expect interface name").val

    local field = {}
    while not self:match("end") do
        local name = self:consume("id", "expect field name").val
        self:consume(":", "expect ':' after field name")
        local type = self:type()
        table.insert(field, { name = name, type = type })
    end

    return make("Interface", line, { name = name, fields = field, pub = pub })
end

function Parser:enum()
    local line = self:advance().line
    local name = self:consume("id", "expect enum name").val
    local values = { self:consume("id", "expect enum value").val }  -- NOTE: Enum cannot be empty.
    while self:match(",") do
        if self:check("end") then break end  -- It means we got trailing comma.
        local value = self:consume("id", "expect enum value").val
        table.insert(values, value)
    end
    self:consume("end", "expect 'end' after enum values")
    return make("EnumDecl", line, { name = name, values = values })
end

function Parser:match_()
    local line = self:advance().line
    local expr = self:expr()
    local cases = {}
    while self:match("|") do
        local expr = self:expr()
        self:consume("->", "expect '->' after | expr")
        local body = {}
        while not self:check_any("end", "|", "EOF") do
            table.insert(body, self:stmt())
        end
        table.insert(cases, { expr = expr, body = body })
    end
    self:consume("end", "expect 'end' after case statement")
    return make("Match", line, { expr = expr, cases = cases })
end

function Parser:return_()
    local line = self:advance().line
    local exprs = {}
    if not self:check_any("end", "elseif", "else", "EOF", "semicolon") then
        repeat
            table.insert(exprs, self:expr())
        until not self:check(",")
    end
    return make("Return", line, { exprs = exprs })
end

function Parser:if_()
    local line = self:advance().line
    local cond = self:expr()
    self:consume("then", "expect 'then' after if condition")

    local then_ = {}
    while not self:check_any("end", "elseif", "else", "EOF") do
        table.insert(then_, self:stmt())
    end

    local elseifs_ = nil
    if self:match("elseif") then
        elseifs_ = {}
        repeat
            local cond = self:expr()
            self:consume("then", "expect 'then' after elseif condition")
            local body = {}
            while not self:check_any("end", "elseif", "else", "EOF") do
                table.insert(body, self:stmt())
            end
            table.insert(elseifs_, { cond = cond, body = body })
        until not self:match("elseif")
    end

    local else_ = nil
    if self:match("else") then
        else_ = {}
        while not self:check_any("end", "EOF") do
            table.insert(else_, self:stmt())
        end
    end
    self:consume("end", "expect 'end' after if body")

    return make("If", line, {cond = cond, then_ = then_, elseifs_ = elseifs_, else_ = else_})
end

function Parser:comptime()
    local line = self:advance().line
    local name = self:consume("id", "expect name for comptime function").val

    self:consume("(", "expect '(' after function name")
    local args = {}
    if not self:check(")") then
        repeat
            table.insert(args, self:expr())
        until not self:match(",")
    end
    self:consume(")", "expect ')' after args")

    return make("Comptime", line, { name = name, args = args })
end

function Parser:use()
    local line = self:advance().line
    local usetype
    if self:match("type") then
        usetype = true
    end
    local names = {}
    repeat
        local name = self:consume("id", "expect id after use").val
        table.insert(names, name)
    until not self:match(",")
    self:consume("from", "expect 'from' after use")
    local source = self:consume("Str", "expect source string after from").val

    local target = util.canonicalize(self.path, source)
    self.imports[target] = parse_module(target)

    return make("Use", line, { names = names, source = source, usetype = usetype })
end

function Parser:while_()
    local line = self:advance().line
    local cond = self:expr()
    local body = self:do_()
    return make("While", line, {cond = cond, body = body})
end

function Parser:repeat_()
    local line = self:advance().line
    local body = {}
    while not self:check_any("until", "EOF") do
        table.insert(body, self:stmt())
    end
    self:consume("until", "expect 'until' after repeat body")
    local cond = self:expr()
    return make("Repeat", line, {cond = cond, body = body})
end

function Parser:for_()
    local line = self:advance().line
    local name = self:consume("id", "expect var name").val
    if self:match("=") then
        local start = self:expr()
        self:consume(",", "expect ',' after start expr")
        local end_ = self:expr()
        local step
        if self:match(",") then
            step = self:expr()
        end
        local body = self:do_()
        return make("ForNumeric", line, {var = name, start = start, end_ = end_, step = step, body = body})
    elseif self:match(",") then
        local names = {name}
        repeat
            local name = self:consume("id", "expect var name").val
            table.insert(names, name)
        until not self:match(",")
        self:consume("in", "expect 'in' after var names")
        local iter = self:expr()
        local body = self:do_()
        return make("ForGeneric", line, {vars = names, iter = iter, body = body})
    else
        error("invalid for syntax")
    end
end

function Parser:do_()
    local tok = self:consume("do", "expect 'do'")
    local stmts = {}
    while not self:check_any("end", "EOF") do
        table.insert(stmts, self:stmt())
    end
    self:consume("end", "expect 'end' after do body")
    return make("Block", tok.line, { stmts = stmts })
end

function Parser:type()
    return self:union()
end

function Parser:union()
    local left = self:prim()

    while self:match("|") do
        local right = self:prim()
        left = { tag = "union", left = left, right = right }
    end

    return left
end

-- TODO: Rename this. Because this is not only for primitives.
function Parser:prim()
    if self:match("id") then
        local name = self.prev.val
        if name == "Obj" then
            return { tag = "obj" }
        end
        if self:match("?") then
            return { tag = "nullable", name = name }
        end
        return name
    elseif self:match("Self") then
        return { tag = "Self"}
    elseif self:match("(") then
        local params = {}
        if not self:check(")") then
            repeat
                local p = self:type()
                table.insert(params, p)
            until not self:match(",") or self:eof()
        end
        self:consume(")", "expect ')' after function params")
        self:consume("->", "expect '->' after function params")
        local ret_t = self:type()
        return { tag = "function", params = params, ret_t = ret_t }
    elseif self:match("Nil") then
        return "Nil"
    elseif self:match("{") then
        local t = self:type()

        if self:match(":") then
            local value_t = self:type()
            self:consume("}", "expect '}' after map type")
            return { tag = "map", key_t = t, value_t = value_t }
        elseif self:match(",") then
            local els = {t}
            repeat
                table.insert(els, self:type())
            until not self:match_any(",", "EOF")
            self:consume("}", "expect '}' after tuple type")
            return { tag = "tuple", els = els }
        else
            self:consume("}", "expect '}' after array type")
            return { tag = "array", type = t }
        end
    end
end


function Parser:expr()
    return self:or_()
end

function Parser:or_()
    local left = self:and_()

    while self:match("or") do
        local op = self.prev.tag
        local right = self:and_()
        left = make("Binary", left.line, {left = left, op = op, right = right})
    end

    return left
end

function Parser:and_()
    local left = self:equality()

    while self:match("and") do
        local op = self.prev.tag
        local right = self:equality()
        left = make("Binary", left.line, {left = left, op = op, right = right})
    end

    return left
end

function Parser:equality()
    local left = self:comparison()

    while self:match_any("==", "~=") do
        local op = self.prev.tag
        local right = self:comparison()
        left = make("Binary", left.line, {left = left, op = op, right = right})
    end

    return left
end

function Parser:comparison()
    local left = self:term()

    while self:match_any(">", ">=", "<", "<=", "in") do
        local op = self.prev.tag
        local right = self:term()
        left = make("Binary", left.line, {left = left, op = op, right = right})
    end

    return left
end

function Parser:term()
    local left = self:factor()

    while self:match_any("-", "+", "..") do
        local op = self.prev.tag
        local right = self:factor()
        left = make("Binary", left.line, {left = left, op = op, right = right})
    end

    return left
end

function Parser:factor()
    local left = self:unary()

    while self:match_any("/", "*", "%") do
        local op = self.prev.tag
        local right = self:unary()
        left = make("Binary", left.line, {left = left, op = op, right = right})
    end

    return left
end

function Parser:unary()
    if self:match_any("not", "-", "#") then
        local op = self.prev.tag
        local right = self:unary()
        return make("Unary", right.line, {op = op, right = right})
    end

    return self:postfix()
end

function Parser:postfix()
    local expr = self:primary()
    return self:complete_postfix(expr)
end

function Parser:primary()
    if self:eof() then error("primary() EOF") end
    local tok = self.curr
    local t = tok.tag

    if t == "Int" then
        self:advance()
        return make("Int", tok.line, {val = tok.val})
    elseif t == "Double" then
        self:advance()
        return make("Double", tok.line, {val = tok.val})
    elseif t == "Str" then
        self:advance()
        return make("Str", tok.line, {val = tok.val})
    elseif t == "id" then
        self:advance()
        return make("Name", tok.line, {name = tok.val})
    elseif t == "false" or t == "true" then
        self:advance()
        return make("Bool", tok.line, {val = t})  -- TODO: Should the val be boolean?
    elseif t == "Nil" then
        self:advance()
        return make("Nil", tok.line, {})
    elseif t == "_" then
        self:advance()
        return make("Anon", tok.line, {})
    elseif t == "Self" then
        self:advance()
        return make("Self", tok.line, {})
    elseif t == "(" then
        self:advance()
        local expr = self:expr()
        if self:match(")") then
            return make("Group", tok.line, {expr = expr})
        elseif self:match(",") then
            local els = {expr}
            while not self:match(")") do
                table.insert(els, self:expr())
            end
            return make("Tuple", tok.line, {els = els})
        end
    elseif t == "{" then
        self:advance()
        local is_array = true

        local items = {}
        if not self:check("}") then
            repeat
                local element = self:expr()
                -- TODO: Fix later.
                if self:match("=") then
                    is_array = false
                    local value = self:expr()
                    element = {key = element, value = value}
                end
                table.insert(items, element)
            until not self:match(",")
        end

        self:consume("}", "expect '}' after array literal")
        if is_array then
            return make("Array", tok.line, {items = items})
        else
            return make("Map", tok.line, {items = items})
        end
    elseif t == "^" then
        self:advance()
        local name
        local isself
        if self:check("id") then
            name = self:advance().val
        elseif self:match("Self") then
            isself = true
        end
        local fields = {}
        if self:match("{") then
            if not self:check("}") then
                repeat
                    local name = self:consume("id", "expect var name").val
                    self:consume("=", "expect '=' before value")
                    local init = self:expr()
                    table.insert(fields, { name = name, init = init })
                until not self:match(",")
            end
            self:consume("}", "expect '}' after struct constructor")
        end
        if name then
            return make("StructCons", tok.line, { name = name, args = fields })
        elseif isself then
            return make("SelfCons", tok.line, { args = fields })
        else
            return make("ObjCons", tok.line, { args = fields })
        end
    elseif t == "fun" then
        -- Lambda.
        self:advance()
        local params, ret_type, body = self:param_body()
        return make("Lambda", tok.line, { params = params, ret_type = ret_type, body = body })
    end

    SyntaxError("expect expression", tok.line, self.path)
end

function parse_module(path)
    local file = io.open(path, "r")
    if not file then
        error("Error: failed to open file: " .. path)
        return
    end

    local src = file:read("*all")
    file:close()
    -- print(src)

    local ok, parser_result, imports = pcall(parser, src, path)
    if not ok then
        print(parser_result)
        return
    end

    local module = { path = path, src = src, ast = parser_result, imports = imports }
    ModuleRegistry[path] = module
    return module
end

function parser(src, path)
    local lex = lexer(src, path)
    local p = Parser.new(lex, path)
    p:advance()  -- Pump the lexer.
    local ast = p:chunk()
    return ast, p.imports
end

return parser