[compiler] Transitively freezing functions marks values as frozen, no…

…t effects

The fixture from the previous PR was getting inconsistent behavior because of the following:
1. Create an object in a useMemo
2. Create a callback in a useCallback, where the callback captures the object from (1) into a local object, then passes that local object into a logging method. We have to assume the logging method could modify the local object, and transitively, the object from (1).
3. Call the callback during render.
4. Pass the callback to JSX.

We correctly infer that the object from (1) is captured and modified in (2). However, in (4) we transitively freeze the callback. When transitively freezing functions we were previously doing two things: updating our internal abstract model of the program values to reflect the values as being frozen *and* also updating function operands to change their effects to freeze.

As the case above demonstrates, that can clobber over information about real potential mutability. The potential fix here is to only walk our abstract value model to mark values as frozen, but _not_ override operand effects. Conceptually, this is a forward data flow propagation — but walking backward to update effects is pushing information backwards in the algorithm. An alternative would be to mark that data was propagated backwards, and trigger another loop over the CFG to propagate information forward again given the updated effects. But the fix in this PR is more correct.

ghstack-source-id: c05e716f37827cb5515a059a1f0e8e8ff94b91df
Pull Request resolved: #30766

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

@@ -453,6 +453,37 @@ class InferenceState {
     }
   }
 
+  freezeValues(values: Set<InstructionValue>, reason: Set<ValueReason>): void {
+    for (const value of values) {
+      this.#values.set(value, {
+        kind: ValueKind.Frozen,
+        reason,
+        context: new Set(),
+      });
+      if (value.kind === 'FunctionExpression') {
+        if (
+          this.#env.config.enablePreserveExistingMemoizationGuarantees ||
+          this.#env.config.enableTransitivelyFreezeFunctionExpressions
+        ) {
+          if (value.kind === 'FunctionExpression') {
+            /*
+             * We want to freeze the captured values, not mark the operands
+             * themselves as frozen. There could be mutations that occur
+             * before the freeze we are processing, and it would be invalid
+             * to overwrite those mutations as a freeze.
+             */
+            for (const operand of eachInstructionValueOperand(value)) {
+              const operandValues = this.#variables.get(operand.identifier.id);
+              if (operandValues !== undefined) {
+                this.freezeValues(operandValues, reason);
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   reference(
     place: Place,
     effectKind: Effect,
@@ -482,29 +513,7 @@ class InferenceState {
             reason: reasonSet,
             context: new Set(),
           };
-          values.forEach(value => {
-            this.#values.set(value, {
-              kind: ValueKind.Frozen,
-              reason: reasonSet,
-              context: new Set(),
-            });
-
-            if (
-              this.#env.config.enablePreserveExistingMemoizationGuarantees ||
-              this.#env.config.enableTransitivelyFreezeFunctionExpressions
-            ) {
-              if (value.kind === 'FunctionExpression') {
-                for (const operand of eachInstructionValueOperand(value)) {
-                  this.referenceAndRecordEffects(
-                    operand,
-                    Effect.Freeze,
-                    ValueReason.Other,
-                    [],
-                  );
-                }
-              }
-            }
-          });
+          this.freezeValues(values, reasonSet);
         } else {
           effect = Effect.Read;
         }
@@ -1241,13 +1250,17 @@ function inferBlock(
       case 'ObjectMethod':
       case 'FunctionExpression': {
         let hasMutableOperand = false;
+        const mutableOperands: Array<Place> = [];
         for (const operand of eachInstructionOperand(instr)) {
           state.referenceAndRecordEffects(
             operand,
             operand.effect === Effect.Unknown ? Effect.Read : operand.effect,
             ValueReason.Other,
             [],
           );
+          if (isMutableEffect(operand.effect, operand.loc)) {
+            mutableOperands.push(operand);
+          }
           hasMutableOperand ||= isMutableEffect(operand.effect, operand.loc);
 
           /**

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-mutate-global-in-effect-fixpoint.expect.md

@@ -46,54 +46,57 @@ import { useEffect, useState } from "react";
 let someGlobal = { value: null };
 
 function Component() {
-  const $ = _c(6);
+  const $ = _c(7);
   const [state, setState] = useState(someGlobal);
-
-  let x = someGlobal;
-  while (x == null) {
-    x = someGlobal;
-  }
-
-  const y = x;
   let t0;
   let t1;
+  let t2;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = () => {
+    let x = someGlobal;
+    while (x == null) {
+      x = someGlobal;
+    }
+
+    const y = x;
+    t0 = useEffect;
+    t1 = () => {
       y.value = "hello";
     };
-    t1 = [];
+    t2 = [];
     $[0] = t0;
     $[1] = t1;
+    $[2] = t2;
   } else {
     t0 = $[0];
     t1 = $[1];
+    t2 = $[2];
   }
-  useEffect(t0, t1);
-  let t2;
+  t0(t1, t2);
   let t3;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
-    t2 = () => {
+  let t4;
+  if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
+    t3 = () => {
       setState(someGlobal.value);
     };
-    t3 = [someGlobal];
-    $[2] = t2;
+    t4 = [someGlobal];
     $[3] = t3;
+    $[4] = t4;
   } else {
-    t2 = $[2];
     t3 = $[3];
+    t4 = $[4];
   }
-  useEffect(t2, t3);
+  useEffect(t3, t4);
 
-  const t4 = String(state);
-  let t5;
-  if ($[4] !== t4) {
-    t5 = <div>{t4}</div>;
-    $[4] = t4;
+  const t5 = String(state);
+  let t6;
+  if ($[5] !== t5) {
+    t6 = <div>{t5}</div>;
     $[5] = t5;
+    $[6] = t6;
   } else {
-    t5 = $[5];
+    t6 = $[6];
   }
-  return t5;
+  return t6;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-repro-missed-memoization-from-inferred-mutation-in-logger.expect.md

@@ -8,16 +8,14 @@ import LogEvent from 'LogEvent';
 import {useCallback, useMemo} from 'react';
 
 component Component(id) {
-  const {data} = useFragment();
-  const items = data.items.edges;
+  const items = useFragment();
 
-  const [prevId, setPrevId] = useState(id);
   const [index, setIndex] = useState(0);
 
   const logData = useMemo(() => {
     const item = items[index];
     return {
-      key: item.key ?? '',
+      key: item.key,
     };
   }, [index, items]);
 
@@ -35,7 +33,6 @@ component Component(id) {
   );
 
   if (prevId !== id) {
-    setPrevId(id);
     setCurrentIndex(0);
   }
 
@@ -55,29 +52,27 @@ component Component(id) {
 ## Error
 
 ```
-  19 |
-  20 |   const setCurrentIndex = useCallback(
-> 21 |     (index: number) => {
-     |     ^^^^^^^^^^^^^^^^^^^^
-> 22 |       const object = {
-     | ^^^^^^^^^^^^^^^^^^^^^^
-> 23 |         tracking: logData.key,
-     | ^^^^^^^^^^^^^^^^^^^^^^
-> 24 |       };
-     | ^^^^^^^^^^^^^^^^^^^^^^
-> 25 |       // We infer that this may mutate `object`, which in turn aliases
-     | ^^^^^^^^^^^^^^^^^^^^^^
-> 26 |       // data from `logData`, such that `logData` may be mutated.
-     | ^^^^^^^^^^^^^^^^^^^^^^
-> 27 |       LogEvent.log(() => object);
-     | ^^^^^^^^^^^^^^^^^^^^^^
-> 28 |       setIndex(index);
-     | ^^^^^^^^^^^^^^^^^^^^^^
-> 29 |     },
-     | ^^^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (21:29)
-  30 |     [index, logData, items]
-  31 |   );
-  32 |
+   9 |   const [index, setIndex] = useState(0);
+  10 |
+> 11 |   const logData = useMemo(() => {
+     |                   ^^^^^^^^^^^^^^^
+> 12 |     const item = items[index];
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 13 |     return {
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 14 |       key: item.key,
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 15 |     };
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 16 |   }, [index, items]);
+     | ^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (11:16)
+
+CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This dependency may be mutated later, which could cause the value to change unexpectedly (28:28)
+
+CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (19:27)
+  17 |
+  18 |   const setCurrentIndex = useCallback(
+  19 |     (index: number) => {
 ```
           
       

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-repro-missed-memoization-from-inferred-mutation-in-logger.js

@@ -4,16 +4,14 @@ import LogEvent from 'LogEvent';
 import {useCallback, useMemo} from 'react';
 
 component Component(id) {
-  const {data} = useFragment();
-  const items = data.items.edges;
+  const items = useFragment();
 
-  const [prevId, setPrevId] = useState(id);
   const [index, setIndex] = useState(0);
 
   const logData = useMemo(() => {
     const item = items[index];
     return {
-      key: item.key ?? '',
+      key: item.key,
     };
   }, [index, items]);
 
@@ -31,7 +29,6 @@ component Component(id) {
   );
 
   if (prevId !== id) {
-    setPrevId(id);
     setCurrentIndex(0);
   }