From 32984349b01b15e1944025e6f16260599389de27 Mon Sep 17 00:00:00 2001
From: Mofei Zhang <feifei0@meta.com>
Date: Thu, 16 Jan 2025 15:20:01 -0500
Subject: [PATCH 1/3] [compiler][ez] rewrite invariant in InferReferenceEffects

Small patch to pass aliased context values into `Object|ArrayExpression`s
---
 .../src/Inference/InferFunctionEffects.ts     |  7 +-
 .../src/Inference/InferReferenceEffects.ts    | 75 +++++++++++--------
 2 files changed, 49 insertions(+), 33 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferFunctionEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferFunctionEffects.ts
index 0ae54839b6fa3..fe209924e4e08 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferFunctionEffects.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferFunctionEffects.ts
@@ -41,11 +41,16 @@ function inferOperandEffect(state: State, place: Place): null | FunctionEffect {
       if (isRefOrRefValue(place.identifier)) {
         break;
       } else if (value.kind === ValueKind.Context) {
+        CompilerError.invariant(value.context.size > 0, {
+          reason:
+            "[InferFunctionEffects] Expected Context-kind value's capture list to be non-empty.",
+          loc: place.loc,
+        });
         return {
           kind: 'ContextMutation',
           loc: place.loc,
           effect: place.effect,
-          places: value.context.size === 0 ? new Set([place]) : value.context,
+          places: value.context,
         };
       } else if (
         value.kind !== ValueKind.Mutable &&
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
index 8cf30a9666e25..70395e58d930e 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts
@@ -857,17 +857,19 @@ function inferBlock(
         break;
       }
       case 'ArrayExpression': {
-        const valueKind: AbstractValue = hasContextRefOperand(state, instrValue)
-          ? {
-              kind: ValueKind.Context,
-              reason: new Set([ValueReason.Other]),
-              context: new Set(),
-            }
-          : {
-              kind: ValueKind.Mutable,
-              reason: new Set([ValueReason.Other]),
-              context: new Set(),
-            };
+        const contextRefOperands = getContextRefOperand(state, instrValue);
+        const valueKind: AbstractValue =
+          contextRefOperands.length > 0
+            ? {
+                kind: ValueKind.Context,
+                reason: new Set([ValueReason.Other]),
+                context: new Set(contextRefOperands),
+              }
+            : {
+                kind: ValueKind.Mutable,
+                reason: new Set([ValueReason.Other]),
+                context: new Set(),
+              };
         continuation = {
           kind: 'initialize',
           valueKind,
@@ -918,17 +920,19 @@ function inferBlock(
         break;
       }
       case 'ObjectExpression': {
-        const valueKind: AbstractValue = hasContextRefOperand(state, instrValue)
-          ? {
-              kind: ValueKind.Context,
-              reason: new Set([ValueReason.Other]),
-              context: new Set(),
-            }
-          : {
-              kind: ValueKind.Mutable,
-              reason: new Set([ValueReason.Other]),
-              context: new Set(),
-            };
+        const contextRefOperands = getContextRefOperand(state, instrValue);
+        const valueKind: AbstractValue =
+          contextRefOperands.length > 0
+            ? {
+                kind: ValueKind.Context,
+                reason: new Set([ValueReason.Other]),
+                context: new Set(contextRefOperands),
+              }
+            : {
+                kind: ValueKind.Mutable,
+                reason: new Set([ValueReason.Other]),
+                context: new Set(),
+              };
 
         for (const property of instrValue.properties) {
           switch (property.kind) {
@@ -1593,15 +1597,21 @@ function inferBlock(
       }
       case 'LoadLocal': {
         const lvalue = instr.lvalue;
-        const effect =
-          state.isDefined(lvalue) &&
-          state.kind(lvalue).kind === ValueKind.Context
-            ? Effect.ConditionallyMutate
-            : Effect.Capture;
+        CompilerError.invariant(
+          !(
+            state.isDefined(lvalue) &&
+            state.kind(lvalue).kind === ValueKind.Context
+          ),
+          {
+            reason:
+              '[InferReferenceEffects] Unexpected LoadLocal with context kind',
+            loc: lvalue.loc,
+          },
+        );
         state.referenceAndRecordEffects(
           freezeActions,
           instrValue.place,
-          effect,
+          Effect.Capture,
           ValueReason.Other,
         );
         lvalue.effect = Effect.ConditionallyMutate;
@@ -1932,19 +1942,20 @@ function inferBlock(
   );
 }
 
-function hasContextRefOperand(
+function getContextRefOperand(
   state: InferenceState,
   instrValue: InstructionValue,
-): boolean {
+): Array<Place> {
+  const result = [];
   for (const place of eachInstructionValueOperand(instrValue)) {
     if (
       state.isDefined(place) &&
       state.kind(place).kind === ValueKind.Context
     ) {
-      return true;
+      result.push(place);
     }
   }
-  return false;
+  return result;
 }
 
 export function getFunctionCallSignature(

From 91b283418caed3a7f637e0bc4459be7ebfc9be1f Mon Sep 17 00:00:00 2001
From: Mofei Zhang <feifei0@meta.com>
Date: Thu, 16 Jan 2025 15:59:37 -0500
Subject: [PATCH 2/3] [compiler] Repro for invalid Array.map type

See test fixture
---
 ...-invalid-mixedreadonly-map-shape.expect.md | 163 ++++++++++++++++++
 .../bug-invalid-mixedreadonly-map-shape.js    |  56 ++++++
 .../packages/snap/src/SproutTodoFilter.ts     |   1 +
 3 files changed, 220 insertions(+)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.js

diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.expect.md
new file mode 100644
index 0000000000000..1418062c33458
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.expect.md
@@ -0,0 +1,163 @@
+
+## Input
+
+```javascript
+import {
+  arrayPush,
+  identity,
+  makeArray,
+  Stringify,
+  useFragment,
+} from 'shared-runtime';
+
+/**
+ * Bug repro showing why it's invalid for function references to be annotated
+ * with a `Read` effect when that reference might lead to the function being
+ * invoked.
+ *
+ * Note that currently, `Array.map` is annotated to have `Read` effects on its
+ * operands. This is incorrect as function effects must be replayed when `map`
+ * is called
+ * - Read:                non-aliasing data dependency
+ * - Capture:             maybe-aliasing data dependency
+ * - ConditionallyMutate: maybe-aliasing data dependency; maybe-write / invoke
+ *     but only if the value is mutable
+ *
+ * Invalid evaluator result: Found differences in evaluator results Non-forget
+ * (expected): (kind: ok)
+ *   <div>{"x":[2,2,2],"count":3}</div><div>{"item":1}</div>
+ *   <div>{"x":[2,2,2],"count":4}</div><div>{"item":1}</div>
+ * Forget:
+ *   (kind: ok)
+ *   <div>{"x":[2,2,2],"count":3}</div><div>{"item":1}</div>
+ *   <div>{"x":[2,2,2,2,2,2],"count":4}</div><div>{"item":1}</div>
+ */
+
+function Component({extraJsx}) {
+  const x = makeArray();
+  const items = useFragment();
+  const jsx = items.a.map((item, i) => {
+    arrayPush(x, 2);
+    return <Stringify item={item} key={i} />;
+  });
+  const offset = jsx.length;
+  for (let i = 0; i < extraJsx; i++) {
+    jsx.push(<Stringify item={0} key={i + offset} />);
+  }
+  const count = jsx.length;
+  identity(count);
+  return (
+    <>
+      <Stringify x={x} count={count} />
+      {jsx[0]}
+    </>
+  );
+}
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{extraJsx: 0}],
+  sequentialRenders: [{extraJsx: 0}, {extraJsx: 1}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import {
+  arrayPush,
+  identity,
+  makeArray,
+  Stringify,
+  useFragment,
+} from "shared-runtime";
+
+/**
+ * Bug repro showing why it's invalid for function references to be annotated
+ * with a `Read` effect when that reference might lead to the function being
+ * invoked.
+ *
+ * Note that currently, `Array.map` is annotated to have `Read` effects on its
+ * operands. This is incorrect as function effects must be replayed when `map`
+ * is called
+ * - Read:                non-aliasing data dependency
+ * - Capture:             maybe-aliasing data dependency
+ * - ConditionallyMutate: maybe-aliasing data dependency; maybe-write / invoke
+ *     but only if the value is mutable
+ *
+ * Invalid evaluator result: Found differences in evaluator results Non-forget
+ * (expected): (kind: ok)
+ *   <div>{"x":[2,2,2],"count":3}</div><div>{"item":1}</div>
+ *   <div>{"x":[2,2,2],"count":4}</div><div>{"item":1}</div>
+ * Forget:
+ *   (kind: ok)
+ *   <div>{"x":[2,2,2],"count":3}</div><div>{"item":1}</div>
+ *   <div>{"x":[2,2,2,2,2,2],"count":4}</div><div>{"item":1}</div>
+ */
+
+function Component(t0) {
+  const $ = _c(9);
+  const { extraJsx } = t0;
+  let t1;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = makeArray();
+    $[0] = t1;
+  } else {
+    t1 = $[0];
+  }
+  const x = t1;
+  const items = useFragment();
+  let jsx;
+  if ($[1] !== extraJsx || $[2] !== items.a) {
+    jsx = items.a.map((item, i) => {
+      arrayPush(x, 2);
+      return <Stringify item={item} key={i} />;
+    });
+    const offset = jsx.length;
+    for (let i_0 = 0; i_0 < extraJsx; i_0++) {
+      jsx.push(<Stringify item={0} key={i_0 + offset} />);
+    }
+    $[1] = extraJsx;
+    $[2] = items.a;
+    $[3] = jsx;
+  } else {
+    jsx = $[3];
+  }
+
+  const count = jsx.length;
+  identity(count);
+  let t2;
+  if ($[4] !== count) {
+    t2 = <Stringify x={x} count={count} />;
+    $[4] = count;
+    $[5] = t2;
+  } else {
+    t2 = $[5];
+  }
+  const t3 = jsx[0];
+  let t4;
+  if ($[6] !== t2 || $[7] !== t3) {
+    t4 = (
+      <>
+        {t2}
+        {t3}
+      </>
+    );
+    $[6] = t2;
+    $[7] = t3;
+    $[8] = t4;
+  } else {
+    t4 = $[8];
+  }
+  return t4;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ extraJsx: 0 }],
+  sequentialRenders: [{ extraJsx: 0 }, { extraJsx: 1 }],
+};
+
+```
+      
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.js
new file mode 100644
index 0000000000000..d038cf6cd38a1
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.js
@@ -0,0 +1,56 @@
+import {
+  arrayPush,
+  identity,
+  makeArray,
+  Stringify,
+  useFragment,
+} from 'shared-runtime';
+
+/**
+ * Bug repro showing why it's invalid for function references to be annotated
+ * with a `Read` effect when that reference might lead to the function being
+ * invoked.
+ *
+ * Note that currently, `Array.map` is annotated to have `Read` effects on its
+ * operands. This is incorrect as function effects must be replayed when `map`
+ * is called
+ * - Read:                non-aliasing data dependency
+ * - Capture:             maybe-aliasing data dependency
+ * - ConditionallyMutate: maybe-aliasing data dependency; maybe-write / invoke
+ *     but only if the value is mutable
+ *
+ * Invalid evaluator result: Found differences in evaluator results Non-forget
+ * (expected): (kind: ok)
+ *   <div>{"x":[2,2,2],"count":3}</div><div>{"item":1}</div>
+ *   <div>{"x":[2,2,2],"count":4}</div><div>{"item":1}</div>
+ * Forget:
+ *   (kind: ok)
+ *   <div>{"x":[2,2,2],"count":3}</div><div>{"item":1}</div>
+ *   <div>{"x":[2,2,2,2,2,2],"count":4}</div><div>{"item":1}</div>
+ */
+
+function Component({extraJsx}) {
+  const x = makeArray();
+  const items = useFragment();
+  const jsx = items.a.map((item, i) => {
+    arrayPush(x, 2);
+    return <Stringify item={item} key={i} />;
+  });
+  const offset = jsx.length;
+  for (let i = 0; i < extraJsx; i++) {
+    jsx.push(<Stringify item={0} key={i + offset} />);
+  }
+  const count = jsx.length;
+  identity(count);
+  return (
+    <>
+      <Stringify x={x} count={count} />
+      {jsx[0]}
+    </>
+  );
+}
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{extraJsx: 0}],
+  sequentialRenders: [{extraJsx: 0}, {extraJsx: 1}],
+};
diff --git a/compiler/packages/snap/src/SproutTodoFilter.ts b/compiler/packages/snap/src/SproutTodoFilter.ts
index 361070739630c..f363fd922e51a 100644
--- a/compiler/packages/snap/src/SproutTodoFilter.ts
+++ b/compiler/packages/snap/src/SproutTodoFilter.ts
@@ -486,6 +486,7 @@ const skipFilter = new Set([
   'bug-aliased-capture-mutate',
   'bug-functiondecl-hoisting',
   'bug-try-catch-maybe-null-dependency',
+  'bug-invalid-mixedreadonly-map-shape',
   'bug-type-inference-control-flow',
   'reduce-reactive-deps/bug-infer-function-cond-access-not-hoisted',
   'bug-invalid-phi-as-dependency',

From a86aaf0f85db6a333e2dc978c5ebf222309179e0 Mon Sep 17 00:00:00 2001
From: Mofei Zhang <feifei0@meta.com>
Date: Thu, 16 Jan 2025 16:01:32 -0500
Subject: [PATCH 3/3] [compiler] Fix invalid Array.map type

See test fixture
---
 .../src/HIR/ObjectShape.ts                    | 10 ++-
 ...d => mixedreadonly-mutating-map.expect.md} | 75 +++++++++----------
 ...shape.js => mixedreadonly-mutating-map.js} |  3 +
 .../packages/snap/src/SproutTodoFilter.ts     |  1 -
 4 files changed, 46 insertions(+), 43 deletions(-)
 rename compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/{bug-invalid-mixedreadonly-map-shape.expect.md => mixedreadonly-mutating-map.expect.md} (78%)
 rename compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/{bug-invalid-mixedreadonly-map-shape.js => mixedreadonly-mutating-map.js} (92%)

diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts
index 4482d17890196..ba2f4e8c4cc79 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/ObjectShape.ts
@@ -545,8 +545,16 @@ addObject(BUILTIN_SHAPES, BuiltInMixedReadonlyId, [
   [
     'map',
     addFunction(BUILTIN_SHAPES, [], {
+      /**
+       * Note `map`'s arguments are annotated as Effect.ConditionallyMutate as
+       * calling `<array>.map(fn)` might invoke `fn`, which means replaying its
+       * effects.
+       *
+       * (Note that Effect.Read / Effect.Capture on a function type means
+       * potential data dependency or aliasing respectively.)
+       */
       positionalParams: [],
-      restParam: Effect.Read,
+      restParam: Effect.ConditionallyMutate,
       returnType: {kind: 'Object', shapeId: BuiltInArrayId},
       calleeEffect: Effect.ConditionallyMutate,
       returnValueKind: ValueKind.Mutable,
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mixedreadonly-mutating-map.expect.md
similarity index 78%
rename from compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.expect.md
rename to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mixedreadonly-mutating-map.expect.md
index 1418062c33458..4867388a864d2 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mixedreadonly-mutating-map.expect.md
@@ -36,6 +36,9 @@ import {
 function Component({extraJsx}) {
   const x = makeArray();
   const items = useFragment();
+  // This closure has the following effects that must be replayed:
+  //   - MaybeFreeze / Capture of `items`
+  //   - ConditionalMutate of x
   const jsx = items.a.map((item, i) => {
     arrayPush(x, 2);
     return <Stringify item={item} key={i} />;
@@ -97,60 +100,47 @@ import {
  */
 
 function Component(t0) {
-  const $ = _c(9);
+  const $ = _c(6);
   const { extraJsx } = t0;
-  let t1;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t1 = makeArray();
-    $[0] = t1;
-  } else {
-    t1 = $[0];
-  }
-  const x = t1;
+  const x = makeArray();
   const items = useFragment();
-  let jsx;
-  if ($[1] !== extraJsx || $[2] !== items.a) {
-    jsx = items.a.map((item, i) => {
-      arrayPush(x, 2);
-      return <Stringify item={item} key={i} />;
-    });
-    const offset = jsx.length;
-    for (let i_0 = 0; i_0 < extraJsx; i_0++) {
-      jsx.push(<Stringify item={0} key={i_0 + offset} />);
-    }
-    $[1] = extraJsx;
-    $[2] = items.a;
-    $[3] = jsx;
-  } else {
-    jsx = $[3];
+
+  const jsx = items.a.map((item, i) => {
+    arrayPush(x, 2);
+    return <Stringify item={item} key={i} />;
+  });
+  const offset = jsx.length;
+  for (let i_0 = 0; i_0 < extraJsx; i_0++) {
+    jsx.push(<Stringify item={0} key={i_0 + offset} />);
   }
 
   const count = jsx.length;
   identity(count);
-  let t2;
-  if ($[4] !== count) {
-    t2 = <Stringify x={x} count={count} />;
-    $[4] = count;
-    $[5] = t2;
+  let t1;
+  if ($[0] !== count || $[1] !== x) {
+    t1 = <Stringify x={x} count={count} />;
+    $[0] = count;
+    $[1] = x;
+    $[2] = t1;
   } else {
-    t2 = $[5];
+    t1 = $[2];
   }
-  const t3 = jsx[0];
-  let t4;
-  if ($[6] !== t2 || $[7] !== t3) {
-    t4 = (
+  const t2 = jsx[0];
+  let t3;
+  if ($[3] !== t1 || $[4] !== t2) {
+    t3 = (
       <>
+        {t1}
         {t2}
-        {t3}
       </>
     );
-    $[6] = t2;
-    $[7] = t3;
-    $[8] = t4;
+    $[3] = t1;
+    $[4] = t2;
+    $[5] = t3;
   } else {
-    t4 = $[8];
+    t3 = $[5];
   }
-  return t4;
+  return t3;
 }
 
 export const FIXTURE_ENTRYPOINT = {
@@ -160,4 +150,7 @@ export const FIXTURE_ENTRYPOINT = {
 };
 
 ```
-      
\ No newline at end of file
+      
+### Eval output
+(kind: ok) <div>{"x":[2,2,2],"count":3}</div><div>{"item":1}</div>
+<div>{"x":[2,2,2],"count":4}</div><div>{"item":1}</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mixedreadonly-mutating-map.js
similarity index 92%
rename from compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.js
rename to compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mixedreadonly-mutating-map.js
index d038cf6cd38a1..858a4ab3dc693 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/bug-invalid-mixedreadonly-map-shape.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/mixedreadonly-mutating-map.js
@@ -32,6 +32,9 @@ import {
 function Component({extraJsx}) {
   const x = makeArray();
   const items = useFragment();
+  // This closure has the following effects that must be replayed:
+  //   - MaybeFreeze / Capture of `items`
+  //   - ConditionalMutate of x
   const jsx = items.a.map((item, i) => {
     arrayPush(x, 2);
     return <Stringify item={item} key={i} />;
diff --git a/compiler/packages/snap/src/SproutTodoFilter.ts b/compiler/packages/snap/src/SproutTodoFilter.ts
index f363fd922e51a..361070739630c 100644
--- a/compiler/packages/snap/src/SproutTodoFilter.ts
+++ b/compiler/packages/snap/src/SproutTodoFilter.ts
@@ -486,7 +486,6 @@ const skipFilter = new Set([
   'bug-aliased-capture-mutate',
   'bug-functiondecl-hoisting',
   'bug-try-catch-maybe-null-dependency',
-  'bug-invalid-mixedreadonly-map-shape',
   'bug-type-inference-control-flow',
   'reduce-reactive-deps/bug-infer-function-cond-access-not-hoisted',
   'bug-invalid-phi-as-dependency',