lua/assembler.lua

-- assembler.lua
-- Mini 6502 assembler in Lua — compiles mnemonic source to hex strings
-- at LuaLaTeX load time. Two-pass: labels resolved on second pass.
--
-- Authors: Ahmad Ali Parr, Jessica L. Williams (SNAPKITTYWEST)

local asm = {}

asm.opcodes = {
    -- LDA
    LDA_IMM=0xA9, LDA_ZP=0xA5, LDA_ZPX=0xB5, LDA_ABS=0xAD,
    LDA_ABSX=0xBD, LDA_ABSY=0xB9, LDA_INDX=0xA1, LDA_INDY=0xB1,
    -- STA
    STA_ZP=0x85, STA_ZPX=0x95, STA_ABS=0x8D,
    STA_ABSX=0x9D, STA_ABSY=0x99, STA_INDX=0x81, STA_INDY=0x91,
    -- LDX / LDY
    LDX_IMM=0xA2, LDX_ZP=0xA6, LDX_ABS=0xAE,
    LDY_IMM=0xA0, LDY_ZP=0xA4, LDY_ABS=0xAC,
    STX_ZP=0x86, STX_ABS=0x8E,
    STY_ZP=0x84, STY_ABS=0x8C,
    -- ALU
    ADC_IMM=0x69, ADC_ZP=0x65, ADC_ABS=0x6D,
    SBC_IMM=0xE9, SBC_ZP=0xE5, SBC_ABS=0xED,
    AND_IMM=0x29, AND_ZP=0x25, AND_ABS=0x2D,
    ORA_IMM=0x09, ORA_ZP=0x05, ORA_ABS=0x0D,
    EOR_IMM=0x49, EOR_ZP=0x45, EOR_ABS=0x4D,
    CMP_IMM=0xC9, CMP_ZP=0xC5, CMP_ABS=0xCD,
    BIT_ZP=0x24,  BIT_ABS=0x2C,
    -- Shifts
    ASL_A=0x0A, ASL_ZP=0x06, ASL_ABS=0x0E,
    LSR_A=0x4A, LSR_ZP=0x46, LSR_ABS=0x4E,
    ROL_A=0x2A, ROL_ZP=0x26, ROL_ABS=0x2E,
    ROR_A=0x6A, ROR_ZP=0x66, ROR_ABS=0x6E,
    -- Inc/Dec
    INC_ZP=0xE6, INC_ABS=0xEE, INC_ZPX=0xF6, INC_ABSX=0xFE,
    DEC_ZP=0xC6, DEC_ABS=0xCE, DEC_ZPX=0xD6, DEC_ABSX=0xDE,
    INX=0xE8, DEX=0xCA, INY=0xC8, DEY=0x88,
    -- Transfer
    TAX=0xAA, TXA=0x8A, TAY=0xA8, TYA=0x98, TXS=0x9A, TSX=0xBA,
    -- Stack
    PHA=0x48, PLA=0x68, PHP=0x08, PLP=0x28,
    -- Flags
    CLC=0x18, SEC=0x38, CLI=0x58, SEI=0x78,
    CLV=0xB8, CLD=0xD8, SED=0xF8,
    -- Branch
    BCC=0x90, BCS=0xB0, BEQ=0xF0, BNE=0xD0,
    BMI=0x30, BPL=0x10, BVC=0x50, BVS=0x70,
    -- Jump/Call
    JMP_ABS=0x4C, JMP_IND=0x6C, JSR_ABS=0x20, RTS=0x60,
    -- Misc
    BRK=0x00, NOP=0xEA,
}

function asm:new()
    local o = {
        labels   = {},
        output   = {},
        addr     = 0x0600,
        patches  = {},   -- {idx, label, is_branch} for second pass
    }
    setmetatable(o, { __index = self })
    return o
end

function asm:emit(byte)
    table.insert(self.output, byte & 0xFF)
    self.addr = self.addr + 1
end

function asm:emit16(val)
    self:emit(val & 0xFF)
    self:emit((val >> 8) & 0xFF)
end

function asm:parse_val(s)
    s = s:gsub("%s", "")
    if s:match("^%$%x+$") then return tonumber(s:sub(2), 16) end
    if s:match("^%d+$")   then return tonumber(s) end
    if s:match("^#")      then return self:parse_val(s:sub(2)) end
    return nil, s  -- might be a label
end

function asm:parse_line(line)
    -- Strip comment
    line = line:gsub(";.*$", ""):match("^%s*(.-)%s*$")
    if line == "" then return end

    -- Label definition
    local lbl, rest = line:match("^(%w+):%s*(.*)")
    if lbl then
        self.labels[lbl] = self.addr
        line = rest
        if line == "" then return end
    end

    -- Mnemonic and operand
    local mn, op = line:match("^(%u+)%s*(.*)")
    if not mn then return end
    op = op:match("^%s*(.-)%s*$")

    -- Implied / Accumulator
    if op == "" or op == "A" then
        local code = self.opcodes[mn] or self.opcodes[mn.."_A"]
        if code then self:emit(code) end
        return
    end

    -- Immediate: #$xx or #nn
    if op:match("^#") then
        local val = self:parse_val(op:sub(2))
        local code = self.opcodes[mn.."_IMM"]
        if code and val then self:emit(code); self:emit(val) end
        return
    end

    -- Indirect X: ($xx,X)
    if op:match("^%(.*,X%)$") then
        local inner = op:match("^%((.+),X%)$")
        local val = self:parse_val(inner)
        if val and val < 256 then
            self:emit(self.opcodes[mn.."_INDX"]); self:emit(val)
        end
        return
    end

    -- Indirect Y: ($xx),Y
    if op:match("^%(.*%),Y$") then
        local inner = op:match("^%((.+)%),Y$")
        local val = self:parse_val(inner)
        if val and val < 256 then
            self:emit(self.opcodes[mn.."_INDY"]); self:emit(val)
        end
        return
    end

    -- Absolute,X / ZP,X
    if op:match(",X$") then
        local base = op:gsub(",X$","")
        local val = self:parse_val(base)
        if val then
            if val < 256 then
                self:emit(self.opcodes[mn.."_ZPX"]); self:emit(val)
            else
                self:emit(self.opcodes[mn.."_ABSX"]); self:emit16(val)
            end
        end
        return
    end

    -- Absolute,Y
    if op:match(",Y$") then
        local base = op:gsub(",Y$","")
        local val = self:parse_val(base)
        if val then
            self:emit(self.opcodes[mn.."_ABSY"]); self:emit16(val)
        end
        return
    end

    -- Branch (relative label or offset)
    local branches = {BCC=true,BCS=true,BEQ=true,BNE=true,BMI=true,BPL=true,BVC=true,BVS=true}
    if branches[mn] then
        local code = self.opcodes[mn]
        self:emit(code)
        local val = self:parse_val(op)
        if val then
            self:emit(val & 0xFF)
        else
            -- label reference — patch in second pass
            table.insert(self.patches, {idx=#self.output+1, label=op, branch_from=self.addr+1})
            self:emit(0x00)  -- placeholder
        end
        return
    end

    -- ZP or ABS
    local val = self:parse_val(op)
    if val then
        if val < 256 then
            local code = self.opcodes[mn.."_ZP"]
            if code then self:emit(code); self:emit(val) return end
        end
        local code = self.opcodes[mn.."_ABS"] or self.opcodes[mn.."_ABS"]
        if not code then code = self.opcodes[mn] end
        if code then self:emit(code); self:emit16(val) end
        return
    end

    -- JSR / JMP with label
    if mn == "JSR" or mn == "JMP" then
        local code = (mn == "JSR") and 0x20 or 0x4C
        self:emit(code)
        local target = self.labels[op]
        if target then
            self:emit16(target)
        else
            table.insert(self.patches, {idx=#self.output, label=op, is16=true})
            self:emit(0x00); self:emit(0x00)
        end
    end
end

function asm:second_pass()
    for _, p in ipairs(self.patches) do
        local target = self.labels[p.label]
        if target then
            if p.is16 then
                self.output[p.idx]   = target & 0xFF
                self.output[p.idx+1] = (target >> 8) & 0xFF
            else
                -- branch: signed 8-bit relative offset
                local offset = target - p.branch_from
                if offset < -128 or offset > 127 then
                    if tex then tex.print("\\PackageWarning{assembler}{Branch out of range: "..p.label.."}") end
                    offset = 0
                end
                self.output[p.idx] = offset & 0xFF
            end
        end
    end
end

function asm:assemble(source)
    local inst = self:new()
    for line in source:gmatch("[^\r\n]+") do
        inst:parse_line(line)
    end
    inst:second_pass()
    local parts = {}
    for _, b in ipairs(inst.output) do
        table.insert(parts, string.format("%02X", b))
    end
    return table.concat(parts, " ")
end

return asm