Finish basic register allocation

lib/data/graph.ts

@@ -111,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/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

@@ -1,5 +1,5 @@
-import { createGraph, Graph } from "./data/graph"
-
+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"
 
@@ -39,10 +39,14 @@ export const locations = (block: Array): Set => {
     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)
-        } else if (typeof ssa.source !== "number") {
-            ls.add(ssa.source)
         }
     })
 
@@ -51,13 +55,53 @@ export const locations = (block: Array): Set => {
 
 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()) {
-                    v(ssa.dest.toString(), ssa.dest)
-                    v(u.toString(), u)
                     e(ssa.dest.toString(), u.toString())
                 }
             })
         )
     })
+
+type RegAlloc = {
+    [s: string]: R8 | StackOffset,
+}
+
+export const allocateRegisters = (block: Array<SSA>): 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(Object.values(R8))
+    const colorMap: { [c: number]: R8 | 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/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 }

test/regalloc.spec.ts

@@ -3,7 +3,7 @@ import { edgeConnects } from "../lib/data/graph"
 
 import { Loc } from "../lib/ir/loc"
 import type { SSA } from "../lib/ir/ssa"
-import { interference, liveness } from "../lib/regalloc"
+import { allocateRegisters, interference, liveness } from "../lib/regalloc"
 
 describe("liveness", () => {
     it("computes liveness", () => {
@@ -42,3 +42,19 @@ describe("interference", () => {
         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)
+
+        expect(alloc.a).to.equal("A")
+        expect(alloc.b).to.equal("B")
+        expect(alloc.x).to.equal("A")
+    })
+})