From c50ee4d80cd20e022349bf09b69508f3d47b3dc1 Mon Sep 17 00:00:00 2001 From: Mofei Zhang Date: Mon, 28 Apr 2025 17:39:55 -0400 Subject: [PATCH] [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 --- .../src/Inference/InferReferenceEffects.ts | 30 +++++- ...eturned-inner-fn-mutates-context.expect.md | 92 ++++++++++++++++ ...repro-returned-inner-fn-mutates-context.js | 37 +++++++ ...urned-inner-fn-reassigns-context.expect.md | 101 ++++++++++++++++++ ...pro-returned-inner-fn-reassigns-context.js | 37 +++++++ 5 files changed, 293 insertions(+), 4 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-mutates-context.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-mutates-context.js create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.js 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 47ddc85a0b..724b67e5fa 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts +++ b/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; reason: Set}; // 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; @@ -267,16 +271,25 @@ class InferenceState { constructor( env: Environment, + isFunctionExpression: boolean, values: Map, variables: Map>, ) { 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; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-mutates-context.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-mutates-context.expect.md new file mode 100644 index 0000000000..0da18e5f65 --- /dev/null +++ b/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 ; +} + +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 = ; + $[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)
{"cb":{"value":2,"a":1},"shouldInvokeFns":true}
+
{"cb":{"value":3,"a":1},"shouldInvokeFns":true}
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-mutates-context.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-mutates-context.js new file mode 100644 index 0000000000..e5bd4a9b8c --- /dev/null +++ b/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 ; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Foo, + params: [{a: 1, b: 2}], + sequentialRenders: [ + {a: 1, b: 2}, + {a: 1, b: 3}, + ], +}; diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.expect.md new file mode 100644 index 0000000000..c09eaf060a --- /dev/null +++ b/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 ; +} + +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 = ; + $[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)
{"cb":{"kind":"Function"},"value":[1],"shouldInvokeFns":true}
+
{"cb":{"kind":"Function"},"value":[2],"shouldInvokeFns":true}
+logs: ['a','a'] \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/repro-returned-inner-fn-reassigns-context.js new file mode 100644 index 0000000000..00bb694b34 --- /dev/null +++ b/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 ; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Foo, + params: [{b: 1}], + sequentialRenders: [{b: 1}, {b: 2}], +};