masm

#!/usr/bin/env python

# a simple macro-less bootstrap assembler for dcpu16

import re
import sys

class SyntaxError(Exception):
    def __init__(self, msg):
        self.msg = msg

OPCODES = {}

class Basic(object):
    def __init__(self, name, code):
        OPCODES[name] = self
        self.name = name
        self.code = code
    def make(self, args):
        if len(args) != 2: raise SyntaxError("expected 2 args")
        b = parse_arg(args[0], False)
        a = parse_arg(args[1], True)
        def emit(program):
            program.append((a.value() << 10) + (b.value() << 5) + self.code)
            a.emit(program)
            b.emit(program)
        return emit

class Special(object):
    def __init__(self, name, code):
        OPCODES[name] = self
        self.name = name
        self.code = code
    def make(self, args):
        a = parse_arg(args[0], True)
        if len(args) != 1: raise SyntaxError("expected 1 arg")
        def emit(program):
            program.append((a.value() << 10) + (self.code << 5) + 0x00)
            a.emit(program)
        return emit

class DW(object):
    def __init__(self):
        OPCODES["DW"] = self
        OPCODES["DAT"] = self
    def make(self, args):
        args = [parse_lit(x) for x in args]
        def emit(program):
            for a in args: program.append(a)
        return emit

class DZ(object):
    def __init__(self):
        OPCODES["DZ"] = self
    def make(self, args):
        if len(args) != 1: raise SyntaxError("expected 1 arg")
        arg = args[0]
        if arg[0] != '"' or arg[-1] != '"':
            raise SyntaxError("expected quoted ascii")
        arg = arg[1:-1]
        def emit(program):
            cs = [ord(x) for x in arg]
            if len(cs) % 2 == 1: cs.append(0)
            for pair in zip(*[iter(cs)]*2):
                program.append((pair[0] << 8) + pair[1])
        return emit

class DZW(object):
    def __init__(self):
        OPCODES["DZW"] = self
    def make(self, args):
        if len(args) != 1: raise SyntaxError("expected 1 arg")
        arg = args[0]
        if arg[0] != '"' or arg[-1] != '"':
            raise SyntaxError("expected quoted ascii")
        arg = arg[1:-1]
        def emit(program):
            for c in [ord(x) for x in arg]:
                program.append(c)
        return emit

# basic ops...
Basic("SET", 0x01)
Basic("ADD", 0x02)
Basic("SUB", 0x03)
Basic("MUL", 0x04)
Basic("MLI", 0x05)
Basic("DIV", 0x06)
Basic("DVI", 0x07)
Basic("MOD", 0x08)
Basic("MDI", 0x09)
Basic("AND", 0x0a)
Basic("BOR", 0x0b)
Basic("XOR", 0x0c)
Basic("SHR", 0x0d)
Basic("ASR", 0x0e)
Basic("SHL", 0x0f)
Basic("IFB", 0x10)
Basic("IFC", 0x11)
Basic("IFE", 0x12)
Basic("IFN", 0x13)
Basic("IFG", 0x14)
Basic("IFA", 0x15)
Basic("IFL", 0x16)
Basic("IFU", 0x17)
Basic("ADX", 0x1a)
Basic("SBX", 0x1b)
Basic("STI", 0x1e)
Basic("STD", 0x1f)

# special ops...
Special("JSR", 0x01)
Special("IMG", 0x02) # custom, dump core to file
Special("DIE", 0x03) # custom, make emulator exit
Special("DBG", 0x04) # custom, enter emulator debugger
Special("INT", 0x08)
Special("IAG", 0x09)
Special("IAS", 0x0a)
Special("RFI", 0x0b)
Special("IAQ", 0x0c)
Special("HWN", 0x10)
Special("HWQ", 0x11)
Special("HWI", 0x12)

# pseudo-ops...
DW()
DZ()
DZW()

class Raw(object):
    def __init__(self, code):
        self.code = code
    def value(self):
        return self.code
    def emit(self, prog):
        pass

class Deref(object):
    def __init__(self, raw):
        self.raw = raw
    def value(self):
        return 0x08 + self.raw.value()
    def emit(self, prog):
        pass

class Pick(object):
    # value is either the actual int value, or a string label
    # which can be emitted into the instruction stream and later
    # rewritten...
    def __init__(self, val):
        self.val = val
    def value(self):
        return 0x1a
    def emit(self, prog):
        prog.append(self.val)

class Lit(object):
    # value is either the actual int value, or a string label
    # which can be emitted into the instruction stream and later
    # rewritten...
    def __init__(self, val, deref, inline=False):
        self.val = val
        self.deref = deref
        self.inline = inline
    def value(self):
        if self.deref: return 0x1e
        if self.inline and -1 <= self.val < 31: return 0x21 + self.val
        # it's simpler just to special-case 0xffff rather than to really
        # deal with un/signed 16-bit values in python...
        if self.inline and self.val == 0xffff: return 0x20
        return 0x1f
    def emit(self, prog):
        if self.inline and -1 <= self.val < 31: return
        if self.inline and self.val == 0xffff: return # see above
        prog.append(self.val)

class Offset(object):
    # offset is either the actual int offset value, or a string label
    # which can be emitted into the instruction stream and later
    # rewritten...
    def __init__(self, offset, reg):
        self.offset = offset 
        self.reg = reg
    def value(self):
        return 0x10 + self.reg.value()
    def emit(self, prog):
        prog.append(self.offset)

REGS = {
  "A": Raw(0x0),
  "B": Raw(0x1),
  "C": Raw(0x2),
  "X": Raw(0x3),
  "Y": Raw(0x4),
  "Z": Raw(0x5),
  "I": Raw(0x6),
  "J": Raw(0x7)
}

OPERANDS = dict(REGS.items() + {
  "PEEK": Raw(0x19),
  "SP": Raw(0x1b),
  "PC": Raw(0x1c),
  "EX": Raw(0x1d)
}.items())

PSHP = Raw(0x18)

hex_re = re.compile('0x[0-9A-Fa-f]+')
dec_re = re.compile('-?[0-9]+')

def parse_lit(arg):
    if hex_re.match(arg): return int(arg, 16)
    if dec_re.match(arg): return int(arg, 10)
    return arg

def parse_pair(lit, reg):
    lit = parse_lit(lit)
    try: return Offset(lit, REGS[reg.upper()])
    except KeyError: raise SyntaxError('bad register: ' + reg)

def parse_arg(arg, a):
    if len(arg) == 0: raise SyntaxError("empty arg")
    if arg[0] == '[':
        if arg[-1] != ']': raise SyntaxError("mismatched brackets")
        arg = arg[1:-1].strip() # trim brackets
        if '+' in arg: return parse_pair(*[x.strip() for x in arg.split('+')])
        try: return Deref(OPERANDS[arg.upper()])
        except KeyError: return Lit(parse_lit(arg), True)
    if arg.upper() in OPERANDS: return OPERANDS[arg.upper()]
    if arg.upper() == 'POP':
        if a: return PSHP
        else: raise SyntaxError("'pop' is allowed only in 'a' position")
    if arg.upper() == 'PUSH':
        if not a: return PSHP
        else: raise SyntaxError("'push' is allowed only in 'b' position")
    if arg.upper().startswith('PICK'):
        return Pick(parse_lit(arg[4:].strip()))
    return Lit(parse_lit(arg), False, a)

# this is a terrible hack, but it was so tempting...
# note that we special-case ";" as 'not a comment' so that we can
# easily define the forth word ';'...
line_re = re.compile('(\s*([A-Za-z0-9_.]+):)?\s*((";"|[^;])+)?(;.*)?')

def parse_line(line):
    _, label, instr, _, _ = line_re.match(line).groups()
    if label and label.upper() in OPERANDS:
        raise SyntaxError("illegal label: " + label)
    if instr:
        op, _, args = instr.partition(' ')
        args = [x.strip() for x in args.split(',')] if args else []
        try: return label, OPCODES[op.upper()].make(args)
        except KeyError: raise SyntaxError("invalid opcode: " + op)
    return label, None


def main():
    if len(sys.argv) == 3:
        input_filename = sys.argv[1]
        output_filename = sys.argv[2]
    else:
        print "usage: masm <input.asm> <output.obj>"
        sys.exit(1)

    program = []
    labels = {}
    with open(input_filename) as f:
        for lineno, line in enumerate(f):
            try:
                label, op = parse_line(line)
                if label is not None: labels[label] = len(program)
                if op is not None: op(program)
            except SyntaxError as e: 
                print ("Syntax error on line %d: %s" % (lineno+1, e.msg))
                sys.exit(1)

    with open(output_filename, "wb") as f:
        for word in program:
            if isinstance(word, str):
                try: word = labels[word]
                except KeyError:
                    print "unknown label: " + word
                    sys.exit(1)
            hi, lo = divmod(word, 0x100)
            f.write(chr(hi % 256))
            f.write(chr(lo))

if __name__ == '__main__':
    main()