Generate code with results from register allocation

example.gby

@@ -1,4 +1,8 @@
+[[ extern ]] u8 z;
+[[ extern ]] u8 w;
+
 u8 x;
 u8 y;
 
-x <- -y + 3;
+x <- -y + j;
+z <- w + 7;

gbuoy.js

@@ -1,2 +1,2 @@
 #!/usr/bin/env node
-require("./build/main")
+require("./build/lib/main")

lib/compile.ts

@@ -1,6 +1,9 @@
 import { convertIR } from "./ir/convert"
 import { prettyPrintBlock } from "./ir/pretty"
 import { parse } from "./parser"
+import { allocateRegisters } from "./regalloc"
+import { generateBlock } from "./sm83/codegen"
+import { R8 } from "./sm83/cpu"
 
 import type { Attr, Decl, Stmt, Type, VarDecl } from "./ast"
 
@@ -38,8 +41,23 @@ export const compile = (fileName: string, source: string): string => {
     console.log(prettyPrintBlock(ir))
 
     // 5. Generate code
+    const alloc = allocateRegisters(ir, Object.values(R8))
+    const insns = generateBlock(alloc, ir)
+
+    let output = ''
+    if (asmDecls.length > 0) {
+        output += `SECTION "${fileName} Data", WRAM0\n\n`
+        output += asmDecls.join("\n")
+        output += "\n\n"
+    }
+
+    if (insns.length > 0) {
+        output += `SECTION "${fileName} Code", ROM0\n\n`
+        output += insns.join("\n")
+    }
 
-    return ''
+    console.log("=== ASM ===")
+    return output
 }
 
 export const processDecls = (decls: Array<Decl>): SymbolMap => decls.reduce(processDecl, {})
@@ -47,6 +65,11 @@ export const processDecls = (decls: Array): SymbolMap => decls.reduce(proc
 export const processDecl = (symbols: SymbolMap, decl: Decl): SymbolMap => {
     const symbol = processVarDecl(symbols, decl)
 
+    // Don't declare extern symbols
+    if (hasAttr(symbol, { name: 'extern', args: [] })) {
+        return symbols
+    }
+
     return {
         ...symbols,
         [symbol.name]: symbol,
@@ -64,3 +87,30 @@ export const processVarDecl = (symbols: SymbolMap, decl: VarDecl): SymbolDefn =>
         type: decl.args.type,
     }
 }
+
+export const hasAttr = (symbol: SymbolDefn, attr: Attr) => {
+    let has = false
+
+    for (let i = 0; i < symbol.attrs.length; ++i) {
+        has = attrEquals(symbol.attrs[i], attr)
+        if (has) {
+            break
+        }
+    }
+
+    return has
+}
+
+const attrEquals = (attr1: Attr, attr2: Attr): boolean => {
+    if (attr1.name !== attr2.name || attr1.args.length !== attr2.args.length) {
+        return false
+    }
+
+    for (let i = 0; i < attr1.args.length; ++i) {
+        if (attr1.args[i] !== attr2.args[i]) {
+            return false
+        }
+    }
+
+    return true
+}

lib/data/graph.ts

@@ -18,6 +18,10 @@ export const createGraph = (cons: (v: AddVertex, e: AddEdge) => void): Gra
     const g: Graph<V> = { numVertices: 0, numEdges: 0, vertices: {}, edges: {} }
 
     const addVertex = (name: string, v: V) => {
+        if (typeof g.vertices[name] !== "undefined") {
+            return
+        }
+
         g.vertices[name] = v
         g.edges[name] = new Set()
         g.numVertices++
@@ -50,6 +54,8 @@ export const neighbors = (g: Graph, name: string): Set => {
 
 export const vertices = <V>(g: Graph<V>): Set<string> => new Set(Object.keys(g.vertices))
 
+export const edgeConnects = <V>(g: Graph<V>, v1: string, v2: string): boolean => g.edges[v1].has(v2)
+
 export const maxCardinalitySearch = <V>(g: Graph<V>): Array<string> => {
     const weights: { [s: string]: number } = {}
     const ordering: Array<string> = []
@@ -105,7 +111,7 @@ export const colorGreedy = (g: Graph, vertexOrdering: Array): Grap
     }
 
     const coloring: GraphColoring = { numColors: 0, colors: {} }
-    const usedColors = {}
+    const usedColors: { [c: number]: boolean } = {}
 
     vertexOrdering.forEach(v => {
         const ns = neighbors(g, v)

lib/data/set.ts

@@ -15,7 +15,7 @@ export const setEquals = (xs: Set, ys: Set): boolean => {
         return false
     }
 
-    for (const x of xs) {
+    for (const x of Array.from(xs)) {
         if (!ys.has(x)) {
             return false
         }

lib/ir/convert.ts

@@ -1,5 +1,6 @@
 import type { AssignStmt, Expr, Stmt } from "../ast"
-import { Loc } from "./loc"
+import { R8 } from "../sm83/cpu"
+import { Loc, LocType } from "./loc"
 import type { Operand, SSA } from "./ssa"
 
 type IRState = {
@@ -26,6 +27,42 @@ export const convertAssignStmt = (state: IRState, stmt: AssignStmt): Array
 }
 
 export const convertExpr = (state: IRState, dest: Loc, expr: Expr): Array<SSA> => {
+    const threeAddress = convertExpr3(state, dest, expr)
+    const output: Array<SSA> = []
+
+    threeAddress.forEach(ssa => {
+        // skip copies
+        if (!("op" in ssa)) {
+            output.push(ssa)
+            return
+        } else if ("source1" in ssa) {
+            let source1 = ssa.source1
+            if (typeof source1 !== 'number' && source1.type === LocType.VARIABLE) {
+                source1 = Loc.temp(`t${state.nextID++}`)
+                output.push(
+                    { dest: Loc.reg(R8.A), source: ssa.source1 },
+                    { dest: source1, source: Loc.reg(R8.A) },
+                )
+            }
+
+            output.push(
+                { dest: Loc.reg(R8.A), source: ssa.source },
+                { dest: Loc.reg(R8.A), source: Loc.reg(R8.A), op: ssa.op, source1 },
+                { dest: ssa.dest, source: Loc.reg(R8.A) },
+            )
+        } else {
+            output.push(
+                { dest: Loc.reg(R8.A), source: ssa.source },
+                { dest: Loc.reg(R8.A), source: Loc.reg(R8.A), op: ssa.op },
+                { dest: ssa.dest, source: Loc.reg(R8.A) },
+            )
+        }
+    })
+
+    return output
+}
+
+export const convertExpr3 = (state: IRState, dest: Loc, expr: Expr): Array<SSA> => {
     let expr_stmts: Array<SSA> = []
 
     if (typeof expr === "number") {
@@ -67,6 +104,6 @@ const getSource = (state: IRState, expr: Expr): [Operand, Array] => {
     }
 
     const source = Loc.temp(`t${state.nextID++}`)
-    const stmts = convertExpr(state, source, expr)
+    const stmts = convertExpr3(state, source, expr)
     return [source, stmts]
 }

lib/live.ts

@@ -1,33 +0,0 @@
-import type { Loc } from "./ir/loc"
-import type { SSA } from "./ir/ssa"
-
-export type LivenessInfo = Array<Set<Loc>>
-
-export const liveness = (block: Array<SSA>): LivenessInfo => {
-    const info: LivenessInfo = []
-
-    for (let i = block.length - 1; i >= 0; --i) {
-        const insn = block[i]
-        const last = info[i + 1] || new Set()
-
-        info[i] = new Set()
-
-        if (typeof insn.source !== "number") {
-            info[i].add(insn.source)
-        }
-
-        if ("source1" in insn && typeof insn.source1 !== "number") {
-            info[i].add(insn.source1)
-        }
-
-        last.forEach(loc => {
-            if (loc === insn.dest) {
-                return
-            }
-
-            info[i].add(loc)
-        })
-    }
-
-    return info
-}

lib/regalloc.ts

@@ -0,0 +1,107 @@
+import { colorGreedy, createGraph, Graph, maxCardinalitySearch } from "./data/graph"
+import { R8, StackOffset } from "./sm83/cpu"
+import type { Loc } from "./ir/loc"
+import type { SSA } from "./ir/ssa"
+
+export type LivenessInfo = Array<Set<Loc>>
+
+export const liveness = (block: Array<SSA>): LivenessInfo => {
+    const info: LivenessInfo = []
+
+    info[block.length] = new Set()
+    for (let i = block.length - 1; i >= 0; --i) {
+        const insn = block[i]
+        const last = info[i + 1]
+
+        info[i] = new Set()
+
+        if (typeof insn.source !== "number") {
+            info[i].add(insn.source)
+        }
+
+        if ("source1" in insn && typeof insn.source1 !== "number") {
+            info[i].add(insn.source1)
+        }
+
+        last.forEach(loc => {
+            if (loc.toString() === insn.dest.toString()) {
+                return
+            }
+
+            info[i].add(loc)
+        })
+    }
+
+    return info
+}
+
+export const locations = (block: Array<SSA>): Set<Loc> => {
+    const ls: Set<Loc> = new Set()
+
+    block.forEach(ssa => {
+        ls.add(ssa.dest)
+
+        if (typeof ssa.source !== "number") {
+            ls.add(ssa.source)
+        }
+
+        if ("source1" in ssa && typeof ssa.source1 !== "number") {
+            ls.add(ssa.source1)
+        }
+    })
+
+    return ls
+}
+
+export const interference = (block: Array<SSA>, live: LivenessInfo): Graph<Loc> =>
+    createGraph((v, e) => {
+        const locs = locations(block)
+        locs.forEach(loc => {
+            v(loc.toString(), loc)
+        })
+
+        block.forEach((ssa, i) =>
+            live[i + 1].forEach(u => {
+                if (ssa.dest.toString() !== u.toString()) {
+                    e(ssa.dest.toString(), u.toString())
+                }
+            })
+        )
+    })
+
+export type RegAlloc = {
+    [s: string]: string | StackOffset,
+}
+
+export const allocateRegisters = (block: Array<SSA>, registers: Array<string>): RegAlloc => {
+    const info = liveness(block)
+    const graph = interference(block, info)
+    const ordering = maxCardinalitySearch(graph)
+    const coloring = colorGreedy(graph, ordering)
+
+    const allocation: RegAlloc = {}
+    const availableRegisters = new Set(registers)
+    const colorMap: { [c: number]: string | StackOffset } = {}
+    let nextStackOffset = 0
+
+    Object.entries(coloring.colors).forEach(([vertex, color]) => {
+        if (typeof colorMap[color] !== 'undefined') {
+            allocation[vertex] = colorMap[color]
+            return
+        }
+
+        let value = null
+        if (availableRegisters.size == 0) {
+            value = { offset: nextStackOffset++ }
+        } else {
+            const result = availableRegisters.values().next()
+            value = result.value
+            availableRegisters.delete(value)
+        }
+
+        allocation[vertex] = value
+        colorMap[color] = value
+    })
+
+    return allocation
+}

lib/sm83/codegen.ts

@@ -0,0 +1,134 @@
+import { R8 } from "../sm83/cpu"
+import { Loc, LocType } from "../ir/loc"
+import type { Operand, SSA, SSABinary, SSACopy, SSAUnary } from "../ir/ssa"
+import type { RegAlloc } from "../regalloc"
+import { UnaryOp } from "../ast"
+
+export const generateBlock = (alloc: RegAlloc, block: Array<SSA>): Array<string> => {
+    const output: Array<string> = []
+
+    block.forEach(ssa => {
+        output.push(...generateSSA(alloc, ssa))
+    })
+
+    return output
+}
+
+export const generateSSA = (alloc: RegAlloc, ssa: SSA): Array<string> => {
+    let output: Array<string> = []
+
+    if ("source1" in ssa) {
+        output = generateSSA_Binary(alloc, ssa)
+    } else if ("op" in ssa) {
+        output = generateSSA_Unary(alloc, ssa)
+    } else {
+        output = generateSSA_Copy(alloc, ssa)
+    }
+
+    return output
+}
+
+export const generateSSA_Copy = (alloc: RegAlloc, ssa: SSACopy): Array<string> => {
+    const output: Array<string> = []
+
+    const source = typeof ssa.source === 'number'
+        ? ssa.source.toString()
+        : getAlloc(alloc, ssa.source).name
+
+    switch (ssa.dest.type) {
+        case LocType.TEMPORARY:
+            const dest = getAlloc(alloc, ssa.dest)
+            // TODO: Remove later with peephole optimizer
+            if (dest.name !== source) {
+                output.push(`LD ${dest.name}, ${source}`)
+            }
+            break
+
+        case LocType.REGISTER:
+            if (isA(ssa.dest) && typeof ssa.source !== 'number' && ssa.source.type === LocType.VARIABLE) {
+                output.push(`LD A, (${ssa.source.name})`)
+
+                // TODO: Remove check with peephole optimizer
+            } else if (source !== ssa.dest.name) {
+                output.push(`LD ${ssa.dest.name}, ${source}`)
+            }
+            break
+
+        case LocType.VARIABLE:
+            if (source !== R8.A) {
+                console.log(source, ssa)
+                throw new Error("unsupported")
+            }
+            output.push(`LD (${ssa.dest.name}), A`)
+            break
+
+        default:
+            throw new Error(`unsupported destination type \`${ssa.dest.type}'`)
+    }
+
+    return output
+}
+
+const getAlloc = (alloc: RegAlloc, loc: Loc): Loc => {
+    const destAlloc = alloc[loc.toString()]
+    if (typeof destAlloc === "object") {
+        throw new Error("stack variables not yet supported")
+    }
+
+    return Loc.reg(destAlloc)
+}
+
+export const generateSSA_Unary = (alloc: RegAlloc, ssa: SSAUnary): Array<string> => {
+    if (!isA(ssa.dest)) {
+        throw new Error("Unexpected form for unary operation")
+    }
+
+    const output: Array<string> = []
+    if (typeof ssa.source === "number") {
+        output.push(`LD A, ${ssa.source}`)
+    } else if (ssa.source.type === LocType.REGISTER) {
+        if (!isA(ssa.dest)) {
+            output.push(`LD A, ${ssa.source.name}`)
+        }
+    } else {
+        // TODO: ??
+        output.push(`LD A, ${alloc[ssa.source.toString()]}`)
+    }
+
+    const ops = unaryOps[ssa.op]
+    if (!ops) {
+        throw new Error(`unsupported unary op \`${ssa.op}'`)
+    }
+
+    return output.concat(ops)
+}
+
+const unaryOps = {
+    "arith_negate": ["CPL", "INC A"],
+    "bit_negate": ["CPL"]
+}
+
+export const generateSSA_Binary = (alloc: RegAlloc, ssa: SSABinary): Array<string> => {
+    if (!isA(ssa.dest) || !isA(ssa.source)) {
+        throw new Error("Unexpected form for binary operation")
+    }
+
+    const output: Array<string> = []
+    const source = typeof ssa.source1 === 'number'
+        ? ssa.source1
+        : getAlloc(alloc, ssa.source1).name
+
+    switch (ssa.op) {
+        case "add":
+            output.push(`ADD ${source}`)
+            break
+
+        default:
+            throw new Error(`unsupported binary op \`${ssa.op}'`)
+    }
+
+    return output
+}
+
+const isA = (loc: Operand): boolean =>
+    typeof loc === "object" && loc.type === LocType.REGISTER && loc.name === R8.A

lib/sm83/cpu.ts

@@ -0,0 +1,12 @@
+export enum R8 {
+    A = "A",
+    B = "B",
+    C = "C",
+    D = "D",
+    E = "E",
+    H = "H",
+    L = "L",
+    HL = "(HL)",
+}
+
+export type StackOffset = { offset: number }

package.json

@@ -19,7 +19,7 @@
     "lib": "lib"
   },
   "scripts": {
-    "build": "mkdir -p build && peggy -o build/parser.js lib/parser.pegjs && tsc",
+    "build": "tsc && peggy -o build/lib/parser.js lib/parser.pegjs",
     "test": "ts-mocha"
   },
   "author": "Forest Belton <forest@homolo.gy>",

test/live.spec.ts

@@ -1,28 +0,0 @@
-import { expect } from "chai"
-
-import { Loc } from "../lib/ir/loc"
-import type { SSA } from "../lib/ir/ssa"
-import { liveness } from "../lib/live"
-
-describe("liveness", () => {
-    it("computes liveness", () => {
-        const x = [0, 1, 2].map(i => Loc.vari("x" + i))
-        const y = [0, 1].map(i => Loc.vari("y" + i))
-
-        const block: Array<SSA> = [
-            { dest: x[0], source: 1 },
-            { dest: x[1], source: x[0], op: "add", source1: x[0] },
-            { dest: x[2], source: x[1], op: "add", source1: x[0] },
-            { dest: y[0], source: x[0], op: "add", source1: x[1] },
-            { dest: y[1], source: y[0], op: "add", source1: x[2] },
-        ]
-
-        const info = liveness(block)
-
-        expect(info[0]).to.deep.equal(new Set())
-        expect(info[1]).to.deep.equal(new Set([x[0]]))
-        expect(info[2]).to.deep.equal(new Set([x[0], x[1]]))
-        expect(info[3]).to.deep.equal(new Set([x[0], x[1], x[2]]))
-        expect(info[4]).to.deep.equal(new Set([y[0], x[2]]))
-    })
-})

test/regalloc.spec.ts

@@ -0,0 +1,61 @@
+import { expect } from "chai"
+
+import { edgeConnects } from "../lib/data/graph"
+import { Loc } from "../lib/ir/loc"
+import type { SSA } from "../lib/ir/ssa"
+import { allocateRegisters, interference, liveness } from "../lib/regalloc"
+import { R8 } from "../lib/sm83/cpu"
+
+describe("liveness", () => {
+    it("computes liveness", () => {
+        const x = [0, 1, 2].map(i => Loc.vari("x" + i))
+        const y = [0, 1].map(i => Loc.vari("y" + i))
+
+        const block: Array<SSA> = [
+            { dest: x[0], source: 1 },
+            { dest: x[1], source: x[0], op: "add", source1: x[0] },
+            { dest: x[2], source: x[1], op: "add", source1: x[0] },
+            { dest: y[0], source: x[0], op: "add", source1: x[1] },
+            { dest: y[1], source: y[0], op: "add", source1: x[2] },
+        ]
+
+        const info = liveness(block)
+
+        expect(info[0]).to.deep.equal(new Set())
+        expect(info[1]).to.deep.equal(new Set([x[0]]))
+        expect(info[2]).to.deep.equal(new Set([x[0], x[1]]))
+        expect(info[3]).to.deep.equal(new Set([x[0], x[1], x[2]]))
+        expect(info[4]).to.deep.equal(new Set([y[0], x[2]]))
+    })
+})
+
+describe("interference", () => {
+    it("computes interference", () => {
+        const block: Array<SSA> = [
+            { dest: Loc.vari("a"), source: 7 },
+            { dest: Loc.vari("b"), source: 3 },
+            { dest: Loc.vari("x"), source: Loc.vari("a") }
+        ]
+
+        const info = liveness(block)
+        const g = interference(block, info)
+
+        expect(edgeConnects(g, "a", "b")).to.be.true
+    })
+})
+
+describe("allocateRegisters", () => {
+    it("allocates registers", () => {
+        const block: Array<SSA> = [
+            { dest: Loc.vari("a"), source: 7 },
+            { dest: Loc.vari("b"), source: 3 },
+            { dest: Loc.vari("x"), source: Loc.vari("a") }
+        ]
+
+        const alloc = allocateRegisters(block, Object.values(R8))
+
+        expect(alloc.a).to.equal("A")
+        expect(alloc.b).to.equal("B")
+        expect(alloc.x).to.equal("A")
+    })
+})