[compiler][bugfix] Returned functions are not always frozen

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

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}],
+};