Capstone: an end-to-end compiler, JITs, correctness and the modern toolchain

# Tern end to end: scanner, parser, checks, constant folding, RISC-V code
# generation and a small RISC-V emulator to run the result. Integer subset:
# functions, var, assignment, if/else, while, return, calls, print.
import re
def lex(src):
    spec = r"\s+|//[^\n]*|(\d+)|([A-Za-z_]\w*)|(==|!=|<=|>=|->|[-+*/%<>=(){},;:])|(.)"
    toks = []
    for m in re.finditer(spec, src):
        num, name, op, bad = m.groups()
        if bad: raise SyntaxError(f"unexpected character {bad!r}")
        if num: toks.append(("num", int(num)))
        elif name: toks.append(("kw" if name in ("fn", "var", "if", "else", "while", "return") else "id", name))
        elif op: toks.append(("op", op))
    return toks + [("eof", None)]
LB, RB = "{", "}"
BP = {"==": 1, "!=": 1, "<": 2, "<=": 2, ">": 2, ">=": 2, "+": 3, "-": 3, "*": 4, "/": 4, "%": 4}
class Parser:
    def __init__(p, toks): p.t, p.i = toks, 0
    def peek(p): return p.t[p.i][1]
    def next(p): p.i += 1; return p.t[p.i - 1]
    def expect(p, v):
        k, x = p.next()
        if x != v: raise SyntaxError(f"expected {v!r}, found {x!r}")
    def program(p):
        fns = []
        while p.t[p.i][0] != "eof": fns.append(p.func())
        return fns
    def func(p):
        p.expect("fn"); name = p.next()[1]; p.expect("("); params = []
        while p.peek() != ")":
            params.append(p.next()[1]); p.expect(":"); p.next()
            if p.peek() == ",": p.next()
        p.expect(")"); p.expect("->"); p.next()
        return ("fn", name, params, p.block())
    def block(p):
        p.expect(LB); body = []
        while p.peek() != RB: body.append(p.stmt())
        p.expect(RB); return body
    def stmt(p):
        v = p.peek()
        if v == "var":
            p.next(); name = p.next()[1]
            if p.peek() == ":": p.next(); p.next()
            p.expect("="); e = p.expr(); p.expect(";"); return ("var", name, e)
        if v == "if":
            p.next(); c = p.expr(); th = p.block(); el = []
            if p.peek() == "else": p.next(); el = [p.stmt()] if p.peek() == "if" else p.block()
            return ("if", c, th, el)
        if v == "while": p.next(); c = p.expr(); return ("while", c, p.block())
        if v == "return": p.next(); e = p.expr(); p.expect(";"); return ("ret", e)
        e = p.expr()
        if p.peek() == "=": p.next(); v2 = p.expr(); p.expect(";"); return ("set", e[1], v2)
        p.expect(";"); return ("expr", e)
    def expr(p, min_bp=0):
        k, v = p.next()
        if k == "num": left = ("num", v)
        elif v == "(": left = p.expr(); p.expect(")")
        elif v == "-": left = ("bin", "-", ("num", 0), p.expr(5))
        elif k == "id" and p.peek() == "(":
            p.next(); args = []
            while p.peek() != ")":
                args.append(p.expr())
                if p.peek() == ",": p.next()
            p.next(); left = ("call", v, args)
        elif k == "id": left = ("name", v)
        else: raise SyntaxError(f"unexpected {v!r}")
        while BP.get(p.peek(), 0) > min_bp:
            op = p.next()[1]; left = ("bin", op, left, p.expr(BP[op]))
        return left
def check(fns):                                   # names declared, arity of calls
    arity = {f[1]: len(f[2]) for f in fns}; arity["print"] = 1; errs = []
    def ex(e, env):
        if e[0] == "name" and e[1] not in env: errs.append(f"undeclared {e[1]!r}")
        if e[0] == "bin": ex(e[2], env); ex(e[3], env)
        if e[0] == "call":
            if arity.get(e[1]) != len(e[2]): errs.append(f"bad call to {e[1]!r}")
            for a in e[2]: ex(a, env)
    def blk(b, env):
        env = set(env)
        for s in b:
            if s[0] == "var": ex(s[2], env); env.add(s[1])
            elif s[0] == "set": ex(("name", s[1]), env); ex(s[2], env)
            elif s[0] in ("if", "while"): ex(s[1], env); blk(s[2], env); s[0] == "if" and blk(s[3], env)
            else: ex(s[1], env)
    for f in fns: blk(f[3], set(f[2]))
    return errs
FOLD = {"+": lambda a, b: a + b, "-": lambda a, b: a - b, "*": lambda a, b: a * b}
def fold(e):                                       # constant folding on the AST
    if e[0] == "bin":
        a, b = fold(e[2]), fold(e[3])
        if a[0] == b[0] == "num" and e[1] in FOLD: return ("num", FOLD[e[1]](a[1], b[1]))
        return ("bin", e[1], a, b)
    if e[0] == "call": return ("call", e[1], [fold(x) for x in e[2]])
    return e
class CodeGen:                                     # every variable in a frame slot (like -O0);
    def __init__(g): g.out, g.n = [], 0            # expression temporaries in t0..t6
    def emit(g, s): g.out.append(s)
    def label(g): g.n += 1; return f".L{g.n}"
    def func(g, f):
        _, name, params, body = f
        slots = {}
        def collect(b):
            for s in b:
                if s[0] == "var" and s[1] not in slots: slots[s[1]] = None
                if s[0] in ("if", "while"): collect(s[2]); s[0] == "if" and collect(s[3])
        for prm in params: slots[prm] = None
        collect(body)
        for k, v in enumerate(slots): slots[v] = -24 - 8 * k
        size = (16 + 8 * len(slots) + 15) // 16 * 16
        g.slots, g.size, g.ret = slots, size, g.label()
        g.emit(f"{name}:")
        for ins in (f"addi sp, sp, -{size}", f"sd ra, {size - 8}(sp)", f"sd s0, {size - 16}(sp)", f"addi s0, sp, {size}"):
            g.emit("    " + ins)
        for k, prm in enumerate(params): g.emit(f"    sd a{k}, {slots[prm]}(s0)")
        for s in body: g.stmt(s)
        g.emit(f"{g.ret}:")
        for ins in (f"ld ra, {size - 8}(sp)", f"ld s0, {size - 16}(sp)", f"addi sp, sp, {size}", "ret"):
            g.emit("    " + ins)
    def expr(g, e, d=0):                           # value of e into register t{d}
        r = f"t{d}"
        if d > 6: raise RuntimeError("expression too deep for t0-t6")
        if e[0] == "num": g.emit(f"    li {r}, {e[1]}")
        elif e[0] == "name": g.emit(f"    ld {r}, {g.slots[e[1]]}(s0)")
        elif e[0] == "call":
            for k, a in enumerate(e[2]):           # evaluate arguments, park them on the frame
                g.expr(a, d); g.emit(f"    addi sp, sp, -16"); g.emit(f"    sd {r}, 0(sp)")
            for k in reversed(range(len(e[2]))):
                g.emit(f"    ld a{k}, 0(sp)"); g.emit(f"    addi sp, sp, 16")
            for k in range(d): g.emit(f"    addi sp, sp, -16"); g.emit(f"    sd t{k}, 0(sp)")   # caller-saved
            g.emit(f"    call {e[1]}")
            for k in reversed(range(d)): g.emit(f"    ld t{k}, 0(sp)"); g.emit(f"    addi sp, sp, 16")
            g.emit(f"    mv {r}, a0")
        else:
            op = e[1]; g.expr(e[2], d); g.expr(e[3], d + 1); s = f"t{d + 1}"
            simple = {"+": "add", "-": "sub", "*": "mul", "/": "div", "%": "rem"}
            if op in simple: g.emit(f"    {simple[op]} {r}, {r}, {s}")
            elif op == "<": g.emit(f"    slt {r}, {r}, {s}")
            elif op == ">": g.emit(f"    slt {r}, {s}, {r}")
            elif op == "<=": g.emit(f"    slt {r}, {s}, {r}"); g.emit(f"    xori {r}, {r}, 1")
            elif op == ">=": g.emit(f"    slt {r}, {r}, {s}"); g.emit(f"    xori {r}, {r}, 1")
            elif op == "==": g.emit(f"    sub {r}, {r}, {s}"); g.emit(f"    seqz {r}, {r}")
            elif op == "!=": g.emit(f"    sub {r}, {r}, {s}"); g.emit(f"    snez {r}, {r}")
    def stmt(g, s):
        k = s[0]
        if k in ("var", "set"): g.expr(fold(s[2])); g.emit(f"    sd t0, {g.slots[s[1]]}(s0)")
        elif k == "expr": g.expr(fold(s[1]))
        elif k == "ret": g.expr(fold(s[1])); g.emit("    mv a0, t0"); g.emit(f"    j {g.ret}")
        elif k == "if":
            els, end = g.label(), g.label()
            g.expr(fold(s[1])); g.emit(f"    beqz t0, {els}")
            for x in s[2]: g.stmt(x)
            g.emit(f"    j {end}"); g.emit(f"{els}:")
            for x in s[3]: g.stmt(x)
            g.emit(f"{end}:")
        elif k == "while":
            top, end = g.label(), g.label()
            g.emit(f"{top}:"); g.expr(fold(s[1])); g.emit(f"    beqz t0, {end}")
            for x in s[2]: g.stmt(x)
            g.emit(f"    j {top}"); g.emit(f"{end}:")
def run(asm, entry="main"):                        # a tiny RV64 emulator for this subset
    lines = [l.strip() for l in asm]
    labels = {l[:-1]: i for i, l in enumerate(lines) if l.endswith(":")}
    R = {f"t{i}": 0 for i in range(7)} | {f"a{i}": 0 for i in range(8)} | {"sp": 1 << 20, "s0": 0, "ra": -1, "zero": 0}
    mem, out, pc, steps = {}, [], labels[entry], 0
    def off(x): o, b = re.match(r"(-?\d+)\((\w+)\)", x).groups(); return R[b] + int(o)
    while pc != -1:
        ins = lines[pc]; pc += 1; steps += 1
        if ins.endswith(":"): continue
        m, *a = re.split(r"[,\s]+", ins)
        if m == "li": R[a[0]] = int(a[1])
        elif m == "mv": R[a[0]] = R[a[1]]
        elif m == "addi": R[a[0]] = R[a[1]] + int(a[2])
        elif m == "xori": R[a[0]] = R[a[1]] ^ int(a[2])
        elif m in ("add", "sub", "mul", "div", "rem", "slt"):
            x, y = R[a[1]], R[a[2]]
            R[a[0]] = {"add": lambda: x + y, "sub": lambda: x - y, "mul": lambda: x * y,
                       "div": lambda: int(x / y), "rem": lambda: x - y * int(x / y),
                       "slt": lambda: int(x < y)}[m]()
        elif m == "seqz": R[a[0]] = int(R[a[1]] == 0)
        elif m == "snez": R[a[0]] = int(R[a[1]] != 0)
        elif m == "sd": mem[off(a[1])] = R[a[0]]
        elif m == "ld": R[a[0]] = mem[off(a[1])]
        elif m == "beqz": pc = labels[a[1]] if R[a[0]] == 0 else pc
        elif m == "j": pc = labels[a[0]]
        elif m == "call":
            if a[0] == "print": out.append(R["a0"])
            else: R["ra"] = pc; pc = labels[a[0]]
        elif m == "ret": pc = R["ra"]
    return R["a0"], out, steps
SRC = """
fn fib(n: int) -> int {
    if n < 2 { return n; }
    return fib(n - 1) + fib(n - 2);
}
fn gcd(a: int, b: int) -> int {
    while b != 0 { var t = a % b; a = b; b = t; }
    return a;
}
fn main() -> int {
    var i = 0;
    while i <= 10 { print(fib(i)); i = i + 1; }
    print(gcd(1071, 462));
    print(2 * 3 * 7);              // folded to 42 at compile time
    return 0;
}"""
fns = Parser(lex(SRC)).program()
errors = check(fns)
if errors: raise SystemExit(errors)
g = CodeGen()
for f in fns: g.func(f)
print(f"{len(fns)} functions, {sum(not l.endswith(':') for l in g.out)} instructions; code for gcd:")
start = g.out.index("gcd:"); end = g.out.index("main:")
print("\n".join(g.out[start:end]))
result, printed, steps = run(g.out)
print("program output:", printed)
print(f"main returned {result} after {steps} emulated instructions")

3 functions, 130 instructions; code for gcd:
gcd:
    addi sp, sp, -48
    sd ra, 40(sp)
    sd s0, 32(sp)
    addi s0, sp, 48
    sd a0, -24(s0)
    sd a1, -32(s0)
.L5:
    ld t0, -32(s0)
    li t1, 0
    sub t0, t0, t1
    snez t0, t0
    beqz t0, .L6
    ld t0, -24(s0)
    ld t1, -32(s0)
    rem t0, t0, t1
    sd t0, -40(s0)
    ld t0, -32(s0)
    sd t0, -24(s0)
    ld t0, -40(s0)
    sd t0, -32(s0)
    j .L5
.L6:
    ld t0, -24(s0)
    mv a0, t0
    j .L4
.L4:
    ld ra, 40(sp)
    ld s0, 32(sp)
    addi sp, sp, 48
    ret
program output: [0, 1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 21, 42]
main returned 0 after 13836 emulated instructions

# Differential testing of an optimizer: generate random expressions, compare
# the optimized program's result with an unoptimized reference, and shrink
# any failing case to a minimal one. One rewrite rule below has a bug.
import random
random.seed(415)
def gen(depth):
    if depth == 0 or random.random() < 0.3:
        return ("num", random.randint(-5, 9)) if random.random() < 0.5 else ("var", random.choice("xyz"))
    return (random.choice("+-*"), gen(depth - 1), gen(depth - 1))
def evaluate(e, env):
    if e[0] == "num": return e[1]
    if e[0] == "var": return env[e[1]]
    a, b = evaluate(e[1], env), evaluate(e[2], env)
    return a + b if e[0] == "+" else a - b if e[0] == "-" else a * b
def simplify(e, buggy):
    if e[0] in ("num", "var"): return e
    op, a, b = e[0], simplify(e[1], buggy), simplify(e[2], buggy)
    if a[0] == b[0] == "num": return ("num", evaluate((op, a, b), {}))
    if op == "*" and ("num", 1) in (a, b): return b if a == ("num", 1) else a
    if op == "*" and ("num", 0) in (a, b): return ("num", 0)
    if op == "+" and b == ("num", 0): return a
    if op == "-" and a == b: return ("num", 0)
    if op == "-" and b[0] == "-":                 # a - (p - q)  ==>  (a - p) + q
        return ("-", ("-", a, b[1]), b[2]) if buggy else ("+", ("-", a, b[1]), b[2])
    return (op, a, b)
def fails(e, buggy):
    for env in ({"x": 3, "y": -2, "z": 7}, {"x": 0, "y": 5, "z": 1}):
        if evaluate(simplify(e, buggy), env) != evaluate(e, env): return True
    return False
def candidates(e):                                # smaller variants: a subtree, or a child shrunk
    if e[0] not in ("+", "-", "*"): return []
    op, a, b = e
    return [a, b] + [(op, s, b) for s in candidates(a)] + [(op, a, s) for s in candidates(b)]
def shrink(e, buggy):                             # greedy test-case reduction
    progress = True
    while progress:
        progress = False
        for cand in candidates(e):
            if fails(cand, buggy): e, progress = cand, True; break
    return e
def show(e):
    if e[0] == "num": return str(e[1]) if e[1] >= 0 else f"({e[1]})"
    if e[0] == "var": return e[1]
    return f"({show(e[1])} {e[0]} {show(e[2])})"
for buggy in (True, False):
    found = None
    for n in range(1, 3001):
        e = gen(5)
        if fails(e, buggy): found = (n, e); break
    label = "buggy optimizer" if buggy else "fixed optimizer"
    if not found: print(f"{label}: 3000 random programs, no differences"); continue
    n, e = found; small = shrink(e, buggy)
    print(f"{label}: mismatch on program {n} ({len(show(e))} characters)")
    print("   minimized:", show(small), " optimized to:", show(simplify(small, buggy)))
    print("   with x=3, y=-2, z=7:", evaluate(small, {"x": 3, "y": -2, "z": 7}), "expected,",
          evaluate(simplify(small, buggy), {"x": 3, "y": -2, "z": 7}), "computed")

buggy optimizer: mismatch on program 8 (73 characters)
   minimized: (y - (6 - x))  optimized to: ((y - 6) - x)
   with x=3, y=-2, z=7: -5 expected, -11 computed
fixed optimizer: 3000 random programs, no differences

cat greet.c
clang --target=riscv64-linux-gnu -O2 -fstack-protector-strong -S greet.c -o - | sed -n "/^greet:/,/^\.Lfunc_end/p" | grep -v "^\s\+\." | grep -v "^#" | grep -v "^\.Lfunc_end"

$ cat greet.c
char *strcpy(char *dst, const char *src);
int puts(const char *s);
void greet(const char *name) {
    char buf[16];
    strcpy(buf, name);   /* unsafe if name is long */
    puts(buf);
}
$ clang --target=riscv64-linux-gnu -O2 -fstack-protector-strong -S greet.c -o - | sed -n "/^greet:/,/^\.Lfunc_end/p" | grep -v "^\s\+\." | grep -v "^#" | grep -v "^\.Lfunc_end"
greet:                                  # @greet
	addi	sp, sp, -48
	sd	ra, 40(sp)                      # 8-byte Folded Spill
	sd	s0, 32(sp)                      # 8-byte Folded Spill
.Lpcrel_hi0:
	auipc	a1, %got_pcrel_hi(__stack_chk_guard)
	ld	s0, %pcrel_lo(.Lpcrel_hi0)(a1)
	ld	a2, 0(s0)
	mv	a1, a0
	sd	a2, 24(sp)
	addi	a0, sp, 8
	call	strcpy
	addi	a0, sp, 8
	call	puts
	ld	a0, 0(s0)
	ld	a1, 24(sp)
	bne	a0, a1, .LBB0_2
	ld	ra, 40(sp)                      # 8-byte Folded Reload
	ld	s0, 32(sp)                      # 8-byte Folded Reload
	addi	sp, sp, 48
	ret
.LBB0_2:
	call	__stack_chk_fail