[compiler][bugfix] Returned functions are not always frozen (facebook#33047)

Fixes an edge case in React Compiler's effects inference model.

Returned values should only be typed as 'frozen' if they are (1) local
and (2) not a function expression which may capture and mutate this
function's outer context. See test fixtures for details
---
[//]: # (BEGIN SAPLING FOOTER)
Stack created with [Sapling](https://sapling-scm.com). Best reviewed
with [ReviewStack](https://reviewstack.dev/facebook/react/pull/33047).
* facebook#32765
* facebook#32747
* __->__ facebook#33047

Author: mofeiZ

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

@@ -111,7 +111,10 @@ export default function inferReferenceEffects(
    * Initial state contains function params
    * TODO: include module declarations here as well
    */
-  const initialState = InferenceState.empty(fn.env);
+  const initialState = InferenceState.empty(
+    fn.env,
+    options.isFunctionExpression,
+  );
   const value: InstructionValue = {
     kind: 'Primitive',
     loc: fn.loc,
@@ -255,6 +258,7 @@ type FreezeAction = {values: Set<InstructionValue>; reason: Set<ValueReason>};
 // Maintains a mapping of top-level variables to the kind of value they hold
 class InferenceState {
   env: Environment;
+  #isFunctionExpression: boolean;
 
   // The kind of each value, based on its allocation site
   #values: Map<InstructionValue, AbstractValue>;
@@ -267,16 +271,25 @@ class InferenceState {
 
   constructor(
     env: Environment,
+    isFunctionExpression: boolean,
     values: Map<InstructionValue, AbstractValue>,
     variables: Map<IdentifierId, Set<InstructionValue>>,
   ) {
     this.env = env;
+    this.#isFunctionExpression = isFunctionExpression;
     this.#values = values;
     this.#variables = variables;
   }
 
-  static empty(env: Environment): InferenceState {
-    return new InferenceState(env, new Map(), new Map());
+  static empty(
+    env: Environment,
+    isFunctionExpression: boolean,
+  ): InferenceState {
+    return new InferenceState(env, isFunctionExpression, new Map(), new Map());
+  }
+
+  get isFunctionExpression(): boolean {
+    return this.#isFunctionExpression;
   }
 
   // (Re)initializes a @param value with its default @param kind.
@@ -613,6 +626,7 @@ class InferenceState {
     } else {
       return new InferenceState(
         this.env,
+        this.#isFunctionExpression,
         nextValues ?? new Map(this.#values),
         nextVariables ?? new Map(this.#variables),
       );
@@ -627,6 +641,7 @@ class InferenceState {
   clone(): InferenceState {
     return new InferenceState(
       this.env,
+      this.#isFunctionExpression,
       new Map(this.#values),
       new Map(this.#variables),
     );
@@ -1781,8 +1796,15 @@ function inferBlock(
     if (block.terminal.kind === 'return' || block.terminal.kind === 'throw') {
       if (
         state.isDefined(operand) &&
-        state.kind(operand).kind === ValueKind.Context
+        ((operand.identifier.type.kind === 'Function' &&
+          state.isFunctionExpression) ||
+          state.kind(operand).kind === ValueKind.Context)
       ) {
+        /**
+         * Returned values should only be typed as 'frozen' if they are both (1)
+         * local and (2) not a function expression which may capture and mutate
+         * this function's outer context.
+         */
         effect = Effect.ConditionallyMutate;
       } else {
         effect = Effect.Freeze;

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-mutates-context.expect.md

@@ -0,0 +1,92 @@
+
+## Input
+
+```javascript
+import {Stringify} from 'shared-runtime';
+
+/**
+ * Example showing that returned inner function expressions should not be
+ * typed with `freeze` effects.
+ */
+function Foo({a, b}) {
+  'use memo';
+  const obj = {};
+  const updaterFactory = () => {
+    /**
+     * This returned function expression *is* a local value. But it might (1)
+     * capture and mutate its context environment and (2) be called during
+     * render.
+     * Typing it with `freeze` effects would be incorrect as it would mean
+     * inferring that calls to updaterFactory()() do not mutate its captured
+     * context.
+     */
+    return newValue => {
+      obj.value = newValue;
+      obj.a = a;
+    };
+  };
+
+  const updater = updaterFactory();
+  updater(b);
+  return <Stringify cb={obj} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{a: 1, b: 2}],
+  sequentialRenders: [
+    {a: 1, b: 2},
+    {a: 1, b: 3},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { Stringify } from "shared-runtime";
+
+/**
+ * Example showing that returned inner function expressions should not be
+ * typed with `freeze` effects.
+ */
+function Foo(t0) {
+  "use memo";
+  const $ = _c(3);
+  const { a, b } = t0;
+  let t1;
+  if ($[0] !== a || $[1] !== b) {
+    const obj = {};
+    const updaterFactory = () => (newValue) => {
+      obj.value = newValue;
+      obj.a = a;
+    };
+
+    const updater = updaterFactory();
+    updater(b);
+    t1 = <Stringify cb={obj} shouldInvokeFns={true} />;
+    $[0] = a;
+    $[1] = b;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ a: 1, b: 2 }],
+  sequentialRenders: [
+    { a: 1, b: 2 },
+    { a: 1, b: 3 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"cb":{"value":2,"a":1},"shouldInvokeFns":true}</div>
+<div>{"cb":{"value":3,"a":1},"shouldInvokeFns":true}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-mutates-context.js

@@ -0,0 +1,37 @@
+import {Stringify} from 'shared-runtime';
+
+/**
+ * Example showing that returned inner function expressions should not be
+ * typed with `freeze` effects.
+ */
+function Foo({a, b}) {
+  'use memo';
+  const obj = {};
+  const updaterFactory = () => {
+    /**
+     * This returned function expression *is* a local value. But it might (1)
+     * capture and mutate its context environment and (2) be called during
+     * render.
+     * Typing it with `freeze` effects would be incorrect as it would mean
+     * inferring that calls to updaterFactory()() do not mutate its captured
+     * context.
+     */
+    return newValue => {
+      obj.value = newValue;
+      obj.a = a;
+    };
+  };
+
+  const updater = updaterFactory();
+  updater(b);
+  return <Stringify cb={obj} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{a: 1, b: 2}],
+  sequentialRenders: [
+    {a: 1, b: 2},
+    {a: 1, b: 3},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md

@@ -0,0 +1,101 @@
+
+## Input
+
+```javascript
+import {makeArray, Stringify, useIdentity} from 'shared-runtime';
+
+/**
+ * Example showing that returned inner function expressions should not be
+ * typed with `freeze` effects.
+ * Also see repro-returned-inner-fn-mutates-context
+ */
+function Foo({b}) {
+  'use memo';
+
+  const fnFactory = () => {
+    /**
+     * This returned function expression *is* a local value. But it might (1)
+     * capture and mutate its context environment and (2) be called during
+     * render.
+     * Typing it with `freeze` effects would be incorrect as it would mean
+     * inferring that calls to updaterFactory()() do not mutate its captured
+     * context.
+     */
+    return () => {
+      myVar = () => console.log('a');
+    };
+  };
+  let myVar = () => console.log('b');
+  useIdentity();
+
+  const fn = fnFactory();
+  const arr = makeArray(b);
+  fn(arr);
+  return <Stringify cb={myVar} value={arr} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{b: 1}],
+  sequentialRenders: [{b: 1}, {b: 2}],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { makeArray, Stringify, useIdentity } from "shared-runtime";
+
+/**
+ * Example showing that returned inner function expressions should not be
+ * typed with `freeze` effects.
+ * Also see repro-returned-inner-fn-mutates-context
+ */
+function Foo(t0) {
+  "use memo";
+  const $ = _c(3);
+  const { b } = t0;
+
+  const fnFactory = () => () => {
+    myVar = _temp;
+  };
+
+  let myVar;
+  myVar = _temp2;
+  useIdentity();
+
+  const fn = fnFactory();
+  const arr = makeArray(b);
+  fn(arr);
+  let t1;
+  if ($[0] !== arr || $[1] !== myVar) {
+    t1 = <Stringify cb={myVar} value={arr} shouldInvokeFns={true} />;
+    $[0] = arr;
+    $[1] = myVar;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+function _temp2() {
+  return console.log("b");
+}
+function _temp() {
+  return console.log("a");
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{ b: 1 }],
+  sequentialRenders: [{ b: 1 }, { b: 2 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"cb":{"kind":"Function"},"value":[1],"shouldInvokeFns":true}</div>
+<div>{"cb":{"kind":"Function"},"value":[2],"shouldInvokeFns":true}</div>
+logs: ['a','a']

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.js

@@ -0,0 +1,37 @@
+import {makeArray, Stringify, useIdentity} from 'shared-runtime';
+
+/**
+ * Example showing that returned inner function expressions should not be
+ * typed with `freeze` effects.
+ * Also see repro-returned-inner-fn-mutates-context
+ */
+function Foo({b}) {
+  'use memo';
+
+  const fnFactory = () => {
+    /**
+     * This returned function expression *is* a local value. But it might (1)
+     * capture and mutate its context environment and (2) be called during
+     * render.
+     * Typing it with `freeze` effects would be incorrect as it would mean
+     * inferring that calls to updaterFactory()() do not mutate its captured
+     * context.
+     */
+    return () => {
+      myVar = () => console.log('a');
+    };
+  };
+  let myVar = () => console.log('b');
+  useIdentity();
+
+  const fn = fnFactory();
+  const arr = makeArray(b);
+  fn(arr);
+  return <Stringify cb={myVar} value={arr} shouldInvokeFns={true} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Foo,
+  params: [{b: 1}],
+  sequentialRenders: [{b: 1}, {b: 2}],
+};