DCPU-16 Assembler

-- DCPU-16 Assembler in Lua, by KillaVanilla

-- If it's an assembler, why is it called DCPUCompiler?

-- Usage:
-- DCPUCompiler <input> <output> <symbols> <arguments>

-- Arguments are denoted by "-" or "--", and MUST come after the input/output/symbolFile arguments.

local args = {...}

if #args < 3 then
    print("Usage: ")
    print(fs.getName(shell.getRunningProgram()).." <input> <output> <symbols> <arguments>")
    print("Required parameters:")
    print("<input> : Input assembly file.")
    print("<output> : Output binary file. ")
    print("<symbols> : Symbols file, useful for debugging. Pass as \"-\" to disable.")
    print("Arguments:")
    print("These must come after the aforementioned parameters.")
    print("-p / --print : Save all output to \"assemblyLog\".")
    print("-a / --ascii / --ascii-out / --text-out / --text : Output as a list of DAT instructions instead of as a binary.")
    print("-d / --debug : Enable debugging output (values emitted, parser output, etc.).")
    print("Do note that this compiler may still have a few bugs in it; if you find any, please report them ASAP, ok?")
    return
end

local flags = {}
for i=4, #args do
    if string.sub(args[i], 1, 2) == "--" then
        flags[string.sub(args[i],3)] = true
    elseif string.sub(args[i], 1, 1) == "-" then
        flags[string.sub(args[i],2)] = true
    end
end

if flags["p"] or flags["print"] then
    local file = fs.open("assemblyLog", "w")
    file.close()
end

-- ooooobbbbbaaaaaa (MSB-0 DCPU Instruction)
-- aaaaaabbbbbooooo (LSB-0 DCPU Instruction)

-- aaaaaaooooo00000 (LSB-0 DCPU Special Instruction)
-- 00000oooooaaaaaa (MSB-0 DCPU Special Instruction)

-- A few tables to help with parsing:

local operandValues = {
    ["A"] = 0,
    ["B"] = 0x01,
    ["C"] = 0x02,
    ["X"] = 0x03,
    ["Y"] = 0x04,
    ["Z"] = 0x05,
    ["I"] = 0x06,
    ["J"] = 0x07,
    ["[A]"] = 0x08,
    ["[B]"] = 0x09,
    ["[C]"] = 0x0A,
    ["[X]"] = 0x0B,
    ["[Y]"] = 0x0C,
    ["[Z]"] = 0x0D,
    ["[I]"] = 0x0E,
    ["[J]"] = 0x0F,
    --[[["[A+]"] = 0x10,
    ["[B+]"] = 0x11,
    ["[C+]"] = 0x12,
    ["[X+]"] = 0x13,
    ["[Y+]"] = 0x14,
    ["[Z+]"] = 0x15,
    ["[I+]"] = 0x16,
    ["[J+]"] = 0x17,]]
    ["PUSH"] = 0x18,
    ["POP"] = 0x18,
    ["PEEK"] = 0x19,
    ["[SP]"] = 0x19,
    ["PICK"] = 0x1A,
    --["[SP+]"] = 0x1A,
    ["SP"] = 0x1B,
    ["PC"] = 0x1C,
    ["EX"] = 0x1D,
    --["[]"] = 0x1E, -- next word (pointer)
    --[""] = 0x1F, -- next word (literal)
    ["-1"] = 0x20,
    ["0"] = 0x21,
    ["1"] = 0x22,
    ["2"] = 0x23,
    ["3"] = 0x24,
    ["4"] = 0x25,
    ["5"] = 0x26,
    ["6"] = 0x27,
    ["7"] = 0x28,
    ["8"] = 0x29,
    ["9"] = 0x2A,
    ["10"] = 0x2B,
    ["11"] = 0x2C,
    ["12"] = 0x2D,
    ["13"] = 0x2E,
    ["14"] = 0x2F,
    ["15"] = 0x30,
    ["16"] = 0x31,
    ["17"] = 0x32,
    ["18"] = 0x33,
    ["19"] = 0x34,
    ["20"] = 0x35,
    ["21"] = 0x36,
    ["22"] = 0x37,
    ["23"] = 0x38,
    ["24"] = 0x39,
    ["25"] = 0x3A,
    ["26"] = 0x3B,
    ["27"] = 0x3C,
    ["28"] = 0x3D,
    ["29"] = 0x3E,
    ["30"] = 0x3F,
}

local opcodes = {
    ["SET"] = 0x01,
    ["ADD"] = 0x02,
    ["SUB"] = 0x03,
    ["MUL"] = 0x04,
    ["MLI"] = 0x05,
    ["DIV"] = 0x06,
    ["DVI"] = 0x07,
    ["MOD"] = 0x08,
    ["MDI"] = 0x09,
    ["AND"] = 0x0A,
    ["BOR"] = 0x0B,
    ["XOR"] = 0x0C,
    ["SHR"] = 0x0D,
    ["ASR"] = 0x0E,
    ["SHL"] = 0x0F,
    ["IFB"] = 0x10,
    ["IFC"] = 0x11,
    ["IFE"] = 0x12,
    ["IFN"] = 0x13,
    ["IFG"] = 0x14,
    ["IFA"] = 0x15,
    ["IFL"] = 0x16,
    ["IFU"] = 0x17,
    ["ADX"] = 0x1A,
    ["SBX"] = 0x1B,
    ["STI"] = 0x1E,
    ["STD"] = 0x1F,
}

local spec_opcodes = {
    ["JSR"] = 0x01,
    ["INT"] = 0x08,
    ["IAG"] = 0x09,
    ["IAS"] = 0x0A,
    ["RFI"] = 0x0B,
    ["IAQ"] = 0x0C,
    ["HWN"] = 0x10,
    ["HWQ"] = 0x11,
    ["HWI"] = 0x12,
}

local literalOperandTransTable = { -- Translates the a-value only literal operands into the literals they represent and vice versa.
    [0xFFFF] = 0x20,
    [0] = 0x21,
    [1] = 0x22,
    [2] = 0x23,
    [3] = 0x24,
    [4] = 0x25,
    [5] = 0x26,
    [6] = 0x27,
    [7] = 0x28,
    [8] = 0x29,
    [9] = 0x2A,
    [10] = 0x2B,
    [11] = 0x2C,
    [12] = 0x2D,
    [13] = 0x2E,
    [14] = 0x2F,
    [15] = 0x30,
    [16] = 0x31,
    [17] = 0x32,
    [18] = 0x33,
    [19] = 0x34,
    [20] = 0x35,
    [21] = 0x36,
    [22] = 0x37,
    [23] = 0x38,
    [24] = 0x39,
    [25] = 0x3A,
    [26] = 0x3B,
    [27] = 0x3C,
    [28] = 0x3D,
    [29] = 0x3E,
    [30] = 0x3F,
    [0x20] = 0xFFFF,
    [0x21] = 0,
    [0x22] = 1,
    [0x23] = 2,
    [0x24] = 3,
    [0x25] = 4,
    [0x26] = 5,
    [0x27] = 6,
    [0x28] = 7,
    [0x29] = 8,
    [0x2A] = 9,
    [0x2B] = 10,
    [0x2C] = 11,
    [0x2D] = 12,
    [0x2E] = 13,
    [0x2F] = 14,
    [0x30] = 15,
    [0x31] = 16,
    [0x32] = 17,
    [0x33] = 18,
    [0x34] = 19,
    [0x35] = 20,
    [0x36] = 21,
    [0x37] = 22,
    [0x38] = 23,
    [0x39] = 24,
    [0x3A] = 25,
    [0x3B] = 26,
    [0x3C] = 27,
    [0x3D] = 28,
    [0x3E] = 29,
    [0x3F] = 30,
}

local function compileMsg(line, severity, text)
    -- Severity:
    -- 0 -- Debug (light grey)
    -- 1 -- Info (white)
    -- 2 -- Warning (yellow; unused)
    -- 3 -- Error (red; stops compile)
    local file = {}
    if flags["p"] or flags["print"] then
        file = fs.open("assemblyLog", "a")
    end
    term.setBackgroundColor(colors.black)
    if severity == 0 then
        term.setTextColor(colors.lightGray)
        if flags["d"] or flags["debug"] then
            print("[DEBUG] "..text)
            if flags["p"] or flags["print"] then
                file.writeLine("[DEBUG] "..text)
            end
        end
    elseif severity == 1 then
        term.setTextColor(colors.white)
        print("[INFO] "..text)
        if flags["p"] or flags["print"] then
            file.writeLine("[INFO] "..text)
        end
    --[[elseif severity == 2 then
        term.setTextColor(colors.yellow)
        print("[WARN] "..text)
        if flags["p"] or flags["print"] then
            file.writeLine("[WARN] "..text)
        end]]
    elseif severity == 3 then
        term.setTextColor(colors.red)
        print("[ERR] "..text)
        if flags["p"] or flags["print"] then
            file.writeLine("[ERR] "..text)
            file.close()
        end
        error("Compile failed at line "..line..".", 0)
    end
    if flags["p"] or flags["print"] then
        file.close()
    end
end

local function parseLine(line, line_num, labels) -- Parse a single lime of DCPU-16 assembly, generating 3 values: the opcode w/ operands, and the two words referred to by the "next word" operands.
    local ins, b, a = string.match(line, "(%a%a%a)%s+([%w_%-%[%]%+]+)%,?%s*([%w_%-%[%]%+]*)")
    compileMsg(line_num, 0, "Line Parser Input: "..line)
    if not ins then
        return
    end
    local ptr_b_reg, ptr_b_lit = string.match(b, "%[(%a+)%+(%w+)%]")
    local ptr_a_reg, ptr_a_lit = string.match(a, "%[(%a+)%+(%w+)%]")
    if operandValues[b] then -- Does B match a recognized operand?
        b = operandValues[b]
        if b >= 0x20 then
            ptr_b_lit = literalOperandTransTable[b]
            b = 0x1F
        end
    elseif labels[b] then -- Is B a label?
        compileMsg(line_num, 0, "B is a label: "..b)
        b = 0x1F
        ptr_b_lit = labels[b][1]
    elseif ptr_b_reg then -- Did we find a [register+offset] combo?
        ptr_b_reg = string.upper(ptr_b_reg)
        if ptr_b_reg == "A" then
            b = 0x10
        elseif ptr_b_reg == "B" then
            b = 0x11
        elseif ptr_b_reg == "C" then
            b = 0x12
        elseif ptr_b_reg == "X" then
            b = 0x13
        elseif ptr_b_reg == "Y" then
            b = 0x14
        elseif ptr_b_reg == "Z" then
            b = 0x15
        elseif ptr_b_reg == "I" then
            b = 0x16
        elseif ptr_b_reg == "J" then
            b = 0x17
        else
            compileMsg(line_num, 3, "Unrecognized register: "..ptr_b_reg)
        end
    elseif string.sub(b, 1, 1) == "[" and string.sub(b, #b, #b) == "]" then -- Is this a reference?
        b = 0x1E
        if tonumber(string.sub(b, 2, #b-1)) then
            ptr_b_lit = tonumber(string.sub(b, 2, #b-1))
        elseif labels[ string.sub(a, 2, #a-1) ] then
            ptr_a_lit = labels[ string.sub(a, 2, #a-1) ][1]
        end
    elseif tonumber(b) then -- Is it a literal?
        ptr_b_lit = tonumber(b)
        b = 0x1F
    else
        compileMsg(line_num, 3, "Unrecognized expression: "..b)
    end
    if a and a ~= "" then
        if operandValues[a] then
            a = operandValues[a]
        elseif labels[a] then
            compileMsg(line_num, 0, "A is a label: "..a)
            if labels[a][1] <= 30 then
                a = literalOperandTransTable[labels[a][1]]
            else
                a = 0x1F
                ptr_a_lit = labels[a][1]
            end
        elseif ptr_a_reg then
            ptr_a_reg = string.upper(ptr_a_reg)
            if ptr_a_reg == "A" then
                a = 0x10
            elseif ptr_a_reg == "B" then
                a = 0x11
            elseif ptr_a_reg == "C" then
                a = 0x12
            elseif ptr_a_reg == "X" then
                a = 0x13
            elseif ptr_a_reg == "Y" then
                a = 0x14
            elseif ptr_a_reg == "Z" then
                a = 0x15
            elseif ptr_a_reg == "I" then
                a = 0x16
            elseif ptr_a_reg == "J" then
                a = 0x17
            elseif ptr_a_reg == "SP" then
                a = 0x1A
            else
                compileMsg(line_num, 3, "Unrecognized register: "..ptr_a_reg)
            end
        elseif string.sub(a, 1, 1) == "[" and string.sub(a, #a, #a) == "]" then
            if tonumber(string.sub(a, 2, #a-1)) then
                ptr_a_lit = tonumber(string.sub(a, 2, #a-1))
            elseif labels[ string.sub(a, 2, #a-1) ] then
                ptr_a_lit = labels[ string.sub(a, 2, #a-1) ][1]
            else
                compileMsg(line_num, 3, "Unrecognized label: "..string.sub(a, 2, #a-1))
            end
            a = 0x1E
        elseif tonumber(a) then
            ptr_a_lit = tonumber(a)
            a = 0x1F
        else
            compileMsg(line_num, 3, "Unrecognized label: "..string.sub(a, 2, #a-1))
        end
    end
    -- Opcode construction:
    local ret = 0
    if opcodes[ins] then -- Use the instruction value as a "base":
        compileMsg(line_num, 0, "Instruction is normal: 0x"..string.format("%X", opcodes[ins]))
        ret = opcodes[ins]
    elseif spec_opcodes[ins] then
        compileMsg(line_num, 0, "Instruction is special: 0x"..string.format("%X", spec_opcodes[ins]))
        -- Shift the instruction value to the middle:
        ret = bit.blshift(spec_opcodes[ins], 5)
    end
    compileMsg(line_num, 0, "B-value: 0x"..string.format("%X", b))
    if a and a ~= "" then
        compileMsg(line_num, 0, "A-value: 0x"..string.format("%X", a))
    end
    -- Add the b-value:
    if opcodes[ins] then
        ret = bit.bor(ret, bit.blshift(b, 5))
    elseif spec_opcodes[ins] then
        ret = bit.bor(ret, bit.blshift(b, 10))
    end

    if opcodes[ins] and (a and a ~= "") then
        -- Add the a-value:
        ret = bit.bor(ret, bit.blshift(a, 10))
    end

    local emit = "Emit: 0x"..string.format("%X", ret)
    if ptr_a_lit then
        emit = emit..", 0x"..string.format("%X", ptr_a_lit)
    end
    if ptr_b_lit then
        emit = emit..", 0x"..string.format("%X", ptr_b_lit)
    end
    compileMsg(line_num, 0, emit)

    return ret, ptr_a_lit, ptr_b_lit
end

local function determineInstructionSize(line)
    local size = 1
    if string.find(line, "DAT ") then
        local dS, dE = string.find(line, "DAT ")
        local data_str = string.sub(line, dE+1)
        for dat in string.gmatch(data_str, "([^%,]+)%,?") do
            dat = string.match(dat, "%s*(%S+)")
            if tonumber(dat) then
                size = size+1
            elseif string.sub(dat, 1,1) == "\"" then
                for i=2, #dat-1 do
                    size = size+1
                end
            end
        end
    else
        local ins, b, a = string.match(line, "(%a%a%a)%s+([%w_%-%[%]%+]+)%,?%s*([%w_%-%[%]%+]*)")
        if not operandValues[b] then
            size = size+1
        else
            if operandValues[b] >= 0x20 then
                size = size+1
            end
        end
        if a and a ~= "" then
            if not operandValues[a] then
                size = size+1
            end
        end
    end
    return size
end

local function generateMacroTable(lines)
    local macros = {}
    local currentMacro = {}
    local inMacro = false
    local cMacroName = ""
    for i, line in ipairs(lines) do -- Label Accumulation
        local no_whtspace = string.match(line, "%s*(.+)")
        if string.sub(no_whtspace, 1, 7) == ".macro " then
            if not inMacro then
                cMacroName = string.sub(no_whtspace, 8)
                inMacro = true
            end
        elseif string.sub(no_whtspace, 1, 9) == ".endmacro" then
            if inMacro then
                macros[cMacroName] = currentMacro
                currentMacro = {}
                cMacroName = ""
                inMacro = false
            end
        elseif inMacro then
            table.insert(currentMacro, line)
        end
    end
    return macros
end

local function readLines(file)
    local lines = {}
    local iHandle = io.open(file, "r")
    for line in iHandle:lines() do
        table.insert(lines, line)
    end
    iHandle:close()
    return lines
end

local function generateSymTable(lines)
    local labels = {}
    local addr = 0
    local sTime = os.clock()
    for i, line in ipairs(lines) do -- Label Accumulation
        if string.match(line, "%:(.+)") then
            local label = string.match(line, "%:(.+)")
            labels[label] = {addr, i}
            compileMsg(0, 0, "Label found: "..label.." at address 0x"..string.format("%X", addr)..".")
        elseif string.match(line, "(.+)%:") then
            local label = string.match(line, "(.+)%:")
            labels[label] = {addr, i}
            compileMsg(0, 0, "Label found: "..label.." at address 0x"..string.format("%X", addr)..".")
        end
        if string.match(line, "(%a%a%a)%s+([%w_%-%[%]%+]+)%,?%s*([%w_%-%[%]%+]*)") then
            local sz = determineInstructionSize(line)
            if sz then
                compileMsg(i, 0, "0x"..string.format("%X", addr)..": "..line.." (size: "..sz..")")
                addr = addr+sz
            else
                compileMsg(i, 3, "Syntax error: "..line)
            end
        end
    end
    return labels, generateMacroTable(lines)
end

local function doSymbolSub(lines, labels, macros)
    local i = 1
    while true do
        if not lines[i] then
            break
        end
        local no_whtspace = string.match(lines[i], "%s*(.+)")
        if string.sub(no_whtspace, 1, 9) == ".insmacro" then
            local macroName = string.sub(no_whtspace, 10)
            if macros[macroName] then
                for i2=#macros[macroName], 1, -1 do
                    table.insert(lines, i2, macros[macroName][i])
                end
            else
                compileMsg(i, 3, "Undefined macro reference: "..macroName)
            end
        end
        local ins, b, a = string.match(lines[i], "(%a%a%a)%s+([%w_%-%[%]%+]+)%,?%s*([%w_%-%[%]%+]*)")
        if ins then
            for index, label in pairs(labels) do
                if b == index then
                    b = label[1]
                end
                if a == index then
                    a = label[1]
                end
                --lines[i] = string.gsub(lines[i], index, label[1]) -- this is the old replacer code
            end
            local repLine = ins.." "..b
            if a and a ~= "" then
                repLine = repLine..", "..a
            end
            if no_whtspace ~= repLine then
                compileMsg(i, 0, "Replacing line "..i.." (\""..no_whtspace.."\") with \""..repLine.."\".")
                lines[i] = repLine
            end
        end
        i=i+1
    end
    return lines
end

local function parseSingleFile(inFile, outFile, labelsFile)
    local mem = {}
    local lines = readLines(inFile)

    compileMsg(0, 1, "Step 1: Symbol Resolution")
    local labels, macros = generateSymTable(lines)
    compileMsg(0, 1, "Symbol resolution complete.")
    compileMsg(0, 1, "Step 2a: Symbol substitution -- Pass 1...")
    lines = doSymbolSub(lines, labels, macros)
    compileMsg(0, 1, "Symbol substitution complete.")
    compileMsg(0, 1, "Step 2b: Label Adjustment...")
    labels, macros = generateSymTable(lines)
    lines = readLines(inFile)
    compileMsg(0, 1, "Label adjustment complete.")
    compileMsg(0, 1, "Step 2c: Symbol substitution -- Pass 2...")
    lines = doSymbolSub(lines, labels, macros)
    compileMsg(0, 1, "Symbol substitution complete.")
    if outFile and (flags["o"] or flags["optimize-only"]) then
        local oHandle = fs.open(outFile, "w")
        for i, line in ipairs(lines) do
            oHandle.writeLine(line)
        end
        oHandle.close()
        return
    end
    compileMsg(0, 1, "Step 3: Final Assembly / Output Generation...")
    for i, line in ipairs(lines) do
        if string.find(line, "DAT ") then
            local dS, dE = string.find(line, "DAT ")
            local data_str = string.sub(line, dE+1)
            for dat in string.gmatch(data_str, "([^%,]+)%,?") do
                dat = string.match(dat, "%s*(%S+)")
                if tonumber(dat) then
                    table.insert(mem, tonumber(dat))
                elseif string.sub(dat, 1,1) == "\"" then
                    for i=2, #dat-1 do
                        table.insert(mem, string.byte(string.sub(dat,i,i)))
                    end
                end
            end
        else
            local at_pc, at_pc_1, at_pc_2 = parseLine(line, i, labels)
            if at_pc then
                table.insert(mem, at_pc)
            end
            if at_pc_1 then
                table.insert(mem, at_pc_1)
            end
            if at_pc_2 then
                table.insert(mem, at_pc_2)
            end
        end
    end
    compileMsg(0, 1, "Assembly complete; "..(#mem*2).." bytes generated.")
    if outFile then
        if flags["a"] or flags["ascii-out"] or flags["ascii"] or flags["text"] or flags["text-out"] then
            local oHandle = fs.open(outFile, "w")
            for index, byte in ipairs(mem) do
                oHandle.writeLine("DAT 0x"..string.format("%X", byte))
            end
            oHandle.close()
        else
            local oHandle = fs.open(outFile, "wb")
            for index, byte in ipairs(mem) do
                oHandle.write(byte)
            end
            oHandle.close()
        end
    end
    if labelsFile and labelsFile ~= "-" then
        local lfHandle = fs.open(labelsFile, "w")
        lfHandle.writeLine("-- BEGIN LABEL DEFINITIONS --")
        lfHandle.writeLine("")
        for i,v in pairs(labels) do
            lfHandle.writeLine(i..": 0x"..string.format("%X", v[1]).." - "..v[2])
        end
        lfHandle.writeLine("-- END LABEL DEFINITIONS --")
        lfHandle.writeLine("")
        lfHandle.writeLine("-- BEGIN MACRO DEFINITIONS --")
        lfHandle.writeLine("")
        for i,v in pairs(macros) do
            lfHandle.writeLine(i..":")
            for i2, line in ipairs(v) do
                lfHandle.writeLine(line)
            end
            lfHandle.writeLine("")
        end
        lfHandle.writeLine("-- END MACRO DEFINITIONS --")
        lfHandle.close()
    end
    return mem
end
parseSingleFile(args[1], args[2], args[3])