Leave the dispose list out of arena constructors when every arena is closed (#2393)

packages/compiler/src/backend/llvm/arenas.test.ts

@@ -0,0 +1,114 @@
+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 } from "../../index.js";
+
+async function llvmFor(source: string): Promise<string> {
+  const dir = await mkdtemp(join(tmpdir(), "scriptc-arenas-"));
+  try {
+    const entry = join(dir, "main.ts"),
+      output = join(dir, "main.ll");
+    await writeFile(entry, source);
+    const result = await compile(entry, { outDir: dir, outPath: output, outputKind: "llvm" });
+    if (!result.ok)
+      throw new Error(result.diagnostics.map((item) => `${item.code}: ${item.message}`).join("\n"));
+    return await readFile(output, "utf8");
+  } finally {
+    await rm(dir, { recursive: true, force: true });
+  }
+}
+
+/** The bodies of the module's arena-aware constructors. */
+function arenaConstructors(llvm: string): string[] {
+  return [...llvm.matchAll(/^define internal ptr @[^(]*\.r\(\)[^\n]*\n[\s\S]*?^}/gm)].map(
+    (m) => m[0],
+  );
+}
+
+const TREE = `
+class Node {
+  left: Node | null;
+  right: Node | null;
+  constructor(left: Node | null, right: Node | null) {
+    this.left = left;
+    this.right = right;
+  }
+  count(): number {
+    return 1 + (this.left === null ? 0 : this.left.count()) + (this.right === null ? 0 : this.right.count());
+  }
+}
+function tree(depth: number): Node {
+  return depth === 0 ? new Node(null, null) : new Node(tree(depth - 1), tree(depth - 1));
+}
+interface Leaf {
+  left: Leaf | null;
+  right: Leaf | null;
+}
+function leaves(depth: number): Leaf {
+  return depth === 0 ? { left: null, right: null } : { left: leaves(depth - 1), right: leaves(depth - 1) };
+}
+function size(leaf: Leaf): number {
+  return leaf.left === null || leaf.right === null ? 1 : 1 + size(leaf.left) + size(leaf.right);
+}
+`;
+
+test("when every arena is closed, arena-aware constructors take no dispose list", async () => {
+  const llvm = await llvmFor(`${TREE}
+let total = 0;
+for (let d = 0; d < 4; d++) total += tree(d).count() + size(leaves(d));
+console.log(total);
+`);
+  expect(llvm).toMatch(/@scr_arena_open\(ptr %[\w.]+, i32 1\)/);
+  expect(llvm).not.toMatch(/@scr_arena_open\(ptr %[\w.]+, i32 0\)/);
+  const constructors = arenaConstructors(llvm);
+  expect(constructors.length).toBe(2);
+  for (const body of constructors) {
+    expect(body).not.toContain("ScrArena.closed");
+    expect(body).not.toContain("@scr_arena_disposal");
+  }
+});
+
+test("an arena that is not closed keeps the dispose list in every arena-aware constructor", async () => {
+  // The rings hold strings and outside categories, so their groups are not closed.
+  const llvm = await llvmFor(`${TREE}
+class Category {
+  name: string;
+  constructor(name: string) {
+    this.name = name;
+  }
+}
+class Member {
+  next: Member | undefined = undefined;
+  category: Category;
+  label: string;
+  constructor(category: Category, label: string) {
+    this.category = category;
+    this.label = label;
+  }
+}
+function ring(categories: Category[], size: number, seed: number): Member {
+  const first = new Member(categories[seed % categories.length]!, \`m\${seed}\`);
+  let last = first;
+  for (let i = 1; i < size; i++) {
+    const member = new Member(categories[(seed + i) % categories.length]!, \`m\${seed}-\${i}\`);
+    last.next = member;
+    last = member;
+  }
+  last.next = first;
+  return first;
+}
+const categories = [new Category("a"), new Category("b")];
+let total = 0;
+for (let round = 0; round < 4; round++) {
+  const r = ring(categories, 30, round);
+  total += r.next!.label.length + r.category.name.length + tree(round).count();
+}
+console.log(total);
+`);
+  expect(llvm).toMatch(/@scr_arena_open\(ptr %[\w.]+, i32 1\)/);
+  expect(llvm).toMatch(/@scr_arena_open\(ptr %[\w.]+, i32 0\)/);
+  const withFields = arenaConstructors(llvm).filter((body) => body.includes(".dispose"));
+  expect(withFields.length).toBeGreaterThan(0);
+  for (const body of withFields) expect(body).toContain("ScrArena.closed");
+});

packages/compiler/src/backend/llvm/arenas.ts

@@ -47,6 +47,10 @@ export interface ArenaPlan {
   /** Record shapes and classes that get an arena-aware constructor. */
   shapes: Set<string>;
   classes: Set<string>;
+  /** Every root and call region opens a closed arena, so the current arena
+   * is always closed or absent, and the arena-aware constructors of record
+   * shapes and classes add no dispose entries. */
+  allClosed: boolean;
 }
 
 export interface ArenaEligibility {
@@ -150,7 +154,8 @@ export function planArenas(mod: IrModule, eligible: ArenaEligibility): ArenaPlan
   };
   const keptSite = (e: Node): boolean => sites.arenaSites.has(e) || groups.roots.has(e);
   for (const r of analyzeClosedGroups(mod, groups, arenaType, keptSite)) groups.closed.add(r);
-  return { groups, sites, shapes, classes };
+  const allClosed = [...groups.roots, ...groups.calls].every((r) => groups.closed.has(r));
+  return { groups, sites, shapes, classes, allClosed };
 }
 
 /** The functions in `reachable` that can allocate a cycle-capable object

packages/compiler/src/backend/llvm/classes.ts

@@ -1527,7 +1527,8 @@ export function emitClassShapes(
         ? `  store i64 ${(BigInt(meta.pre + 1) * 2n ** 32n + BigInt(ARENA_RC)).toString()}, ptr %o.vt`
         : `  store ${host.sizeType} ${ARENA_RC}, ptr %o`;
       const newSym = mangleClassNew(cls.name);
-      const dispose = refFields.length > 0 ? arenaDisposeSym(newSym) : null;
+      const dispose =
+        refFields.length > 0 && host.arenasClosed !== true ? arenaDisposeSym(newSym) : null;
       defs.push(
         ...arenaNewLines(
           host,

packages/compiler/src/backend/llvm/emitter.ts

@@ -2520,6 +2520,7 @@ export class LlEmitter {
       if (this.arenas !== null) {
         this.shapeHost.arenaShapes = this.arenas.shapes;
         this.shapeHost.arenaClasses = this.arenas.classes;
+        this.shapeHost.arenasClosed = this.arenas.allClosed;
       }
     }
     for (const [name, rec] of RUNTIME_ERROR_CLASSES) {

packages/compiler/src/backend/llvm/shapes.ts

@@ -152,6 +152,9 @@ export interface ShapeHost {
    * (arenas.ts). */
   arenaShapes?: ReadonlySet<string>;
   arenaClasses?: ReadonlySet<string>;
+  /** Every arena the program opens is closed (ArenaPlan.allClosed): the
+   * arena-aware constructors take no dispose list. */
+  arenasClosed?: boolean;
   /** The target C ABI's `size_t` / refcount integer type. */
   readonly sizeType: "i32" | "i64";
   /** Worker executables: the inline allocator state is thread-local. */
@@ -2347,7 +2350,8 @@ export function emitRecordShapes(
     defs.push(...nw);
     if (host.arenaShapes?.has(shape.id) === true && !shape.indexValue) {
       const newSym = mangleRecordNew(shape.id);
-      const dispose = refMembers.length > 0 ? arenaDisposeSym(newSym) : null;
+      const dispose =
+        refMembers.length > 0 && host.arenasClosed !== true ? arenaDisposeSym(newSym) : null;
       const init = [`  store ${host.sizeType} ${ARENA_RC}, ptr %o`];
       if (shapeNumber !== 0)
         init.push(