Factor out register assignment from code generation

lib/codegen/allocate.ts

@@ -0,0 +1,33 @@
+import type { RegAlloc } from '../regalloc'
+import { AbsInsn2, CopyInsn, insnReduce2, Operand, UnaryInsn } from '../ir/insn'
+import { Loc, LocType } from '../ir/loc'
+import { UnaryOp, BinaryOp } from '../ast'
+
+export const allocate = (alloc: RegAlloc, block: Array<AbsInsn2>): Array<AbsInsn2> =>
+    block.map(insn => allocateInsn(alloc, insn))
+
+export const allocateInsn = (alloc: RegAlloc, insn: AbsInsn2): AbsInsn2 => insnReduce2<AbsInsn2>(
+    (dest: Loc, source: Operand) => CopyInsn(
+        allocateInsnDest(alloc, dest),
+        typeof source === 'number' ? source : allocateInsnDest(alloc, source),
+    ),
+    (dest: Loc, op: UnaryOp | BinaryOp, source: Operand) => UnaryInsn(
+        allocateInsnDest(alloc, dest),
+        op,
+        typeof source === 'number' ? source : allocateInsnDest(alloc, source),
+    ),
+    insn,
+)
+
+export const allocateInsnDest = (alloc: RegAlloc, loc: Loc): Loc =>
+    loc.type === LocType.TEMPORARY
+        ? lookupAllocation(alloc, loc.ppName())
+        : loc
+
+const lookupAllocation = (alloc: RegAlloc, name: string): Loc => {
+    const loc = alloc[name]
+    if (typeof loc === 'object') {
+        throw new Error('stack offsets not supported yet')
+    }
+    return Loc.reg(loc)
+}

lib/compile.ts

@@ -2,8 +2,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 { realizeBlock } from "./sm83/realize"
 import { R8 } from "./sm83/cpu"
+import { allocate } from "./codegen/allocate"
 
 import type { Attr, Decl, Stmt, Type, VarDecl } from "./ast"
 
@@ -42,7 +43,7 @@ export const compile = (fileName: string, source: string): string => {
 
     // 5. Generate code
     const alloc = allocateRegisters(ir, Object.values(R8))
-    const insns = generateBlock(alloc, ir)
+    const insns = realizeBlock(allocate(alloc, ir))
 
     let output = ''
     if (asmDecls.length > 0) {

lib/ir/insn.ts

@@ -22,6 +22,15 @@ export type AbsInsn2Unary = UnaryInsn
 // Abstract instruction in two-address form
 export type AbsInsn2 = AbsInsnCopy | AbsInsn2Unary
 
+type CopyFn<A> = (dest: Loc, source: Operand) => A
+type UnaryFn<Op, A> = (dest: Loc, op: Op, source: Operand) => A
+type UnaryFn2<A> = UnaryFn<UnaryOp | BinaryOp, A>
+
+export const insnReduce2 = <A>(copy: CopyFn<A>, unary: UnaryFn2<A>, insn: AbsInsn2): A =>
+    insn.type === 'copy'
+        ? copy(insn.dest, insn.source)
+        : unary(insn.dest, insn.op, insn.source)
+
 export type AbsInsn3Unary = UnaryInsn<UnaryOp>
 
 export type AbsInsnBinary = {

lib/ir/loc.ts

@@ -13,20 +13,37 @@ export class Loc {
         this.name = name
     }
 
-    toString(): string {
-        let sigil = ''
-
+    reduce<A>(
+        vari: (name: string) => A,
+        temp: (name: string) => A,
+        reg: (name: string) => A,
+    ): A {
         switch (this.type) {
-            case 'register':
-                sigil = '%'
-                break
-
-            case 'temporary':
-                sigil = '#'
-                break
+            case LocType.VARIABLE:
+                return vari(this.name)
+            case LocType.TEMPORARY:
+                return temp(this.name)
+            case LocType.REGISTER:
+                return reg(this.name)
         }
+    }
+
+    ppName(): string {
+        return this.reduce(
+            (name: string) => name,
+            (name: string) => '#' + name,
+            (name: string) => '%' + name,
+        )
+    }
 
-        return sigil + this.name
+    asmName(): string {
+        return this.reduce(
+            (name: string) => `(${name})`,
+            (_name: string) => {
+                throw new Error("temporaries do not exist in assembly!")
+            },
+            (name: string) => name.toUpperCase(),
+        )
     }
 
     static vari(name: string): Loc {

lib/regalloc.ts

@@ -20,7 +20,7 @@ export const liveness = (block: Array): LivenessInfo => {
         }
 
         last.forEach(loc => {
-            if (loc.toString() === insn.dest.toString()) {
+            if (loc.ppName() === insn.dest.ppName()) {
                 return
             }
 
@@ -49,13 +49,13 @@ export const interference = (block: Array, live: LivenessInfo): Graph<
     createGraph((v, e) => {
         const locs = locations(block)
         locs.forEach(loc => {
-            v(loc.toString(), loc)
+            v(loc.ppName(), loc)
         })
 
         block.forEach((insn, i) =>
             live[i + 1].forEach(u => {
-                if (insn.dest.toString() !== u.toString()) {
-                    e(insn.dest.toString(), u.toString())
+                if (insn.dest.ppName() !== u.ppName()) {
+                    e(insn.dest.ppName(), u.ppName())
                 }
             })
         )

lib/sm83/codegen.ts

@@ -1,105 +0,0 @@
-import { R8 } from "../sm83/cpu"
-import { Loc, LocType } from "../ir/loc"
-import type { Operand, AbsInsn2, AbsInsnCopy, AbsInsn2Unary } from "../ir/insn"
-import type { RegAlloc } from "../regalloc"
-
-export const generateBlock = (alloc: RegAlloc, block: Array<AbsInsn2>): Array<string> => {
-    const output: Array<string> = []
-
-    block.forEach(ssa => {
-        output.push(...generateSSA(alloc, ssa))
-    })
-
-    return output
-}
-
-export const generateSSA = (alloc: RegAlloc, ssa: AbsInsn2): Array<string> => {
-    let output: Array<string> = []
-
-    switch (ssa.type) {
-        case 'copy':
-            output = generateSSA_Copy(alloc, ssa)
-            break
-
-        case 'unary':
-            output = generateSSA_Unary(alloc, ssa)
-            break
-    }
-
-    return output
-}
-
-export const generateSSA_Copy = (alloc: RegAlloc, ssa: AbsInsnCopy): 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 later with peephole optimizer
-            } else if (source !== ssa.dest.name) {
-                output.push(`LD ${ssa.dest.name}, ${source}`)
-            }
-            break
-
-        case LocType.VARIABLE:
-            if (source !== R8.A) {
-                throw new Error(`can only deref address into A (had ${source})`)
-            }
-            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: AbsInsn2Unary): Array<string> => {
-    if (!isA(ssa.dest)) {
-        throw new Error("Unexpected form for unary operation")
-    }
-
-    const source = typeof ssa.source === 'number'
-        ? ssa.source.toString()
-        : getAlloc(alloc, ssa.source).name
-
-    switch (ssa.op) {
-        case "add":
-            return [`ADD ${source}`]
-
-        case "arith_negate":
-            return ["CPL", "INC A"]
-
-        case "bit_negate":
-            return ["CPL"]
-
-        default:
-            throw new Error(`unsupported unary op \`${ssa.op}'`)
-    }
-}
-
-const isA = (loc: Operand): boolean =>
-    typeof loc === "object" && loc.type === LocType.REGISTER && loc.name === R8.A

lib/sm83/insn.ts

@@ -0,0 +1 @@
+export type SM83Insn = string

lib/sm83/realize.ts

@@ -0,0 +1,40 @@
+import { R8 } from "./cpu"
+import { Loc, LocType } from "../ir/loc"
+import { Operand, AbsInsn2, insnReduce2 } from "../ir/insn"
+import { BinaryOp, UnaryOp } from "../ast"
+import type { SM83Insn } from "./insn"
+
+export const realizeBlock = (block: Array<AbsInsn2>): Array<SM83Insn> => block.flatMap(realizeInsn)
+
+export const realizeInsn = (insn: AbsInsn2): Array<SM83Insn> => insnReduce2(realizeCopy, realizeUnary, insn)
+
+const getSourceName = (source: Operand): string => typeof source === "number"
+    ? source.toString()
+    : source.asmName()
+
+export const realizeCopy = (dest: Loc, source: Operand): Array<SM83Insn> => [
+    `LD ${dest.asmName()}, ${getSourceName(source)}`
+]
+
+export const realizeUnary = (dest: Loc, op: UnaryOp | BinaryOp, source: Operand): Array<SM83Insn> => {
+    if (!isA(dest)) {
+        throw new Error("unexpected form for unary operation")
+    }
+
+    switch (op) {
+        case "add":
+            return [`ADD ${getSourceName(source)}`]
+
+        case "arith_negate":
+            return ["CPL", "INC A"]
+
+        case "bit_negate":
+            return ["CPL"]
+
+        default:
+            throw new Error(`unsupported unary op \`${op}'`)
+    }
+}
+
+const isA = (loc: Operand): boolean =>
+    typeof loc === "object" && loc.type === LocType.REGISTER && loc.name === R8.A