transform.lua

local x = require("libraries.inspect").label

-- x = function() end

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

local ilua_to_lua = {
    Int = "number",
    Double = "number",
    Str = "string",
    Bool = "boolean",
    Nil = "nil",
}

local lovemap = {
    ["newFont_default"] = function(expr)
        local aux = {}
        local fn = k(expr.fn, aux)
        fn.name = "newFont"

        local args = {}
        for _, e in ipairs(expr.args) do
            local arg = k(e, aux)
            table.insert(args, arg)
        end
        return nil, make("Call", { fn = fn, args = args })
    end
}

local function counter()
    local n = 0
    return function() n = n + 1; return "_t" .. n end
end

local fresh = counter()

-- TODO: Handle lambda callback and unnamed array.
function inline_map(expr)
    local arr_name = expr.args[2].name
    local fn_name = expr.args[1].name
    local tmp = fresh()
    local s = {}
    s[1] = make("NameDecl", { name = tmp, init = {tag = "Array", items = {}} })
    s[2] = make("ForNumeric", {
        var = "i",
        start = {tag = "Int", val = 1},
        end_ = {tag = "Unary", op = "#", right = {tag = "Name", name = arr_name}},
        body = {tag = "Block", stmts = { {
            lhs = {
                tag = "ArrayIndex",
                base = {tag = "Name", name = tmp},
                index = {tag = "Name", name = "i"},
            },
            rhs = {
                args = { {
                    tag = "ArrayIndex",
                    base = {tag = "Name", name = arr_name},
                    index = {tag = "Name", name = "i"},
                } },
                fn = {tag = "Name", name = fn_name},
                tag = "Call"
            },
            tag = "IndexAssign"
        } }}
    })

    return s, make("Name", {name = tmp})
end

-- Leave the line number.
local function shallow_copy(t)
    local o = {}
    for k, v in pairs(t) do
        if k ~= "line" then
            o[k] = v
        end
    end
    return o
end

function tf_expr(expr)
    local t = expr.tag

    -- TODO: Is copying necessary?
    if t == "Int" or t == "Double" or t == "Str" or t == "Nil" or t == "Bool" or t == "Name" then
        return nil, shallow_copy(expr)

    elseif t == "Self" then
        return nil, {tag = "Name", name = expr.ref}

    elseif t == "Group" then
        local l_s, l_e = tf_expr(expr.expr)
        return nil, make("Group", {expr = l_e})

    elseif t == "Lambda" then
        local params = {}
        for _, p in ipairs(expr.params) do
            table.insert(params, p.name)
        end

        local body = {}
        for _, s in ipairs(expr.body) do
            local ts = tf_stmt(s)
            for _, x in ipairs(ts) do
                table.insert(body, x)
            end
        end
        return nil, make("Lambda", { params = params, body = body })

    elseif t == "Map" then
        local items = {}
        for _, item in ipairs(expr.items) do
            local k_s, k_e = tf_expr(item.key)
            local v_s, v_e = tf_expr(item.value)
            table.insert(items, {key = k_e, value = v_e})
        end
        return nil, make("Map", {items = items})

    elseif t == "Array" then
        local items = {}
        for _, el in ipairs(expr.items) do
            local l_s, l_e = tf_expr(el)
            table.insert(items, l_e)
        end
        return nil, make("Array", {items = items})

    elseif t == "Tuple" then
        local els = {}
        for _, el in ipairs(expr.els) do
            local l_s, l_e = tf_expr(el)
            table.insert(els, l_e)
        end
        return nil, make("Tuple", {els = els})

    elseif t == "SelfCons" then
        local fields = {}
        for _, arg in ipairs(expr.args) do
            local _, init = tf_expr(arg.init)
            table.insert(fields, { name = arg.name, init = init })
        end
        return nil, make("SelfCons", { ref = expr.ref, args = fields })

    elseif t == "StructCons" then
        local fields = {}
        for _, arg in ipairs(expr.args) do
            local _, init = tf_expr(arg.init)
            table.insert(fields, { name = arg.name, init = init })
        end
        return nil, make("StructCons", { name = expr.name, args = fields })

    elseif t == "ObjCons" then
        local fields = {}
        for _, arg in ipairs(expr.args) do
            local _, init = tf_expr(arg.init)
            table.insert(fields, { name = arg.name, init = init })
        end
        return nil, make("ObjCons", { args = fields })

    -- Desugar s[i] to string.sub(s, i, i)
    elseif t == "Index" then
        local b_s, b_e = tf_expr(expr.base)
        local i_s, i_e = tf_expr(expr.index)
        local s_s, s_e
        if expr.stop then
            s_s, s_e = tf_expr(expr.stop)
        end

        if expr.index_kind == "string" then
            local expr = make("Member", { base = {tag = "Name", name = "string"}, name = "sub" })
            local args
            if s_e then
                args = {b_e, i_e, s_e}
            else
                args = {b_e, i_e, i_e}
            end
            return nil, make("Call", { fn = expr, args = args })
        elseif expr.index_kind == "array" then
            if s_e then
                return nil, make("ArrayIndex", { base = b_e, index = i_e, stop = s_e })
            else
                return nil, make("ArrayIndex", { base = b_e, index = i_e })
            end
        elseif expr.index_kind == "tuple" then
            return nil, make("ArrayIndex", { base = b_e, index = i_e })
        elseif expr.index_kind == "map" then
            return nil, make("MapIndex", { base = b_e, index = i_e })
        end

    elseif t == "MethodCall" then
        local _, r_e = tf_expr(expr.base)
        return nil, make("MethodCall", { base = r_e, name = expr.name, args = expr.args })

    elseif t == "Member" then
        -- TODO: Handle l_s
        local r_s, r_e = tf_expr(expr.base)
        if expr.enumno then
            return nil, make("Int", { val = expr.enumno })
        else
            return nil, make("Member", { base = r_e, name = expr.name })
        end

    -- TODO: Clean up.
    elseif t == "Call" then
        local n = expr.fn.name

        if expr.fn.tag == "Name" and expr.fn.name == "push" then
            local fn = make("Member", { base = {tag = "Name", name = "table"}, name = "insert" })
            local expr = make("Call", { fn = fn, args = expr.args })
            return nil, expr
        elseif expr.fn.name == "map" then
            return inline_map(expr)
        elseif lovemap[n] then
            return lovemap[n](expr)
        elseif expr.kind == "constructor" then
            local aux = {}
            local fn = k(expr.fn, aux)

            local args = {}
            for _, e in ipairs(expr.args) do
                local arg = k(e, aux)
                table.insert(args, arg)
            end
            local member = make("Member", { base = fn, name = "new" })
            return nil, make("Call", { fn = member, args = args })
        else
            local aux = {}
            local fn = k(expr.fn, aux)

            local args = {}
            for _, e in ipairs(expr.args) do
                local arg = k(e, aux)
                table.insert(args, arg)
            end
            return nil, make("Call", { fn = fn, args = args })
        end

    elseif t == "Unary" then
        local r_s, r_e = tf_expr(expr.right)
        return nil, make("Unary", {op = expr.op, right = r_e})

    elseif t == "Is" then
        local pred = expr.pred
        if pred.tag == "not" then
            pred = pred.sub
        end
        if pred.type.tag ~= "prim" then
            error("only support primitive predicate for now")
        end
        local t = ilua_to_lua[pred.type.name]
        local lhs, rhs
        if t == "nil" then
            lhs = pred.subject
            rhs = {tag = "Nil"}
        else
            lhs = make("Call", { fn = {tag = "Name", name = "type"}, args = {pred.subject} })
            rhs = {tag = "Str", val = t}
        end
        local op = expr.negate and "~=" or "=="
        return nil, make("Binary", {left = lhs, op = op, right = rhs})

    elseif t == "Binary" then
        local l_s, l_e = tf_expr(expr.left)
        local r_s, r_e = tf_expr(expr.right)
        return nil, make("Binary", {left = l_e, op = expr.op, right = r_e})

    else
        error("expr type not implemented: " .. tostring(t))
    end
end

function k(expr, aux)
    local stmts, expr = tf_expr(expr)
    if stmts then
        for _, stmt in ipairs(stmts) do
            table.insert(aux, stmt)
        end
    end
    return expr
end

function tf_stmt(stmt)
    local t = stmt.tag

    -- NOTE: Currently all declarations doesn't need termporaries.
    if t == "NameDecl" then
        local aux = {}
        local init
        if stmt.init then
            init = k(stmt.init, aux)
        end

        table.insert(aux, make("NameDecl", { name = stmt.name, init = init }))
        return aux

    elseif t == "Unpack" then
        local aux = {}
        local init
        if stmt.init then
            init = k(stmt.init, aux)
        end

        table.insert(aux, make("Unpack", { tuple = stmt.tuple, init = init, lualike = stmt.lualike }))
        return aux

    elseif t == "EnumDecl" then
        -- It is erased.
        return {}

    elseif t == "Interface" then
        -- It is erased.
        return {}

    elseif t == "StructDecl" then
        -- It is erased.
        return {}

    elseif t == "FunctionDecl" then
        -- Erase param types and return type.
        local params = {}
        for _, p in ipairs(stmt.params) do
            table.insert(params, p.name)
        end

        local body = {}
        for _, s in ipairs(stmt.body) do
            local ts = tf_stmt(s)
            for _, x in ipairs(ts) do
                table.insert(body, x)
            end
        end
        return { make("FunctionDecl", { name = stmt.name, params = params, body = body }) }

    -- Similar to FunctionDecl.
    elseif t == "Generic" then
        -- Erase param types and return type.
        local params = {}
        for _, p in ipairs(stmt.params) do
            table.insert(params, p.name)
        end

        local body = {}
        for _, s in ipairs(stmt.body) do
            local ts = tf_stmt(s)
            for _, x in ipairs(ts) do
                table.insert(body, x)
            end
        end
        return { make("Generic", { name = stmt.name, params = params, body = body }) }

    elseif t == "Use" then
        local module = ModuleRegistry[stmt.target]
        local ok, transform_result = pcall(tf, module)
        if not ok then
            error(transform_result)
        end
        module.tf_ast = transform_result

        if stmt.usetype then
            -- Erased.
            return {}
        end

        local dummy = "_m"
        local req = make("Require", { source = stmt.source })
        local decl = make("NameDecl", { name = dummy, init = req })

        local inits = {}
        for i, name in ipairs(stmt.names) do
            table.insert(inits, make("Member", { base = {tag = "Name", name = dummy}, name = name }))
        end
        local a = make("MultiDecl", { names = stmt.names, inits = inits })

        return { decl, a }

    elseif t == "Return" then
        local aux = {}
        local exprs = {}
        for _, e in ipairs(stmt.exprs) do
            local expr = k(e, aux)
            table.insert(exprs, expr)
        end

        table.insert(aux, make("Return", { exprs = exprs }))
        return aux

    elseif t == "While" then
        local aux = {}
        local cond = k(stmt.cond, aux)
        local body = tf_stmt(stmt.body)[1]
        table.insert(aux, make("While", { cond = cond, body = body }))
        return aux

    elseif t == "Break" then
        return { make("Break") }

    -- TODO: Complete later.
    elseif t == "ForNumeric" then
        local body = tf_stmt(stmt.body)[1]
        return { make("ForNumeric", { var = stmt.var, start = stmt.start, end_ = stmt.end_, step = stmt.step, body = body }) }

    elseif t == "ForGeneric" then
        local body = tf_stmt(stmt.body)[1]
        return { make("ForGeneric", {vars = stmt.vars, iter = stmt.iter, body = body, kind = stmt.kind}) }

    elseif t == "Block" then
        local stmts = {}
        for i, stmt in ipairs(stmt.stmts) do
            local ts = tf_stmt(stmt)
            for _, t in ipairs(ts) do
                table.insert(stmts, t)
            end
        end
        return { make("Block", { stmts = stmts }) }

    elseif t == "ClassDecl" then
        local n = stmt.name
        local methods = stmt.methods
        local s = {}
        s[1] = make("NameDecl", { name = stmt.name, init = {tag = "Array", items = {}} })
        s[2] = make("SetProp", { lhs = {tag = "Member", base = {tag = "Name", name = stmt.name}, name = "__index"}, rhs = {tag = "Name", name = stmt.name} })

        -- Constructor
        local params = {}
        for _, p in ipairs(stmt.new.params) do
            table.insert(params, p.name)
        end
        local body = {}
        for _, s in ipairs(stmt.new.body) do
            local ts = tf_stmt(s)
            for _, x in ipairs(ts) do
                table.insert(body, x)
            end
        end
        s[#s+1] = make("Function", { name = {tag = "Member", base = {tag = "Name", name = n}, name = "new"}, params = params, body = body })

        for i, m in ipairs(methods) do
            local params = {}
            for _, p in ipairs(m.params) do
                table.insert(params, p.name)
            end

            local body = {}
            for _, s in ipairs(m.body) do
                local ts = tf_stmt(s)
                for _, x in ipairs(ts) do
                    table.insert(body, x)
                end
            end
            s[#s+1] = make("Function", { name = {tag = "Member", base = {tag = "Name", name = n}, name = m.name}, params = params, body = body })
        end
        return s

    elseif t == "Match" then
        local case = stmt.cases[1]
        local expr_s, expr_e = tf_expr(case.expr)
        local comp_s, comp_e = tf_expr(stmt.expr)
        local cond = make("Binary", {left = comp_e, op = "==", right = expr_e})

        local then_ = {}
        for _, s in ipairs(case.body) do
            local ts = tf_stmt(s)
            for _, t in ipairs(ts) do
                table.insert(then_, t)
            end
        end

        local others
        local elseifs_ = nil
        if #stmt.cases >= 2 then
            elseifs_ = {}
            for i = 2, #stmt.cases do
                local case = stmt.cases[i]
                if case.expr.tag == "Anon" then
                    others = case
                    break
                end
                local expr_s, expr_e = tf_expr(case.expr)
                local cond = make("Binary", {left = comp_e, op = "==", right = expr_e})

                local body = {}
                for _, s in ipairs(case.body) do
                    local ts = tf_stmt(s)
                    for _, t in ipairs(ts) do
                        table.insert(body, t)
                    end
                end
                table.insert(elseifs_, { cond = cond, body = body })
            end
        end

        local else_ = nil
        if others then
            else_ = {}
            for _, s in ipairs(others.body) do
                local ts = tf_stmt(s)
                for _, t in ipairs(ts) do
                    table.insert(else_, t)
                end
            end
        end
        return { make("If", {cond = cond, then_ = then_, elseifs_ = elseifs_, else_ = else_}) }

    elseif t == "If" then
        local aux = {}
        local cond = k(stmt.cond, aux)

        local new_then = {}
        for _, s in ipairs(stmt.then_) do
            local ts = tf_stmt(s)
            for _, t in ipairs(ts) do
                table.insert(new_then, t)
            end
        end

        local elseifs_
        if stmt.elseifs_ then
            elseifs_ = {}
            for _, stmt in ipairs(stmt.elseifs_) do
                local cond = k(stmt.cond, aux)

                local body = {}
                for _, s in ipairs(stmt.body) do
                    local ts = tf_stmt(s)
                    for _, t in ipairs(ts) do
                        table.insert(body, t)
                    end
                end
                table.insert(elseifs_, { cond = cond, body = body })
            end
        end

        local else_
        if stmt.else_ then
            else_ = {}
            for _, s in ipairs(stmt.else_) do
                local ts = tf_stmt(s)
                for _, t in ipairs(ts) do
                    table.insert(else_, t)
                end
            end
        end

        table.insert(aux, make("If", { cond = cond, then_ = new_then, elseifs_ = elseifs_, else_ = else_ }))
        return aux

    elseif t == "SetProp" then
        local aux = {}
        local lhs = k(stmt.lhs, aux)
        local rhs = k(stmt.rhs, aux)

        table.insert(aux, make("SetProp", { lhs = lhs, rhs = rhs }))
        return aux

    elseif t == "Assign" then
        local aux = {}
        local lhs = k(stmt.lhs, aux)
        local rhs = k(stmt.rhs, aux)

        table.insert(aux, make("Assign", { lhs = lhs, rhs = rhs }))
        return aux

    elseif t == "IndexAssign" then
        local aux = {}
        local lhs = k(stmt.lhs, aux)
        local rhs = k(stmt.rhs, aux)

        table.insert(aux, make("IndexAssign", { lhs = lhs, rhs = rhs }))
        return aux

    elseif t == "ExprStmt" then
        local _, expr_e = tf_expr(stmt.expr)
        return { make("ExprStmt", { expr = expr_e }) }

    else
        error("stmt transform not implemented: " .. t)
    end
end

function tf(module)
    module.export._types = nil
    if module.ast.tag ~= "Chunk" then error("Expected Chunk root") end

    local body = {}
    for _, stmt in ipairs(module.ast.body) do
        local transformed = tf_stmt(stmt)
        for _, t in ipairs(transformed) do
            table.insert(body, t)
        end
    end

    local ast = make("Chunk", { body = body })
    return ast
end

return tf