luacode-generator.lua

--
-- luacode-generator.lua
--
-- This file is part of the LuaJIT Language Toolkit.
--
-- Module to generate the Lua code that corresponds to a given Lua AST Tree.
-- Can be used as an alternative to the bytecode generator.

local operator = require("sci-lang.operator")

local strbyte, strsub = string.byte, string.sub

local LuaReservedKeyword = {['and'] = 1, ['break'] = 2, ['do'] = 3, ['else'] = 4, ['elseif'] = 5, ['end'] = 6, ['false'] = 7, ['for'] = 8, ['function'] = 9, ['goto'] = 10, ['if'] = 11, ['in'] = 12, ['local'] = 13, ['nil'] = 14, ['not'] = 15, ['or'] = 16, ['repeat'] = 17, ['return'] = 18, ['then'] = 19, ['true'] = 20, ['until'] = 21, ['while'] = 22 }

local ASCII_0, ASCII_9 = 48, 57
local ASCII_a, ASCII_z = 97, 122
local ASCII_A, ASCII_Z = 65, 90

local function char_isletter(c)
    local b = strbyte(c)
    if b >= ASCII_a and b <= ASCII_z then
        return true
    elseif b >= ASCII_A and b <= ASCII_Z then
        return true
    else
        return (c == '_')
    end
end

local function char_isdigit(c)
    local b = strbyte(c)
    return b >= ASCII_0 and b <= ASCII_9
end

local function replace_cc(c)
    local esc = {
        ['\a'] = [[\a]], ['\b'] = [[\b]], ['\f'] = [[\f]], ['\n'] = [[\n]], ['\r'] = [[\r]], ['\t'] = [[\t]], ['\v'] = [[\v]]
    }
    return esc[c] and esc[c] or ('\\' .. string.format("%d", string.byte(c)))
end

local function escape(s)
    s = string.gsub(s, "[\"\\]", "\\%1")
    return string.gsub(s, "%c", replace_cc)
end

local StatementRule = { }
local ExpressionRule = { }

local concat = table.concat
local format = string.format

local function is_string(node)
    return node.kind == "Literal" and type(node.value) == "string"
end

local function is_const(node, val)
    return node.kind == "Literal" and node.value == val
end

local function is_literal(node)
    local k = node.kind
    return (k == "Literal" or k == "Table")
end

local function string_is_ident(str)
    local c = strsub(str, 1, 1)
    if c == '' or not char_isletter(c) then
        return false
    end
    for k = 2, #str do
        c = strsub(str, k, k)
        if not char_isletter(c) and not char_isdigit(c) then
            return false
        end
    end
    return not LuaReservedKeyword[str]
end

local function comma_sep_list(ls, f)
    local strls
    if f then
        strls = { }
        for k = 1, #ls do strls[k] = f(ls[k]) end
    else
        strls = ls
    end
    return concat(strls, ", ")
end

local function as_parameter(node)
    return node.kind == "Vararg" and "..." or node.name
end

function ExpressionRule:Identifier(node)
    return node.name, operator.ident_priority
end

function ExpressionRule:Literal(node)
    local val = node.value
    local str = type(val) == "string" and format("\"%s\"", escape(val)) or tostring(val)
    return str, operator.ident_priority
end

function ExpressionRule:MemberExpression(node)
    local object, prio = self:expr_emit(node.object)
    if prio < operator.ident_priority or is_literal(node.object) then
        object = "(" .. object .. ")"
    end
    local exp
    if node.computed then
        local prop = self:expr_emit(node.property)
        exp = format("%s[%s]", object, prop)
    else
        exp = format("%s.%s", object, node.property.name)
    end
    return exp, operator.ident_priority
end

function ExpressionRule:Vararg()
    return "...", operator.ident_priority
end

function ExpressionRule:ExpressionValue(node)
    return "(" .. self:expr_emit(node.value) .. ")"
end

function ExpressionRule:BinaryExpression(node)
    local oper = node.operator
    local lprio = operator.left_priority(oper)
    local rprio = operator.right_priority(oper)
    local a, alprio, arprio = self:expr_emit(node.left)
    local b, blprio, brprio = self:expr_emit(node.right)
    if not arprio then arprio = alprio end
    if not brprio then brprio = blprio end
    local ap = arprio <  lprio and format("(%s)", a) or a
    local bp = blprio <= rprio and format("(%s)", b) or b
    return format("%s %s %s", ap, oper, bp), lprio, rprio
end

function ExpressionRule:UnaryExpression(node)
    local arg, arg_prio = self:expr_emit(node.argument)
    local op_prio = operator.unary_priority(node.operator)
    if arg_prio < op_prio then arg = format("(%s)", arg) end
    return format("%s %s", node.operator, arg), operator.unary_priority(node.operator)
end

ExpressionRule.LogicalExpression = ExpressionRule.BinaryExpression

function ExpressionRule:ConcatenateExpression(node)
    local ls = { }
    local cat_prio = operator.left_priority("..")
    for k = 1, #node.terms do
        local kprio
        ls[k], kprio = self:expr_emit(node.terms[k])
        if kprio < cat_prio then ls[k] = format("(%s)", ls[k]) end
    end
    return concat(ls, " .. "), cat_prio
end

function ExpressionRule:Table(node)
    local hash = { }
    for i = 1, #node.keyvals do
        local kv = node.keyvals[i]
        local val = self:expr_emit(kv[1])
        local key = kv[2]
        if key then
            if is_string(key) and string_is_ident(key.value) then
                hash[i] = format("%s = %s", key.value, val)
            else
                hash[i] = format("[%s] = %s", self:expr_emit(key), val)
            end
        else
            hash[i] = format("%s", val)
        end
    end
    local content = ""
    if #hash > 0 then
        content = comma_sep_list(hash)
    end
    return "{" .. content .. "}", operator.ident_priority
end

function ExpressionRule:CallExpression(node)
    local callee, prio = self:expr_emit(node.callee)
    if prio < operator.ident_priority then
        callee = "(" .. callee .. ")"
    end
    local exp = format("%s(%s)", callee, self:expr_list(node.arguments))
    return exp, operator.ident_priority
end

function ExpressionRule:SendExpression(node)
    local rec, prio = self:expr_emit(node.receiver)
    if prio < operator.ident_priority or is_literal(node.receiver) then
        rec = "(" .. rec .. ")"
    end
    local method = node.method.name
    local exp = format("%s:%s(%s)", rec, method, self:expr_list(node.arguments))
    return exp, operator.ident_priority
end

function StatementRule:FunctionDeclaration(node)
    self:proto_enter(0)
    local name = self:expr_emit(node.id)
    local header = format("function %s(%s)", name, comma_sep_list(node.params, as_parameter))
    if node.locald then
        header = "local " .. header
    end
    self:add_section(header, node.body)
    local child_proto = self:proto_leave()
    self.proto:merge(child_proto)
end

function ExpressionRule:FunctionExpression(node)
    self:proto_enter()
    local header = format("function(%s)", comma_sep_list(node.params, as_parameter))
    self:add_section(header, node.body)
    local child_proto = self:proto_leave()
    return child_proto:inline(), 0
end

function StatementRule:CallExpression(node)
    local line = self:expr_emit(node)
    self:add_line(line)
end

function StatementRule:ForStatement(node)
    local init = node.init
    local istart = self:expr_emit(init.value)
    local iend = self:expr_emit(node.last)
    local header
    if node.step and not is_const(node.step, 1) then
        local step = self:expr_emit(node.step)
        header = format("for %s = %s, %s, %s do", init.id.name, istart, iend, step)
    else
        header = format("for %s = %s, %s do", init.id.name, istart, iend)
    end
    self:add_section(header, node.body)
end

function StatementRule:ForInStatement(node)
    local vars = comma_sep_list(node.namelist.names, as_parameter)
    local explist = self:expr_list(node.explist)
    local header = format("for %s in %s do", vars, explist)
    self:add_section(header, node.body)
end

function StatementRule:DoStatement(node)
    self:add_section("do", node.body)
end

function StatementRule:WhileStatement(node)
    local test = self:expr_emit(node.test)
    local header = format("while %s do", test)
    self:add_section(header, node.body)
end

function StatementRule:RepeatStatement(node)
    self:add_section("repeat", node.body, true)
    local test = self:expr_emit(node.test)
    local until_line = format("until %s", test)
    self:add_line(until_line)
end

function StatementRule:BreakStatement()
    self:add_line("break")
end

function StatementRule:IfStatement(node)
    local ncons = #node.tests
    for i = 1, ncons do
        local header_tag = i == 1 and "if" or "elseif"
        local test = self:expr_emit(node.tests[i])
        local header = format("%s %s then", header_tag, test)
        self:add_section(header, node.cons[i], true)
    end
    if node.alternate then
        self:add_section("else", node.alternate, true)
    end
    self:add_line("end")
end

function StatementRule:LocalDeclaration(node)
    local line
    local names = comma_sep_list(node.names, as_parameter)
    if #node.expressions > 0 then
        line = format("local %s = %s", names, self:expr_list(node.expressions))
    else
        line = format("local %s", names)
    end
    self:add_line(line)
end

function StatementRule:AssignmentExpression(node)
    local line = format("%s = %s", self:expr_list(node.left), self:expr_list(node.right))
    self:add_line(line)
end

function StatementRule:Chunk(node)
    self:list_emit(node.body)
end

function StatementRule:ExpressionStatement(node)
    local line = self:expr_emit(node.expression)
    self:add_line(line)
end

function StatementRule:ReturnStatement(node)
    local line = format("return %s", self:expr_list(node.arguments))
    self:add_line(line)
end

function StatementRule:LabelStatement(node)
   self:add_line("::" .. node.label .. "::")
end

function StatementRule:GotoStatement(node)
   self:add_line("goto " .. node.label)
end

local function proto_inline(proto)
    -- remove leading whitespaces from first line
    if #proto.code > 0 then
        proto.code[1] = string.gsub(proto.code[1], "^%s*", "")
    end
    return concat(proto.code, "\n")
end

local function proto_merge(proto, child)
    for k = 1, #child.code do
        local line = child.code[k]
        local indent_str = string.rep("    ", proto.indent)
        proto.code[#proto.code + 1] = indent_str .. line
    end
end

local function proto_new(parent, indent)
    local ind = 0
    if indent then
        ind = indent
    elseif parent then
        ind = parent.indent
    end
    local proto = { code = { }, indent = ind, parent = parent }
    proto.inline = proto_inline
    proto.merge = proto_merge
    return proto
end

local function generate(tree, name)

    local self = { line = 0 }
    self.proto = proto_new()
    self.chunkname = tree.chunkname

    function self:proto_enter(indent)
        self.proto = proto_new(self.proto, indent)
    end

    function self:proto_leave()
        local proto = self.proto
        self.proto = proto.parent
        return proto
    end

    local function to_expr(node)
        return self:expr_emit(node)
    end

    function self:compile_code()
        return concat(self.code, "\n")
    end

    function self:indent_more()
        local proto = self.proto
        proto.indent = proto.indent + 1
    end

    function self:indent_less()
        local proto = self.proto
        proto.indent = proto.indent - 1
    end

    function self:line(line)
        -- FIXME: ignored for the moment
    end

    function self:add_line(line)
        local proto = self.proto
        local indent = string.rep("    ", proto.indent)
        proto.code[#proto.code + 1] = indent .. line
    end

    function self:add_section(header, body, omit_end)
        self:add_line(header)
        self:indent_more()
        self:list_emit(body)
        self:indent_less()
        if not omit_end then
            self:add_line("end")
        end
    end

    function self:expr_emit(node)
        local rule = ExpressionRule[node.kind]
        if not rule then error("cannot find an expression rule for " .. node.kind) end
        return rule(self, node)
    end

    function self:expr_list(exps)
        return comma_sep_list(exps, to_expr)
    end

    function self:emit(node)
        local rule = StatementRule[node.kind]
          if not rule then error("cannot find a statement rule for " .. node.kind) end
          rule(self, node)
          if node.line then self:line(node.line) end
    end

    function self:list_emit(node_list)
        for i = 1, #node_list do
            self:emit(node_list[i])
        end
    end

    self:emit(tree)

    return self:proto_leave():inline()
end

return generate