diff --git a/packages/compiler/src/backend/llvm/shape-union-tags.test.ts b/packages/compiler/src/backend/llvm/shape-union-tags.test.ts new file mode 100644 index 000000000..e07626e77 --- /dev/null +++ b/packages/compiler/src/backend/llvm/shape-union-tags.test.ts @@ -0,0 +1,80 @@ +import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { expect, test } from "vitest"; +import { compile, deserializeModule } from "../../index.js"; +import { validateModule } from "../../ir/validate.js"; +import { emitLlvmModule } from "./emitter.js"; + +async function llvmFor(source: string): Promise { + const dir = await mkdtemp(join(tmpdir(), "scriptc-shape-union-tags-")); + try { + const entry = join(dir, "main.ts"), + output = join(dir, "main.ir.json"); + await writeFile(entry, source); + const result = await compile(entry, { outDir: dir, outPath: output, outputKind: "ir" }); + if (!result.ok) + throw new Error(result.diagnostics.map((item) => `${item.code}: ${item.message}`).join("\n")); + const module = deserializeModule(await readFile(output, "utf8")); + expect(validateModule(module)).toEqual([]); + return emitLlvmModule(module); + } finally { + await rm(dir, { recursive: true, force: true }); + } +} + +function tagHelpers(ir: string): string[] { + return [...ir.matchAll(/^define internal i32 @sc_stag\.[^\n]*\{[^]*?^}/gm)].map( + (match) => match[0], + ); +} + +test("a union whose arms span few shape numbers switches on the shape word", async () => { + const ir = await llvmFor(` +type Num = { kind: "num"; value: number }; +type Neg = { kind: "neg"; operand: Node }; +type Add = { kind: "add"; left: Node; right: Node }; +type Node = Num | Neg | Add; +function evaluate(node: Node): number { + switch (node.kind) { + case "num": return node.value; + case "neg": return -evaluate(node.operand); + case "add": return evaluate(node.left) + evaluate(node.right); + } +} +const one: Node = { kind: "num", value: 1 }; +console.log(evaluate({ kind: "add", left: one, right: { kind: "neg", operand: one } })); +`); + expect(ir).not.toMatch(/@sc_stagtab\./); + const helpers = tagHelpers(ir); + expect(helpers.length).toBeGreaterThan(0); + for (const helper of helpers) expect(helper).toMatch(/switch i32 %w/); +}); + +test("a union whose arms span many shape numbers reads the arm from its table", async () => { + const arms = Array.from({ length: 40 }, (_, i) => i); + const ir = await llvmFor(` +${arms.map((i) => `type K${i} = { kind: "k${i}"; v${i}: number };`).join("\n")} +type Any = ${arms.map((i) => `K${i}`).join(" | ")}; +function value(item: Any): number { + switch (item.kind) { +${arms.map((i) => ` case "k${i}": return item.v${i};`).join("\n")} + } +} +const items: Any[] = [${arms.map((i) => `{ kind: "k${i}", v${i}: ${i} }`).join(", ")}]; +let total = 0; +for (const item of items) total += value(item); +console.log(total); +`); + // One byte per shape number from the arms' lowest to highest; 255 names no arm. + const table = + /^(@sc_stagtab\.\S+) = private unnamed_addr constant \[(\d+) x i8\] \[([^\]]*)\]$/m.exec(ir); + expect(table, "arm table").not.toBeNull(); + const entries = table![3]!.split(", ").map((entry) => Number(entry.replace("i8 ", ""))); + expect(entries).toHaveLength(Number(table![2])); + expect(entries.filter((entry) => entry !== 255).sort((a, b) => a - b)).toEqual(arms); + const helper = tagHelpers(ir).find((text) => text.includes(table![1]!)); + expect(helper, "tag helper").toBeDefined(); + expect(helper).toContain(`getelementptr inbounds [${table![2]} x i8], ptr ${table![1]}`); + expect(helper).not.toMatch(/switch i32 %w/); +}); diff --git a/packages/compiler/src/backend/llvm/shapes.ts b/packages/compiler/src/backend/llvm/shapes.ts index 6d03db222..79a79f36c 100644 --- a/packages/compiler/src/backend/llvm/shapes.ts +++ b/packages/compiler/src/backend/llvm/shapes.ts @@ -313,19 +313,46 @@ function shapeUnionRc(host: ShapeHost, union: ShapeUnion): { retain: string; rel }; } +/** The widest span of shape numbers a union's arms may cover for its tag + * helper to switch on the shape word (emitShapeUnionHelpers). */ +const SHAPE_TAG_SWITCH_SPAN = 32; + /** The tag, release and trace helpers of every shape union * (nullable-unions.ts). The tag helper maps a value to its arm: NULL and * the null sentinel are the unit arms, anything else a tagged record whose - * shape word names its arm (the switch folds to a lookup table). */ + * shape word names its arm. + * - Arms within SHAPE_TAG_SWITCH_SPAN shape numbers: the helper switches on + * the shape word. Inlined, the switch folds into the caller's branch on + * the arm, one indirect jump with a jump table of at most that many + * entries per site. + * - Wider unions: the shape word indexes the union's arm table, one entry + * per shape number from the arms' lowest to highest, and callers branch + * on the small arm tag. Folded switches would give every site of such a + * union a jump table spanning all of its arms' shape numbers. */ export function emitShapeUnionHelpers(host: ShapeHost): string[] { const out: string[] = []; const S = host.sizeType; for (const union of host.nullableUnions.allShapeUnions()) { const header = shapeWordHeader(S); const tagFn = shapeTagFn(union.unionId); - const cases = union.arms - .map((arm) => `i32 ${host.nullableUnions.shapeNumber(arm.shapeId)}, label %a${arm.tag}`) - .join(" "); + const numbers = union.arms.map((arm) => host.nullableUnions.shapeNumber(arm.shapeId)); + const lo = Math.min(...numbers); + const span = Math.max(...numbers) - lo + 1; + const byTable = span > SHAPE_TAG_SWITCH_SPAN; + // Entries are bytes, or 16-bit words for a union of 255 arms or more. + // A shape number that names no arm maps to the all-ones entry, which is + // above every arm tag, so callers take their path for no arm as they do + // for the -1 of a shape number outside the table. + const wide = union.arms.some((arm) => arm.tag >= 255); + const entry = wide ? "i16" : "i8"; + const none = wide ? 65535 : 255; + const table = new Array(span).fill(none); + union.arms.forEach((arm, i) => (table[numbers[i]! - lo] = arm.tag)); + const tableName = tagFn.replace("@sc_stag.", "@sc_stagtab."); + if (byTable) + out.push( + `${tableName} = private unnamed_addr constant [${span} x ${entry}] [${table.map((t) => `${entry} ${t}`).join(", ")}]`, + ); out.push( `define internal i32 ${tagFn}(ptr %u) alwaysinline ${FN_ATTRS} { ; shape union tag`, `entry:`, @@ -343,16 +370,35 @@ export function emitShapeUnionHelpers(host: ShapeHost): string[] { ` ret i32 ${union.nullTag}`, `rec:`, ); - out.push( - ` %wp = getelementptr inbounds ${header}, ptr %u, i32 0, i32 ${SHAPE_WORD_INDEX}`, - ` %w = load i32, ptr %wp`, - ` switch i32 %w, label %bad [ ${cases} ]`, - ...union.arms.flatMap((arm) => [`a${arm.tag}:`, ` ret i32 ${arm.tag}`]), - `bad:`, - ` ret i32 -1`, - `}`, - ``, - ); + const cases = union.arms.map((arm, i) => `i32 ${numbers[i]}, label %a${arm.tag}`).join(" "); + if (!byTable) + out.push( + ` %wp = getelementptr inbounds ${header}, ptr %u, i32 0, i32 ${SHAPE_WORD_INDEX}`, + ` %w = load i32, ptr %wp`, + ` switch i32 %w, label %bad [ ${cases} ]`, + ...union.arms.flatMap((arm) => [`a${arm.tag}:`, ` ret i32 ${arm.tag}`]), + `bad:`, + ` ret i32 -1`, + `}`, + ``, + ); + else + out.push( + ` %wp = getelementptr inbounds ${header}, ptr %u, i32 0, i32 ${SHAPE_WORD_INDEX}`, + ` %w = load i32, ptr %wp`, + ` %i = sub i32 %w, ${lo}`, + ` %inrange = icmp ult i32 %i, ${span}`, + ` br i1 %inrange, label %look, label %bad`, + `look:`, + ` %tp = getelementptr inbounds [${span} x ${entry}], ptr ${tableName}, i32 0, i32 %i`, + ` %t8 = load ${entry}, ptr %tp`, + ` %tag = zext ${entry} %t8 to i32`, + ` ret i32 %tag`, + `bad:`, + ` ret i32 -1`, + `}`, + ``, + ); // Release and trace: one definition for the unions over the same // record arms (helperOwner). Unit encodings are NULL or the immortal // null sentinel; the shape word picks the arm.