types.lua

local TypeError = require("error").Type
local util = require("util")
local inspect = require("libraries.inspect").inspect
local x = require("libraries.inspect").label

local x = function() end

-- Primitives.
local BOOL   = { tag = "prim", name = "Bool" }
local DOUBLE = { tag = "prim", name = "Double" }
local INT    = { tag = "prim", name = "Int" }
local NIL    = { tag = "prim", name = "Nil" }
local STR    = { tag = "prim", name = "Str" }
local TOP    = { tag = "prim", name = "Top" }


local function Array(elem_t)
    return { tag = "array", elem_t = elem_t }
end

local function Class(name, data, new, methods)
    return { tag = "class", name = name, data = data, new = new, methods = methods }
end

local function Enum(name, values)
    return { tag = "enum", name = name, values = values }
end

local function Fn(params, ret_t, is_pure)
    return { tag = "fn", params = params, ret_t = ret_t, is_pure = is_pure or false }
end

local function Generic(tparams, params, ret_t)
    return { tag = "generic", tparams = tparams, params = params, ret_t = ret_t }
end

local function Interface(name, fields, more)
    return { tag = "interface", name = name, fields = fields, more = more or {} }
end

local function Map(key_t, value_t)
    return { tag = "map", key_t = key_t, value_t = value_t }
end

local function Method(params, ret_t)
    return { tag = "method", params = params, ret_t = ret_t }
end

local function New(params, ret_t)
    return { tag = "new", params = params, ret_t = ret_t }
end

local function Obj(fields)
    return { tag = "obj", fields = fields or {} }
end

local function Struct(name, fields, methods)
    return { tag = "struct", name = name, fields = fields, methods = methods }
end

local function Tuple(els, lualike)
    return { tag = "tuple", els = els, lualike = lualike }
end

local function Union(parts)
    return { tag = "union", parts = parts }
end

local function Var(name, link)
    return { tag = "var", name = name, link = link }
end


-- NOTE: These three funcitons are helpers. Not actual types.
local function uninit(t)
    return { tag = "uninit", type = t }
end

local function mut(t)
    return { tag = "mut", type = t }
end

local function unmut(t)
    return t.tag == "mut" and t.type or t
end


local function equal(a, b)
    if a == b then
        return true

    elseif a.tag == "mut" then
        return equal(a.type, b)

    elseif a.tag == "mut" then
        return equal(a, b.type)

    elseif a.tag ~= b.tag then
        return false

    elseif a.tag == "prim" then
        return a.name == b.name

    elseif a.tag == "array" then
        return equal(a.elem_t, b.elem_t)

    elseif a.tag == "tuple" then
        for i, t in ipairs(a.els) do
            if not equal(t, b.els[i]) then
                return false
            end
        end
        return true

    elseif a.tag == "fn" and b.tag == "fn" then
        if #a.params ~= #b.params then
            return false
        end
        for i = 1, #a.params do
            if not equal(a.params[i], b.params[i]) then
                return false
            end
        end
        return equal(a.ret_t, b.ret_t)

    else
        error("equal: unhandled " .. inspect(a) .. inspect(b))
    end
end

local function can_assign(a, b)
    -- TODO: Handle mut elsewhere.
    if a.tag == "mut" then
        return can_assign(a.type, b)

    elseif b.tag == "mut" then
        return can_assign(a, b.type)

    elseif equal(b, TOP) then
        return not equal(a, NIL)

    elseif a.tag == "prim" and b.tag == "prim" then
        if a.name == "Int" and b.name == "Double" then
            return true
        end
        return a.name == b.name

    elseif a.tag == "array" and b.tag == "array" then
        return can_assign(a.elem_t, b.elem_t)

    elseif a.tag == "tuple" and b.tag == "tuple" then
        for i, t in ipairs(a.els) do
            if not can_assign(t, b.els[i]) then
                return false
            end
        end
        return true

    elseif a.tag == "map" and b.tag == "map" then
        return can_assign(a.key_t, b.key_t) and can_assign(a.value_t, b.value_t)

    elseif a.tag == "enum" and b.tag == "enum" then
        return a.name == b.name

    elseif a.tag == "struct" and b.tag == "struct" then
        return a.name == b.name

    elseif a.tag == "class" and b.tag == "class" then
        return a.name == b.name

    elseif a.tag == "union" then
        if b.tag ~= "union" then return false end
        for _, t in ipairs(a.parts) do
            if not can_assign(t, b) then
                return false
            end
        end
        return true

    elseif b.tag == "union" then
        for _, t in ipairs(b.parts) do
            if can_assign(a, t) then
                return true
            end
        end
        return false

    elseif a.tag == "fn" and b.tag == "fn" then
        if #a.params ~= #b.params then
            return false
        end
        for i = 1, #a.params do
            can_assign(a.params[i], b.params[i])
        end
        return can_assign(a.ret_t, b.ret_t)

    elseif a.tag == "named" then
        return b.name and b.name == a.name

    elseif b.tag == "named" then
        return a.name and b.name == a.name

    -- TODO: Fix this.
    elseif b.tag == "var" then
        if not b.link then
            if not a.link then
                return a.name == b.name
            end
            b.link = a
            return true
        else
            return can_assign(b.link, a)
        end

    elseif a.tag == "var" then
        return can_assign(b, a)

    elseif a.tag == "obj" and b.tag == "obj" then
        return true

    -- TODO: This only handles init of empty array of interface, like let a: {IStuff} = {}.
    elseif a.tag == "interface" and b.tag == "interface" then
        return a.name == b.name

    else
        return false
    end
end

local function show(t)
    local tag = t.tag

    if tag == "prim" then
        return t.name

    elseif tag == "enum" then
        return t.name .. " enum"

    elseif tag == "struct" then
        return t.name

    elseif tag == "array" then
        return "{" .. show(t.elem_t) .. "}"

    elseif tag == "tuple" then
        local els = {}
        for i, t in ipairs(t.els) do
            els[i] = show(t)
        end
        return "{" .. table.concat(els, ", ") .. "}"

    elseif tag == "var" then
        if t.link then
            return show(t.link)
        else
            return "var " .. t.name
        end

    elseif tag == "union" then
        local parts = {}
        for _, p in ipairs(t.parts) do
            table.insert(parts, show(p))
        end
        return table.concat(parts, " | ")

    elseif tag == "fn" then
        local params = {}
        for _, p in ipairs(t.params) do
            table.insert(params, show(p))
        end
        return "(" .. table.concat(params, ", ") .. ") -> " .. show(t.ret_t)

    elseif tag == "obj" then
        return "Obj"

    elseif tag == "class" then
        return "class " .. t.name

    -- TODO: There should be function show(a, b) to handle below case. It need to show
    --       what methods are not implemented.
    elseif tag == "interface" then
        local f = {}
        local keys = {}
        for k in pairs(t.fields) do table.insert(keys, k) end
        table.sort(keys)
        for _, k in ipairs(keys) do
            f[#f+1] = "  " .. k .. ": " .. show(t.fields[k])
        end

        local keys = {}
        for k in pairs(t.more) do table.insert(keys, k) end
        table.sort(keys)
        for _, k in ipairs(keys) do
            f[#f+1] = "  " .. k .. ": " .. show(t.more[k])
        end
        return "interface " .. t.name .. " {\n" .. table.concat(f, "\n") .. "\n}"

    -- TODO: Remove this.
    elseif tag == "mut" then
        return "mut " .. show(t.type)

    else
        error("show: unhandled " .. inspect(t))
    end
end

-- NOTE: Don't inline this. This will support flattening and deduplication later.
local function union_of(parts)
    if #parts == 1 then
        return parts[1]
    end
    return Union(parts)
end

local function intersect_with_pred(t, pred_t, line)
    if t.tag == "union" then
        local keep = {}
        for _, p in ipairs(t.parts) do
            if can_assign(p, pred_t) then
                table.insert(keep, p)
            end
        end
        if #keep == 0 then
            TypeError("type intersection is empty", line, path)  -- TODO: path is undefined, moved back to typecheck.lua
        end
        return union_of(keep)
    elseif t.tag == "interface" then
        local name = pred_t.val
        if t.fields[name] then
            return t
        elseif t.more[name] then
            local o = util.copy_map(t)
            local fields = util.copy_map(t.fields)
            local more = util.copy_map(t.more)
            fields[name] = t.more[name]
            more[name] = nil
            o.fields = fields
            o.more = more
            return o
        else
            TypeError("no field named " .. name .. " on interface " .. t.name, line, path)  -- TODO: path is undefined, moved back to typecheck.lua
        end
    else
        if can_assign(t, pred_t) then
            return t
        else
            TypeError("type intersection is empty", line, path)  -- TODO: path is undefined, moved back to typecheck.lua
        end
    end
end

local function subtract_pred(t, pred_t, line)
    if t.tag == "union" then
        local keep = {}
        for _, p in ipairs(t.parts) do
            if not can_assign(p, pred_t) then
                table.insert(keep, p)
            end
        end
        -- NOTE: it is impossile to have an empty keep.
        return union_of(keep)
    elseif t.tag == "interface" then
        local name = pred_t.val
        if t.fields[name] then
            local o = util.copy_map(t)
            local fields = util.copy_map(t.fields)
            fields[name] = nil
            o.fields = fields
            return o
        elseif t.more[name] then
            local o = util.copy_map(t)
            local more = util.copy_map(t.more)
            more[name] = nil
            o.more = more
            return o
        else
            return nil
        end
    else
        if can_assign(t, pred_t) then
            -- NOTE: This is handled by the caller. This is different from intersect_with_pred().
            -- It is only error if the false branch exist.
            return nil
        else
            return t
        end
    end
end

local function refine_on_pred(tenv, pred, line, names)
    local pred_t = pred.type
    local sub = pred.subject
    local name
    local base
    local cur

    if pred.tag == "in" then
        name = sub.name
        cur = tenv:get(name)
    else
        if sub.tag == "Name" then
            name = sub.name
            cur = tenv:get(name)
        elseif sub.tag == "Member" then
            name = sub.base.name
            base = tenv:get(name)
            base = unmut(base)
            cur = base.fields[sub.name]
        else
            error("refine_on_pred: unhandled tag " .. sub.tag)
        end
    end
    table.insert(names, name)

    local trueT, falseT
    if cur.tag == "mut" then
        cur = cur.type
        trueT = mut(intersect_with_pred(cur, pred_t), line)
        falseT = mut(subtract_pred(cur, pred_t), line)
    else
        trueT = intersect_with_pred(cur, pred_t, line)
        falseT = subtract_pred(cur, pred_t, line)
    end

    local ttrue = tenv:copy()
    if pred.tag == "in" then
        ttrue:set(name, trueT)
    elseif sub.tag == "Name" then
        ttrue:set(name, trueT)
    elseif sub.tag == "Member" then
        local t = util.copy_map(base)
        t.fields = util.copy_map(base.fields)
        t.more = util.copy_map(base.more)
        t.fields[sub.name] = trueT
        ttrue:set(name, t)
    end

    if falseT == nil then
        return ttrue, nil
    else
        local tfalse = tenv:copy()
        if pred.tag == "in" then
            tfalse:set(name, falseT)
        elseif sub.tag == "Name" then
            tfalse:set(name, falseT)
        elseif sub.tag == "Member" then
            local t = util.copy_map(base)
            t.fields = util.copy_map(base.fields)
            t.more = util.copy_map(base.more)
            t.fields[sub.name] = falseT
            tfalse:set(name, t)
        end
        return ttrue, tfalse
    end
end

local function refine(tenv, pred, line, names)
    if pred == nil then
        return tenv:copy(), tenv:copy()
    end

    local t = pred.tag
    if t == "isnil" or t == "is" then
        return refine_on_pred(tenv, pred, line, names)
    elseif t == "in" then
        return refine_on_pred(tenv, pred, line, names)
    elseif t == "not" then
        local ttrue, tfalse = refine_on_pred(tenv, pred.sub, line, names)
        return tfalse, ttrue
    else
        error("refine: unhandled pred tag " .. t)
    end
end


return {
    BOOL   = BOOL,
    DOUBLE = DOUBLE,
    INT    = INT,
    NIL    = NIL,
    STR    = STR,
    TOP    = TOP,

    Array     = Array,
    Class     = Class,
    Enum      = Enum,
    Fn        = Fn,
    Generic   = Generic,
    Interface = Interface,
    Map       = Map,
    Method    = Method,
    New       = New,
    Obj       = Obj,
    Struct    = Struct,
    Tuple     = Tuple,
    Union     = Union,
    Var       = Var,

    uninit = uninit,
    mut    = mut,
    unmut  = unmut,

    can_assign = can_assign,
    equal      = equal,
    show       = show,
    refine     = refine,
}