Keep Map.get results owned when a copy or a temporary map can outlive the entry

A borrowed Map.get local holds the entry's own reference, so every use
must run before anything that can release the entry. Two cases escaped
that proof:

- A local declared from the borrowed local (`const copy = held`) is a
  stable alias of it, with no reference of its own. Its reads after a
  write to the map were not counted as uses. Reference locals declared
  from the local, at any depth of the statement that declares them, now
  count as its uses.
- A read from a map that no local holds (`fresh().get(k)`) borrowed an
  entry of a map released at the end of the declaration. The map must
  now be a local or a field chain rooted at one, whose changes the proof
  sees as edge changes.

Both keep the owned read. Corpus 8181 covers the two cases, and the
emission test checks that their reads retain the value.

packages/compiler/src/backend/llvm/map-read-lifetimes.ts

@@ -166,6 +166,17 @@ export function findMapReadLifetimes(
   return result;
 }
 
+/** Whether a map read's receiver keeps the map past its own statement: a
+ * local or a field of one. Changing either is an edge change the borrow
+ * proof sees. A temporary map (a call's result, `new Map()`) is released at
+ * the end of the declaration, and its entries with it. */
+function heldReceiver(expr: IrExpr): boolean {
+  if (expr.kind === "varRef") return true;
+  if (expr.kind === "fieldGet" || expr.kind === "recordGet") return heldReceiver(expr.obj);
+  if (expr.kind === "unionNarrow" || expr.kind === "downcast") return heldReceiver(expr.value);
+  return false;
+}
+
 /** The locals, declared by statements `candidate` accepts, that can hold a
  * map entry's own reference instead of a +1 copy. The statements after the
  * declaration run in order, and every one that names the local must
@@ -174,7 +185,9 @@ export function findMapReadLifetimes(
  * while the local is read. An `if` passes its condition's state to both
  * branches and joins them; any other statement that names the local is
  * judged whole. A later pass of an enclosing loop declares the local again,
- * so only the rest of its own statement list matters. An update such as
+ * so only the rest of its own statement list matters. The map must be held
+ * by a local or a field (heldReceiver), and locals declared from the local
+ * count as its uses. An update such as
  * `const c = m.get(k); if (c === undefined) m.set(k, v); else c.n++` keeps
  * the borrow: the write runs only on the branch that does not read `c`. */
 export function findMapReadBorrows(
@@ -184,12 +197,14 @@ export function findMapReadBorrows(
 ): Set<string> {
   const result = new Set<string>();
   if (fn.async || fn.generator) return result;
-  const names = (node: IrStmt | IrExpr, id: string, isStmt: boolean): boolean => {
+  // `ids` holds the local and the locals declared from it: a stable alias
+  // (`const d = c`) reads the same entry without a reference of its own.
+  const names = (node: IrStmt | IrExpr, ids: ReadonlySet<string>, isStmt: boolean): boolean => {
     let found = false;
     const visitor = {
       stmt: () => !found,
       expr: (expr: IrExpr) => {
-        if (expr.kind === "varRef" && expr.localId === id) found = true;
+        if (expr.kind === "varRef" && ids.has(expr.localId)) found = true;
         return !found;
       },
     };
@@ -197,28 +212,48 @@ export function findMapReadBorrows(
     else everyExpr(node as IrExpr, visitor);
     return found;
   };
+  // Adds the reference locals that `stmt` declares from one in `ids`, at any
+  // depth: a `var` inside a loop or a try stays in scope after it.
+  const alias = (stmt: IrStmt, ids: Set<string>): void => {
+    everyStmt(stmt, {
+      stmt: (node) => {
+        if (
+          node.kind === "varDecl" &&
+          node.init !== null &&
+          isRefCounted(node.init.type) &&
+          names(node.init, ids, false)
+        )
+          ids.add(node.localId);
+        return true;
+      },
+      expr: () => true,
+    });
+  };
   // Whether an edge may have changed after `list`, or null when a statement
   // names the local after such a change (or changes edges itself).
-  const walk = (list: readonly IrStmt[], id: string, changed: boolean): boolean | null => {
+  const walk = (list: readonly IrStmt[], ids: Set<string>, changed: boolean): boolean | null => {
     for (const stmt of list) {
       if (stmt.kind === "if") {
         const preserving = effects.preserves(stmt.cond);
-        if (names(stmt.cond, id, false) && (changed || !preserving)) return null;
+        if (names(stmt.cond, ids, false) && (changed || !preserving)) return null;
         const condition = changed || !preserving;
-        const then = walk(stmt.then, id, condition);
-        const otherwise = stmt.else_ === null ? condition : walk(stmt.else_, id, condition);
+        const then = walk(stmt.then, ids, condition);
+        const otherwise = stmt.else_ === null ? condition : walk(stmt.else_, ids, condition);
         if (then === null || otherwise === null) return null;
         changed = then || otherwise;
         continue;
       }
       if (stmt.kind === "block") {
-        const inner = walk(stmt.body, id, changed);
+        const inner = walk(stmt.body, ids, changed);
         if (inner === null) return null;
         changed = inner;
         continue;
       }
       const preserving = effects.preservesScope([stmt]);
-      if (names(stmt, id, true) && (changed || !preserving)) return null;
+      if (names(stmt, ids, true)) {
+        if (changed || !preserving) return null;
+        alias(stmt, ids);
+      }
       changed ||= !preserving;
     }
     return changed;
@@ -227,8 +262,10 @@ export function findMapReadBorrows(
     list.forEach((stmt, index) => {
       if (
         stmt.kind === "varDecl" &&
+        stmt.init?.kind === "mapIntrinsic" &&
+        heldReceiver(stmt.init.receiver) &&
         candidate(stmt) &&
-        walk(list.slice(index + 1), stmt.localId, false) !== null
+        walk(list.slice(index + 1), new Set([stmt.localId]), false) !== null
       )
         result.add(stmt.localId);
       for (const nested of nestedStmtLists(stmt)) visit(nested);

packages/compiler/test/map-get-borrow-emission.test.ts

@@ -32,9 +32,26 @@ function replaced(m: Map<number, Counter>): number {
   m.set(1, new Counter());
   return old === undefined ? 0 : old.num;
 }
+function aliased(m: Map<number, Counter>): number {
+  const held = m.get(1);
+  const copy = held;
+  m.set(1, new Counter());
+  return copy === undefined ? 0 : copy.num;
+}
+function fresh(): Map<number, Counter> {
+  return new Map([[1, new Counter()]]);
+}
+function temporary(keys: number[]): number {
+  let total = 0;
+  for (const key of keys) {
+    const held = fresh().get(key);
+    if (held !== undefined) total += held.num;
+  }
+  return total;
+}
 const numbers = new Map<number, Counter>();
 count(numbers, [1, 2, 1]);
-console.log(named(new Map([["a", new Counter()]]), "a"), replaced(numbers));
+console.log(named(new Map([["a", new Counter()]]), "a"), replaced(numbers), aliased(numbers), temporary([1, 2]));
 `,
     );
     const result = await compile(entry, { outDir: dir, outPath, outputKind: "ir" });
@@ -50,8 +67,11 @@ console.log(named(new Map([["a", new Counter()]]), "a"), replaced(numbers));
       expect(fn(name)).not.toContain("@sc_retain_Counter");
       expect(fn(name)).not.toContain("@sc_release_Counter");
     }
-    // A write between the read and its use keeps the read's own reference.
-    expect(fn("replaced")).toContain("@sc_retain_Counter");
+    // A write between the read and its use keeps the read's own reference,
+    // also when the use goes through a local declared from the read, and so
+    // does a read from a map that is released after the declaration.
+    for (const name of ["replaced", "aliased", "temporary"])
+      expect(fn(name)).toContain("@sc_retain_Counter");
   } finally {
     await rm(dir, { recursive: true, force: true });
   }

tests/corpus/8181-map-get-borrows.ts

@@ -1,7 +1,8 @@
 // Map.get results held in locals across the operations that can and cannot
 // release the entry's value: an update on the absent branch, a write, delete
 // or clear before a later read, a call that may change the map, loops,
-// string keys, and values that are stored, returned or compared.
+// string keys, values that are stored, returned or compared, a local declared
+// from the read, and a read from a temporary map.
 class Counter {
   num: number;
   label: string;
@@ -128,3 +129,27 @@ function loopWrites(m: Map<number, Counter>): number {
   return total;
 }
 console.log("loop", loopWrites(new Map([[0, new Counter(1, "a")], [1, new Counter(2, "b")]])));
+
+// A local declared from the read is read after a write that replaces the
+// entry; it reads the value the map held when the local was declared.
+function aliasThenWrite(m: Map<number, Counter>): string {
+  const held = m.get(1);
+  const copy = held;
+  m.set(1, new Counter(10, "ten"));
+  m.delete(1);
+  const filler = [new Counter(11, "filler"), new Counter(12, "filler")];
+  return copy === undefined ? "none" : `${copy.num}:${copy.label}:${filler.length}`;
+}
+console.log("alias", aliasThenWrite(new Map([[1, new Counter(1, "one")]])));
+
+// A read from a map that no local holds: the map is released after the
+// declaration, and the value stays readable.
+function freshMap(): Map<number, Counter> {
+  return new Map([[1, new Counter(5, "five")]]);
+}
+function fromTemporary(): string {
+  const held = freshMap().get(1);
+  const filler = [new Counter(6, "filler"), new Counter(7, "filler")];
+  return held === undefined ? "none" : `${held.num}:${held.label}:${filler.length}`;
+}
+console.log("temporary", fromTemporary());