asm.py

#!/usr/bin/env python3
"""Minimal RV32I assembler producing $readmemh hex files.

Supports: labels, the full RV32I user subset used in the course
(plus mul for demonstration is NOT included -- pure RV32I),
pseudo-instructions li (12-bit and full 32-bit), mv, nop, j, ret,
call (jal ra), beqz/bnez, and .word.
"""
import re
import sys

REGS = {f"x{i}": i for i in range(32)}
ABI = ("zero ra sp gp tp t0 t1 t2 s0 s1 a0 a1 a2 a3 a4 a5 a6 a7 "
       "s2 s3 s4 s5 s6 s7 s8 s9 s10 s11 t3 t4 t5 t6").split()
for i, n in enumerate(ABI):
    REGS[n] = i
REGS["fp"] = 8

R_OPS = {  # name: (funct7, funct3)
    "add": (0b0000000, 0b000), "sub": (0b0100000, 0b000),
    "sll": (0b0000000, 0b001), "slt": (0b0000000, 0b010),
    "sltu": (0b0000000, 0b011), "xor": (0b0000000, 0b100),
    "srl": (0b0000000, 0b101), "sra": (0b0100000, 0b101),
    "or": (0b0000000, 0b110), "and": (0b0000000, 0b111),
}
I_OPS = {  # name: funct3
    "addi": 0b000, "slti": 0b010, "sltiu": 0b011,
    "xori": 0b100, "ori": 0b110, "andi": 0b111,
}
SHIFT_OPS = {"slli": (0b0000000, 0b001), "srli": (0b0000000, 0b101),
             "srai": (0b0100000, 0b101)}
LOAD_OPS = {"lb": 0b000, "lh": 0b001, "lw": 0b010, "lbu": 0b100, "lhu": 0b101}
STORE_OPS = {"sb": 0b000, "sh": 0b001, "sw": 0b010}
BR_OPS = {"beq": 0b000, "bne": 0b001, "blt": 0b100,
          "bge": 0b101, "bltu": 0b110, "bgeu": 0b111}


def parse_imm(tok, labels=None, pc=None):
    tok = tok.strip()
    if labels is not None and tok in labels:
        return labels[tok] - (pc if pc is not None else 0)
    return int(tok, 0)


def enc_r(f7, rs2, rs1, f3, rd):
    return (f7 << 25) | (rs2 << 20) | (rs1 << 15) | (f3 << 12) | (rd << 7) | 0b0110011


def enc_i(imm, rs1, f3, rd, opcode):
    return ((imm & 0xFFF) << 20) | (rs1 << 15) | (f3 << 12) | (rd << 7) | opcode


def enc_s(imm, rs2, rs1, f3):
    return (((imm >> 5) & 0x7F) << 25) | (rs2 << 20) | (rs1 << 15) | \
           (f3 << 12) | ((imm & 0x1F) << 7) | 0b0100011


def enc_b(imm, rs2, rs1, f3):
    assert imm % 2 == 0
    return (((imm >> 12) & 1) << 31) | (((imm >> 5) & 0x3F) << 25) | \
           (rs2 << 20) | (rs1 << 15) | (f3 << 12) | \
           (((imm >> 1) & 0xF) << 8) | (((imm >> 11) & 1) << 7) | 0b1100011


def enc_u(imm20, rd, opcode):
    return ((imm20 & 0xFFFFF) << 12) | (rd << 7) | opcode


def enc_j(imm, rd):
    assert imm % 2 == 0
    return (((imm >> 20) & 1) << 31) | (((imm >> 1) & 0x3FF) << 21) | \
           (((imm >> 11) & 1) << 20) | (((imm >> 12) & 0xFF) << 12) | \
           (rd << 7) | 0b1101111


def expand_pseudo(mnem, args):
    """Return list of (mnem, args) after pseudo expansion."""
    if mnem == "nop":
        return [("addi", ["x0", "x0", "0"])]
    if mnem == "mv":
        return [("addi", [args[0], args[1], "0"])]
    if mnem == "not":
        return [("xori", [args[0], args[1], "-1"])]
    if mnem == "neg":
        return [("sub", [args[0], "x0", args[1]])]
    if mnem == "seqz":
        return [("sltiu", [args[0], args[1], "1"])]
    if mnem == "snez":
        return [("sltu", [args[0], "x0", args[1]])]
    if mnem == "j":
        return [("jal", ["x0", args[0]])]
    if mnem == "call":
        return [("jal", ["ra", args[0]])]
    if mnem == "ret":
        return [("jalr", ["x0", "0(ra)"])]
    if mnem == "jr":
        return [("jalr", ["x0", f"0({args[0]})"])]
    if mnem == "beqz":
        return [("beq", [args[0], "x0", args[1]])]
    if mnem == "bnez":
        return [("bne", [args[0], "x0", args[1]])]
    if mnem == "li":
        val = int(args[1], 0) & 0xFFFFFFFF
        if val >= 0x80000000:
            sval = val - (1 << 32)
        else:
            sval = val
        if -2048 <= sval <= 2047:
            return [("addi", [args[0], "x0", str(sval)])]
        upper = (val + 0x800) >> 12
        lower = val - (upper << 12)
        if lower >= 2048:
            lower -= 4096
        ops = [("lui", [args[0], str(upper & 0xFFFFF)])]
        if lower != 0:
            ops.append(("addi", [args[0], args[0], str(lower)]))
        return ops
    return [(mnem, args)]


def assemble(lines):
    # pass 1: layout
    labels = {}
    prog = []  # (pc, mnem, args, lineno)
    pc = 0
    for ln, raw in enumerate(lines, 1):
        line = raw.split("#")[0].strip()
        if not line:
            continue
        while ":" in line:
            lbl, line = line.split(":", 1)
            labels[lbl.strip()] = pc
            line = line.strip()
        if not line:
            continue
        parts = line.split(None, 1)
        mnem = parts[0].lower()
        args = [a.strip() for a in parts[1].split(",")] if len(parts) > 1 else []
        if mnem == ".word":
            prog.append((pc, ".word", args, ln))
            pc += 4
            continue
        for m, a in expand_pseudo(mnem, args):
            prog.append((pc, m, a, ln))
            pc += 4
    # pass 2: encode
    words = []
    for pc, mnem, args, ln in prog:
        try:
            words.append(encode(pc, mnem, args, labels))
        except Exception as e:
            raise SystemExit(f"line {ln}: cannot assemble {mnem} {args}: {e}")
    return words


def encode(pc, mnem, args, labels):
    if mnem == ".word":
        return parse_imm(args[0]) & 0xFFFFFFFF
    if mnem in R_OPS:
        f7, f3 = R_OPS[mnem]
        return enc_r(f7, REGS[args[2]], REGS[args[1]], f3, REGS[args[0]])
    if mnem in I_OPS:
        return enc_i(parse_imm(args[2]), REGS[args[1]], I_OPS[mnem],
                     REGS[args[0]], 0b0010011)
    if mnem in SHIFT_OPS:
        f7, f3 = SHIFT_OPS[mnem]
        sh = parse_imm(args[2]) & 0x1F
        return enc_i((f7 << 5) | sh, REGS[args[1]], f3, REGS[args[0]], 0b0010011)
    if mnem in LOAD_OPS:
        m = re.match(r"(-?\w+)\((\w+)\)", args[1])
        return enc_i(parse_imm(m.group(1)), REGS[m.group(2)],
                     LOAD_OPS[mnem], REGS[args[0]], 0b0000011)
    if mnem in STORE_OPS:
        m = re.match(r"(-?\w+)\((\w+)\)", args[1])
        return enc_s(parse_imm(m.group(1)), REGS[args[0]],
                     REGS[m.group(2)], STORE_OPS[mnem])
    if mnem in BR_OPS:
        off = parse_imm(args[2], labels, pc)
        return enc_b(off, REGS[args[1]], REGS[args[0]], BR_OPS[mnem])
    if mnem == "lui":
        return enc_u(parse_imm(args[1]), REGS[args[0]], 0b0110111)
    if mnem == "auipc":
        return enc_u(parse_imm(args[1]), REGS[args[0]], 0b0010111)
    if mnem == "jal":
        if len(args) == 1:
            args = ["ra", args[0]]
        off = parse_imm(args[1], labels, pc)
        return enc_j(off, REGS[args[0]])
    if mnem == "jalr":
        m = re.match(r"(-?\w+)\((\w+)\)", args[1])
        return enc_i(parse_imm(m.group(1)), REGS[m.group(2)], 0b000,
                     REGS[args[0]], 0b1100111)
    if mnem == "ecall":
        return 0x00000073
    if mnem == "ebreak":
        return 0x00100073
    if mnem == "fence":
        return 0x0FF0000F
    raise ValueError("unknown mnemonic")


def main():
    if len(sys.argv) != 3:
        print("usage: asm.py in.s out.hex")
        sys.exit(1)
    with open(sys.argv[1]) as f:
        words = assemble(f.readlines())
    with open(sys.argv[2], "w") as f:
        for w in words:
            f.write(f"{w:08x}\n")
    print(f"{len(words)} instructions -> {sys.argv[2]}")


if __name__ == "__main__":
    main()