Trace Array subclass instances in array-typed slots and test their class on element reads (#2372)

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

@@ -1097,9 +1097,15 @@ export function emitClassShapes(
     // A traced class whose trace visits nothing carries the collector header
     // only because its hierarchy is traced: it is never on a cycle. Its
     // descriptor names scr_cyc_no_children and its objects start out
-    // SCR_CYC_ACYCLIC, so releases never buffer them.
+    // SCR_CYC_ACYCLIC, so releases never buffer them. An Array subclass's
+    // trace also visits its elements by their current element ops
+    // (scr_arr_subclass_trace), so it is never acyclic.
     const acyclic =
-      traced && runtimePrefix === null && !isStreamRooted && tracedFields.length === 0;
+      traced &&
+      runtimePrefix === null &&
+      !isStreamRooted &&
+      tracedFields.length === 0 &&
+      cls.arrayElem === undefined;
     const bounded =
       runtimePrefix !== null ||
       isStreamRooted ||

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

@@ -2355,6 +2355,39 @@ export class LlEmitter {
       programLifetime: this.programLifetime,
       regionAliasTag: (word) => this.memoryAliasAttachment(`region.${word}`),
     };
+    // An instance of a class that extends Array is an array: a slot of its
+    // element array type can hold it, and the collector visits it there as
+    // an array, which has a collector header when its element ops trace
+    // (scr_arr_make_ops_kind). Such a class carries a header too, and so
+    // does the rest of its hierarchy. A class whose elements are instances
+    // of another such class follows that class, so repeat until no class
+    // adds one.
+    const arrayClasses = (mod.classes ?? []).filter((c) => c.arrayElem !== undefined);
+    if (arrayClasses.length > 0) {
+      const baseOf = new Map((mod.classes ?? []).map((c) => [c.name, c.base] as const));
+      const rootOf = (name: string): string => {
+        let root = name;
+        for (
+          let base = baseOf.get(root);
+          base !== undefined && baseOf.has(base);
+          base = baseOf.get(root)
+        )
+          root = base;
+        return root;
+      };
+      for (let changed = true; changed;) {
+        changed = false;
+        for (const cls of arrayClasses) {
+          if (this.tracedShapes.has(`object:${cls.name}`)) continue;
+          if (traceAdapter(this.shapeHost, { kind: "array", elem: cls.arrayElem! }) === null)
+            continue;
+          const root = rootOf(cls.name);
+          for (const member of mod.classes ?? [])
+            if (rootOf(member.name) === root) this.tracedShapes.add(`object:${member.name}`);
+          changed = true;
+        }
+      }
+    }
     this.walkers = new LlWalkers(this.shapeHost);
     this.dyn = new LlDyn(this.shapeHost);
     for (const g of mod.globals ?? []) {

packages/compiler/src/frontend/lowering/builtins/array-classes.ts

@@ -10,6 +10,7 @@ import * as ts from "../../ts7/adapter.js";
 import { type Lowerer } from "../lowerer.js";
 import { isJsSourceFile, locOf } from "../../program.js";
 import { stripTypeWrappers } from "../typed-arguments.js";
+import { keyedHelperName } from "../stable-names.js";
 import { dynUndefinedExpr, nodeThrowExpr } from "../../../ir/build.js";
 import { lowerArrayFromCall, lowerArrayOfCall } from "../containers/array-construction.js";
 import { boundedArrayHofHelper } from "../containers/array-callbacks.js";
@@ -305,6 +306,7 @@ export function lowerArrayClassInstanceOf(
   const value = lowerer.lowerExpr(expr.left);
   if (value.type.kind === "array")
     return { kind: "instanceOf", value, className: info.def.name, type: BOOL, loc };
+  if (value.type.kind === "union") return arrayClassUnionInstanceOf(lowerer, value, info, loc);
   if (value.type.kind === "dyn")
     return {
       kind: "libCall",
@@ -324,6 +326,62 @@ export function lowerArrayClassInstanceOf(
   return null;
 }
 
+/** `x instanceof X` for a class X declared `extends Array` on a union (an
+ * element read of an array of X instances is `X | undefined`): an array arm
+ * tests the array's class at run time, a checked-dynamic arm asks the
+ * runtime, and every other arm is no instance. */
+function arrayClassUnionInstanceOf(
+  lowerer: Lowerer,
+  value: IrExpr,
+  info: ClassInfo,
+  loc: SrcLoc,
+): IrExpr | null {
+  if (value.type.kind !== "union") return null;
+  const unionId = value.type.unionId;
+  const arms = lowerer.unions.get(unionId)?.arms;
+  if (arms === undefined) return null;
+  const className = info.def.name;
+  const key = `instanceof.arrayClass:${unionId}:${className}`;
+  let helper = lowerer.valueHelpers.get(key);
+  if (!helper) {
+    helper = keyedHelperName(lowerer, "%instanceof.arrayClass", key);
+    lowerer.valueHelpers.set(key, helper);
+    const param: IrExpr = { kind: "varRef", localId: "value.0", type: value.type, loc };
+    const body: IrStmt[] = [];
+    arms.forEach((arm, tag) => {
+      if (arm.kind !== "array" && arm.kind !== "dyn") return;
+      const narrowed: IrExpr = { kind: "unionNarrow", unionId, tag, value: param, type: arm, loc };
+      const test: IrExpr =
+        arm.kind === "array"
+          ? { kind: "instanceOf", value: narrowed, className, type: BOOL, loc }
+          : {
+              kind: "libCall",
+              fn: "array.dynIsClass",
+              args: [narrowed, { kind: "strLit", value: className, type: STRING, loc }],
+              type: BOOL,
+              loc,
+            };
+      body.push({
+        kind: "if",
+        cond: { kind: "unionIsTag", unionId, tag, negated: false, value: param, type: BOOL, loc },
+        then: [{ kind: "return", value: test, loc }],
+        else_: null,
+        loc,
+      });
+    });
+    body.push({ kind: "return", value: { kind: "boolLit", value: false, type: BOOL, loc }, loc });
+    lowerer.liftedFns.push({
+      name: helper,
+      params: [{ localId: "value.0", name: "value", type: value.type }],
+      returnType: BOOL,
+      locals: [{ id: "value.0", name: "value", type: value.type, mutable: false }],
+      body,
+      loc,
+    });
+  }
+  return { kind: "call", callee: helper, args: [value], type: BOOL, loc };
+}
+
 /** The class declared `extends Array` that a value of `type` is an instance
  * of, or null. */
 export function arrayClassOfType(lowerer: Lowerer, type: ts.Type): ClassInfo | null {

packages/compiler/src/frontend/lowering/expressions/relational.ts

@@ -840,9 +840,10 @@ export function lowerInstanceOf(lowerer: Lowerer, expr: ts.BinaryExpression, loc
   if (left.type.kind === "union") {
     const unionId = left.type.unionId;
     const arms = lowerer.unions.get(unionId)?.arms;
-    // Ordinary records, arrays and scalar values do not carry a class
-    // prototype (no program class extends Array). Dynamic/opaque runtime
-    // objects keep their existing instanceof fences.
+    // Records, arrays and scalar values are no instances of a class that
+    // does not extend Array (lowerArrayClassInstanceOf tests the classes
+    // that do). Dynamic/opaque runtime objects keep their existing
+    // instanceof fences.
     if (
       arms?.every(
         (arm) =>

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/array-subclass-element-cycles.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/array-subclass-element-cycles.ts"
+ ],
+ "diags": []
+}

packages/compiler/test/ts7/baselines/order-parity/tests/corpus/array-subclass-element-instanceof.ts.json

@@ -0,0 +1,6 @@
+{
+ "order": [
+  "<repo>/tests/corpus/array-subclass-element-instanceof.ts"
+ ],
+ "diags": []
+}

tests/corpus/array-subclass-element-cycles.ts

@@ -0,0 +1,60 @@
+// A class that extends Array whose elements can lead back to it only
+// through a field of plain array type (TreeNode.kids holds a Kids, typed as
+// an array of TreeNode): its instances carry the collector header that
+// arrays of its element type carry, so the collector can visit them through
+// that field. Forest's elements are Kids instances, which carry one too.
+class Kids extends Array<TreeNode> {
+  label = "kids";
+  total(): number {
+    let sum = 0;
+    for (const kid of this) sum += kid.size();
+    return sum;
+  }
+}
+
+class TreeNode {
+  kids = new Kids();
+  value: number;
+  constructor(value: number) {
+    this.value = value;
+  }
+  size(): number {
+    return 1 + this.kids.total();
+  }
+}
+
+class Forest extends Array<Kids> {
+  name = "forest";
+}
+
+function grow(depth: number): TreeNode {
+  const node = new TreeNode(depth);
+  if (depth > 0) for (let i = 0; i < 3; i++) node.kids.push(grow(depth - 1));
+  return node;
+}
+
+// Garbage cycles through the field of array type make the collector run
+// while the trees are alive.
+function churn(rounds: number): number {
+  let made = 0;
+  for (let i = 0; i < rounds; i++) {
+    const node = new TreeNode(i);
+    node.kids.push(node);
+    made++;
+  }
+  return made;
+}
+
+let checks = 0;
+for (let round = 0; round < 20; round++) {
+  checks += grow(4).size();
+  churn(500);
+}
+console.log(checks);
+
+const forest = new Forest();
+for (let i = 0; i < 4; i++) forest.push(grow(2).kids);
+churn(5000);
+let leaves = 0;
+for (const kids of forest) leaves += kids.total();
+console.log(forest.name, forest.length, leaves, forest[0]!.label, forest[0] instanceof Kids);

tests/corpus/array-subclass-element-instanceof.ts

@@ -0,0 +1,37 @@
+// Instances of classes that extend Array read back from an array of them
+// (an element read is `Row | undefined`) still answer instanceof for their
+// own class and Array, false for another Array subclass, and plain arrays
+// answer false.
+class Row extends Array<number> {
+  sum(): number {
+    let s = 0;
+    for (const v of this) s += v;
+    return s;
+  }
+}
+class Column extends Array<number> {}
+
+function rows(n: number): Row[] {
+  const out: Row[] = [];
+  for (let i = 0; i < n; i++) {
+    const row = new Row();
+    for (let j = 0; j <= i; j++) row.push(j);
+    out.push(row);
+  }
+  return out;
+}
+
+const lines: number[][] = [new Row(), new Column(), [1, 2]];
+
+const table = rows(6);
+console.log(table.length, table[5]!.sum(), table[5] instanceof Row, table[5] instanceof Column);
+console.log(table[4] instanceof Row, table[4] instanceof Array);
+console.log(table.at(-1) instanceof Row, table.at(10) instanceof Row);
+const plain: number[][] = [[1, 2]];
+console.log(plain[0] instanceof Row, plain[0] instanceof Array);
+const kinds = lines.map((r) => (r instanceof Column ? "column" : r instanceof Row ? "row" : "plain"));
+console.log(kinds.join(","));
+let checks = 0;
+for (let round = 0; round < 20; round++)
+  checks += rows(6).reduce((s, r) => s + (r instanceof Row ? r.sum() : 0), 0);
+console.log(checks);

tests/harness/corpus-costs/array-subclass-element-cycles.ts.json

@@ -0,0 +1 @@
+{ "plain": 3200, "san": 4800 }

tests/harness/corpus-costs/array-subclass-element-instanceof.ts.json

@@ -0,0 +1 @@
+{ "plain": 2500, "san": 3300 }