top-sort-asm/topsort_asm/isa.py

"""A tiny load/store ISA, just rich enough to exhibit every kind of dependence.

Syntax (one instruction per line, ``;`` starts a comment)::

    li    rd, imm          rd <- imm
    mov   rd, rs           rd <- rs
    add   rd, rs, rt       rd <- rs + rt        (also sub, mul)
    load  rd, [rs+off]     rd <- mem[rs + off]
    store rs, [rt+off]     mem[rt + off] <- rs
    cmp   rs, rt           flags <- (rs < rt, rs == rt)
    setlt rd               rd <- 1 if flags.lt else 0

``cmp`` and ``setlt`` communicate through an implicit ``flags`` register that
never appears in the source text. The parser makes it explicit (``cmp``
writes ``flags``; ``setlt`` reads it) so that dependence analysis treats it
exactly like any other register -- forgetting hidden operands like this is a
classic way to write a scheduler that silently reorders code incorrectly.

Memory is word-addressed: each address holds one value, so two accesses
conflict only if their addresses are equal.
"""

import re
from dataclasses import dataclass, replace
from typing import List, Optional, Tuple

FLAGS = "flags"

# Cycles from issue until the result can be consumed.
LATENCY = {
    "li": 1,
    "mov": 1,
    "add": 1,
    "sub": 1,
    "mul": 3,
    "load": 4,
    "store": 1,
    "cmp": 1,
    "setlt": 1,
}

ALU3 = ("add", "sub", "mul")


class AsmError(ValueError):
    pass


@dataclass(frozen=True)
class Instr:
    """One instruction with all of its operands made explicit.

    ``dest`` is the register written (``None`` for ``store``) and ``srcs`` the
    registers read, both including the implicit ``flags``. For memory
    operations ``srcs`` lists the value register (store only) before the base
    register, and ``imm`` holds the offset.
    """

    index: int
    op: str
    dest: Optional[str]
    srcs: Tuple[str, ...]
    imm: Optional[int] = None

    @property
    def reads(self) -> Tuple[str, ...]:
        return self.srcs

    @property
    def writes(self) -> Tuple[str, ...]:
        return (self.dest,) if self.dest is not None else ()

    @property
    def is_mem(self) -> bool:
        return self.op in ("load", "store")

    @property
    def mem_base(self) -> str:
        assert self.is_mem
        return self.srcs[-1]

    @property
    def latency(self) -> int:
        return LATENCY[self.op]

    def text(self) -> str:
        """Render back to assembly, hiding the implicit flags operand."""
        op = self.op
        if op == "li":
            args = f"{self.dest}, {self.imm}"
        elif op == "mov":
            args = f"{self.dest}, {self.srcs[0]}"
        elif op in ALU3:
            args = f"{self.dest}, {self.srcs[0]}, {self.srcs[1]}"
        elif op == "load":
            args = f"{self.dest}, {_addr(self.srcs[0], self.imm)}"
        elif op == "store":
            args = f"{self.srcs[0]}, {_addr(self.srcs[1], self.imm)}"
        elif op == "cmp":
            args = f"{self.srcs[0]}, {self.srcs[1]}"
        elif op == "setlt":
            args = f"{self.dest}"
        else:  # pragma: no cover - parse() rejects unknown ops
            raise AsmError(op)
        return f"{op:<5} {args}"

    def __str__(self) -> str:
        return self.text()


def _addr(base: str, off: int) -> str:
    if off == 0:
        return f"[{base}]"
    return f"[{base}{off:+d}]"


_REG = r"[a-z][a-z0-9_.]*"
_MEM = re.compile(rf"^\[\s*({_REG})\s*(?:([+-])\s*(\d+))?\s*\]$")


def _reg(tok: str, lineno: int) -> str:
    if not re.fullmatch(_REG, tok) or tok == FLAGS:
        raise AsmError(f"line {lineno}: bad register {tok!r}")
    return tok


def _imm(tok: str, lineno: int) -> int:
    try:
        return int(tok, 0)
    except ValueError:
        raise AsmError(f"line {lineno}: bad immediate {tok!r}") from None


def _mem(tok: str, lineno: int) -> Tuple[str, int]:
    m = _MEM.match(tok)
    if not m:
        raise AsmError(f"line {lineno}: bad memory operand {tok!r}")
    base, sign, off = m.groups()
    n = int(off) if off else 0
    return _reg(base, lineno), -n if sign == "-" else n


ARITY = {"li": 2, "mov": 2, "add": 3, "sub": 3, "mul": 3, "load": 2,
         "store": 2, "cmp": 2, "setlt": 1}


def parse(source: str) -> List[Instr]:
    """Parse a straight-line block (no labels or branches)."""
    out: List[Instr] = []
    for lineno, raw in enumerate(source.splitlines(), 1):
        line = raw.split(";", 1)[0].strip()
        if not line:
            continue
        op, _, rest = line.partition(" ")
        op = op.lower()
        if op not in ARITY:
            raise AsmError(f"line {lineno}: unknown instruction {op!r}")
        args = [a.strip() for a in rest.split(",")] if rest.strip() else []
        if len(args) != ARITY[op]:
            raise AsmError(
                f"line {lineno}: {op} takes {ARITY[op]} operands, got {len(args)}")
        i = len(out)
        if op == "li":
            ins = Instr(i, op, _reg(args[0], lineno), (), _imm(args[1], lineno))
        elif op == "mov":
            ins = Instr(i, op, _reg(args[0], lineno), (_reg(args[1], lineno),))
        elif op in ALU3:
            ins = Instr(i, op, _reg(args[0], lineno),
                        (_reg(args[1], lineno), _reg(args[2], lineno)))
        elif op == "load":
            base, off = _mem(args[1], lineno)
            ins = Instr(i, op, _reg(args[0], lineno), (base,), off)
        elif op == "store":
            base, off = _mem(args[1], lineno)
            ins = Instr(i, op, None, (_reg(args[0], lineno), base), off)
        elif op == "cmp":
            ins = Instr(i, op, FLAGS, (_reg(args[0], lineno), _reg(args[1], lineno)))
        else:  # setlt
            ins = Instr(i, op, _reg(args[0], lineno), (FLAGS,))
        out.append(ins)
    return out


def reindex(instrs: List[Instr]) -> List[Instr]:
    """Renumber instructions 0..n-1 in list order."""
    return [replace(ins, index=i) for i, ins in enumerate(instrs)]