[compiler] Move co-mutation range extension to InferMutableRanges

compiler/packages/babel-plugin-react-compiler/src/HIR/PrintHIR.ts

@@ -731,7 +731,7 @@ function isMutable(range: MutableRange): boolean {
 }
 
 const DEBUG_MUTABLE_RANGES = false;
-function printMutableRange(identifier: Identifier): string {
+export function printMutableRange(identifier: Identifier): string {
   if (DEBUG_MUTABLE_RANGES) {
     // if debugging, print both the identifier and scope range if they differ
     const range = identifier.mutableRange;

compiler/packages/babel-plugin-react-compiler/src/Inference/AnalyseFunctions.ts

@@ -25,8 +25,6 @@ import {
   eachInstructionLValue,
   eachInstructionValueOperand,
 } from '../HIR/visitors';
-import prettyFormat from 'pretty-format';
-import {printIdentifier} from '../HIR/PrintHIR';
 import {Iterable_some} from '../Utils/utils';
 
 export default function analyseFunctions(func: HIRFunction): void {
@@ -69,16 +67,6 @@ function lower(func: HIRFunction): DisjointSet<Identifier> {
   return aliases;
 }
 
-export function debugAliases(aliases: DisjointSet<Identifier>): void {
-  console.log(
-    prettyFormat(
-      aliases
-        .buildSets()
-        .map(set => [...set].map(ident => printIdentifier(ident))),
-    ),
-  );
-}
-
 /**
  * The alias sets returned by InferMutableRanges() accounts only for aliases that
  * are known to mutate together. Notably this skips cases where a value is captured

compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRanges.ts

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import prettyFormat from 'pretty-format';
 import {HIRFunction, Identifier} from '../HIR/HIR';
 import DisjointSet from '../Utils/DisjointSet';
 import {inferAliasForUncalledFunctions} from './InerAliasForUncalledFunctions';
@@ -14,7 +15,9 @@ import {inferAliasForPhis} from './InferAliasForPhis';
 import {inferAliasForStores} from './InferAliasForStores';
 import {inferMutableLifetimes} from './InferMutableLifetimes';
 import {inferMutableRangesForAlias} from './InferMutableRangesForAlias';
+import {inferMutableRangesForComutation} from './InferMutableRangesForComutation';
 import {inferTryCatchAliases} from './InferTryCatchAliases';
+import {printIdentifier, printMutableRange} from '../HIR/PrintHIR';
 
 export function inferMutableRanges(ir: HIRFunction): DisjointSet<Identifier> {
   // Infer mutable ranges for non fields
@@ -32,11 +35,13 @@ export function inferMutableRanges(ir: HIRFunction): DisjointSet<Identifier> {
    * Eagerly canonicalize so that if nothing changes we can bail out
    * after a single iteration
    */
-  let prevAliases: Map<Identifier, Identifier> = aliases.canonicalize();
+  let prevAliases: Map<Identifier, string> = canonicalize(aliases);
   while (true) {
     // Infer mutable ranges for aliases that are not fields
     inferMutableRangesForAlias(ir, aliases);
 
+    inferMutableRangesForComutation(ir);
+
     // Update aliasing information of fields
     inferAliasForStores(ir, aliases);
 
@@ -45,7 +50,7 @@ export function inferMutableRanges(ir: HIRFunction): DisjointSet<Identifier> {
     // Update aliasing information of phis
     inferAliasForPhis(ir, aliases);
 
-    const nextAliases = aliases.canonicalize();
+    const nextAliases = canonicalize(aliases);
     if (areEqualMaps(prevAliases, nextAliases)) {
       break;
     }
@@ -77,12 +82,13 @@ export function inferMutableRanges(ir: HIRFunction): DisjointSet<Identifier> {
    * but does not modify values that `y` "contains" such as the
    * object literal or `z`.
    */
-  prevAliases = aliases.canonicalize();
+  prevAliases = canonicalize(aliases);
   while (true) {
     inferMutableRangesForAlias(ir, aliases);
+    inferMutableRangesForComutation(ir);
     inferAliasForPhis(ir, aliases);
     inferAliasForUncalledFunctions(ir, aliases);
-    const nextAliases = aliases.canonicalize();
+    const nextAliases = canonicalize(aliases);
     if (areEqualMaps(prevAliases, nextAliases)) {
       break;
     }
@@ -92,7 +98,41 @@ export function inferMutableRanges(ir: HIRFunction): DisjointSet<Identifier> {
   return aliases;
 }
 
-function areEqualMaps<T>(a: Map<T, T>, b: Map<T, T>): boolean {
+export function debugAliases(aliases: DisjointSet<Identifier>): void {
+  console.log(
+    prettyFormat(
+      aliases
+        .buildSets()
+        .map(set =>
+          [...set].map(
+            ident => printIdentifier(ident) + printMutableRange(ident),
+          ),
+        ),
+    ),
+  );
+}
+
+/**
+ * Canonicalizes the alias set and mutable range information calculated at the current time.
+ * The returned value maps each identifier in the program to the root identifier of its alias
+ * set and the the mutable range of that set.
+ *
+ * This ensures that we fixpoint the mutable ranges themselves and not just the alias sets.
+ */
+function canonicalize(
+  aliases: DisjointSet<Identifier>,
+): Map<Identifier, string> {
+  const entries = new Map<Identifier, string>();
+  aliases.forEach((item, root) => {
+    entries.set(
+      item,
+      `${root.id}:${root.mutableRange.start}:${root.mutableRange.end}`,
+    );
+  });
+  return entries;
+}
+
+function areEqualMaps<T, U>(a: Map<T, U>, b: Map<T, U>): boolean {
   if (a.size !== b.size) {
     return false;
   }

compiler/packages/babel-plugin-react-compiler/src/Inference/InferMutableRangesForComutation.ts

@@ -0,0 +1,91 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {
+  HIRFunction,
+  Identifier,
+  isRefOrRefValue,
+  makeInstructionId,
+} from '../HIR';
+import {eachInstructionOperand} from '../HIR/visitors';
+import {isMutable} from '../ReactiveScopes/InferReactiveScopeVariables';
+
+/**
+ * Finds instructions with operands that co-mutate and extends all their mutable ranges
+ * to end at the same point (the highest `end` value of the group). Note that the
+ * alias sets used in InferMutableRanges are meant for values that strictly alias:
+ * a mutation of one value in the set would directly modify the same object as some
+ * other value in the set.
+ *
+ * However, co-mutation can cause an alias to one object to be stored within another object,
+ * for example:
+ *
+ * ```
+ * const a = {};
+ * const b = {};
+ * const f = () => b.c; //
+ * setProperty(a, 'b', b); // equiv to a.b = b
+ *
+ * a.b.c = 'c'; // this mutates b!
+ * ```
+ *
+ * Here, the co-mutation in `setProperty(a, 'b', b)` means that a reference to b may be stored
+ * in a, vice-versa, or both. We need to extend the mutable range of both a and b to reflect
+ * the fact the values may mutate together.
+ *
+ * Previously this was implemented in InferReactiveScopeVariables, but that is too late:
+ * we need this to be part of the InferMutableRanges fixpoint iteration to account for functions
+ * like `f` in the example, which capture a reference to a value that may change later. `f`
+ * cannot be independently memoized from the `setProperty()` call due to the co-mutation.
+ *
+ * See aliased-capture-mutate and aliased-capture-aliased-mutate fixtures for examples.
+ */
+export function inferMutableRangesForComutation(fn: HIRFunction): void {
+  for (const block of fn.body.blocks.values()) {
+    for (const instr of block.instructions) {
+      let operands: Array<Identifier> | null = null;
+      for (const operand of eachInstructionOperand(instr)) {
+        if (
+          isMutable(instr, operand) &&
+          operand.identifier.mutableRange.start > 0
+        ) {
+          if (
+            instr.value.kind === 'FunctionExpression' ||
+            instr.value.kind === 'ObjectMethod'
+          ) {
+            if (operand.identifier.type.kind === 'Primitive') {
+              continue;
+            }
+          }
+          operands ??= [];
+          operands.push(operand.identifier);
+        }
+      }
+      if (operands != null) {
+        // Find the last instruction which mutates any of the mutable operands
+        let lastMutatingInstructionId = makeInstructionId(0);
+        for (const id of operands) {
+          if (id.mutableRange.end > lastMutatingInstructionId) {
+            lastMutatingInstructionId = id.mutableRange.end;
+          }
+        }
+
+        /**
+         * Update all mutable operands's mutable ranges to end at the same point
+         */
+        for (const id of operands) {
+          if (
+            id.mutableRange.end < lastMutatingInstructionId &&
+            !isRefOrRefValue(id)
+          ) {
+            id.mutableRange.end = lastMutatingInstructionId;
+          }
+        }
+      }
+    }
+  }
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/aliased-capture-aliased-mutate.expect.md

@@ -0,0 +1,71 @@
+
+## Input
+
+```javascript
+// @flow @enableTransitivelyFreezeFunctionExpressions:false
+import {arrayPush, setPropertyByKey, Stringify} from 'shared-runtime';
+
+function useFoo({a, b}: {a: number, b: number}) {
+  const x = [];
+  const y = {value: a};
+
+  arrayPush(x, y); // x and y co-mutate
+  const y_alias = y;
+  const cb = () => y_alias.value;
+  setPropertyByKey(x[0], 'value', b); // might overwrite y.value
+  return <Stringify cb={cb} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{a: 2, b: 10}],
+  sequentialRenders: [
+    {a: 2, b: 10},
+    {a: 2, b: 11},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { arrayPush, setPropertyByKey, Stringify } from "shared-runtime";
+
+function useFoo(t0) {
+  const $ = _c(3);
+  const { a, b } = t0;
+  let t1;
+  if ($[0] !== a || $[1] !== b) {
+    const x = [];
+    const y = { value: a };
+
+    arrayPush(x, y);
+    const y_alias = y;
+    const cb = () => y_alias.value;
+    setPropertyByKey(x[0], "value", b);
+    t1 = <Stringify cb={cb} shouldInvokeFns={true} />;
+    $[0] = a;
+    $[1] = b;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{ a: 2, b: 10 }],
+  sequentialRenders: [
+    { a: 2, b: 10 },
+    { a: 2, b: 11 },
+  ],
+};
+
+```
+
+### Eval output
+(kind: ok) <div>{"cb":{"kind":"Function","result":10},"shouldInvokeFns":true}</div>
+<div>{"cb":{"kind":"Function","result":11},"shouldInvokeFns":true}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/aliased-capture-aliased-mutate.js

@@ -0,0 +1,22 @@
+// @flow @enableTransitivelyFreezeFunctionExpressions:false
+import {arrayPush, setPropertyByKey, Stringify} from 'shared-runtime';
+
+function useFoo({a, b}: {a: number, b: number}) {
+  const x = [];
+  const y = {value: a};
+
+  arrayPush(x, y); // x and y co-mutate
+  const y_alias = y;
+  const cb = () => y_alias.value;
+  setPropertyByKey(x[0], 'value', b); // might overwrite y.value
+  return <Stringify cb={cb} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: useFoo,
+  params: [{a: 2, b: 10}],
+  sequentialRenders: [
+    {a: 2, b: 10},
+    {a: 2, b: 11},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-aliased-capture-mutate.expect.md → compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/aliased-capture-mutate.expect.md

@@ -5,19 +5,6 @@
 // @flow @enableTransitivelyFreezeFunctionExpressions:false
 import {setPropertyByKey, Stringify} from 'shared-runtime';
 
-/**
- * Variation of bug in `bug-aliased-capture-aliased-mutate`
- * Found differences in evaluator results
- * Non-forget (expected):
- *   (kind: ok)
- *   <div>{"cb":{"kind":"Function","result":2},"shouldInvokeFns":true}</div>
- *   <div>{"cb":{"kind":"Function","result":3},"shouldInvokeFns":true}</div>
- * Forget:
- *   (kind: ok)
- *   <div>{"cb":{"kind":"Function","result":2},"shouldInvokeFns":true}</div>
- *   <div>{"cb":{"kind":"Function","result":2},"shouldInvokeFns":true}</div>
- */
-
 function useFoo({a}: {a: number, b: number}) {
   const arr = [];
   const obj = {value: a};
@@ -46,7 +33,7 @@ import { c as _c } from "react/compiler-runtime";
 import { setPropertyByKey, Stringify } from "shared-runtime";
 
 function useFoo(t0) {
-  const $ = _c(4);
+  const $ = _c(2);
   const { a } = t0;
   let t1;
   if ($[0] !== a) {
@@ -55,15 +42,7 @@ function useFoo(t0) {
 
     setPropertyByKey(obj, "arr", arr);
     const obj_alias = obj;
-    let t2;
-    if ($[2] !== obj_alias.arr.length) {
-      t2 = () => obj_alias.arr.length;
-      $[2] = obj_alias.arr.length;
-      $[3] = t2;
-    } else {
-      t2 = $[3];
-    }
-    const cb = t2;
+    const cb = () => obj_alias.arr.length;
     for (let i = 0; i < a; i++) {
       arr.push(i);
     }
@@ -84,4 +63,7 @@ export const FIXTURE_ENTRYPOINT = {
 };
 
 ```
-
+
+### Eval output
+(kind: ok) <div>{"cb":{"kind":"Function","result":2},"shouldInvokeFns":true}</div>
+<div>{"cb":{"kind":"Function","result":3},"shouldInvokeFns":true}</div>