From fde49869b6b1c099ab2251f7a4af8d2c4b6edbe5 Mon Sep 17 00:00:00 2001 From: Mike Vitousek Date: Wed, 29 May 2024 12:29:55 -0700 Subject: [PATCH] [compiler] Option for preserving calls to useMemo/useCallback Summary: This adds a compiler option to not drop existing manual memoization and leaving useMemo/useCallback in the generated source. Why do we need this, given that we also have options to validate or ensure that existing memoization is preserved? It's because later diffs on this stack are designed to alter the behavior of the memoization that the compiler emits, in order to detect rules of react violations and debug issues. We don't want to change the behavior of user-level memoization, however, since doing so would be altering the semantics of the user's program in an unacceptable way. [ghstack-poisoned] --- .../src/Entrypoint/Pipeline.ts | 6 +- .../src/HIR/Environment.ts | 7 ++ .../useMemo-simple-preserved.expect.md | 66 +++++++++++++++++++ .../compiler/useMemo-simple-preserved.js | 13 ++++ compiler/packages/snap/src/compiler.ts | 5 ++ 5 files changed, 95 insertions(+), 2 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js diff --git a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts index 5f50b7a0f1a..258cb9199e6 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/Entrypoint/Pipeline.ts @@ -148,8 +148,10 @@ function* runWithEnvironment( validateContextVariableLValues(hir); validateUseMemo(hir); - dropManualMemoization(hir); - yield log({ kind: "hir", name: "DropManualMemoization", value: hir }); + if (!env.config.enablePreserveExistingManualUseMemo) { + dropManualMemoization(hir); + yield log({ kind: "hir", name: "DropManualMemoization", value: hir }); + } inlineImmediatelyInvokedFunctionExpressions(hir); yield log({ diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts index f950068f15e..7375c35c765 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -165,6 +165,13 @@ const EnvironmentConfigSchema = z.object({ */ validatePreserveExistingMemoizationGuarantees: z.boolean().default(true), + /** + * When this is true, rather than pruning existing manual memoization but ensuring or validating + * that the memoized values remain memoized, the compiler will simply not prune existing calls to + * useMemo/useCallback. + */ + enablePreserveExistingManualUseMemo: z.boolean().default(false), + // 🌲 enableForest: z.boolean().default(false), diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md new file mode 100644 index 00000000000..6c813c27a68 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.expect.md @@ -0,0 +1,66 @@ + +## Input + +```javascript +// @enablePreserveExistingManualUseMemo +import { useMemo } from "react"; + +function Component({ a }) { + let x = useMemo(() => [a], []); + return
{x}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ a: 42 }], + isComponent: true, +}; + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @enablePreserveExistingManualUseMemo +import { useMemo } from "react"; + +function Component(t0) { + const $ = _c(5); + const { a } = t0; + let t1; + if ($[0] !== a) { + t1 = () => [a]; + $[0] = a; + $[1] = t1; + } else { + t1 = $[1]; + } + let t2; + if ($[2] === Symbol.for("react.memo_cache_sentinel")) { + t2 = []; + $[2] = t2; + } else { + t2 = $[2]; + } + const x = useMemo(t1, t2); + let t3; + if ($[3] !== x) { + t3 =
{x}
; + $[3] = x; + $[4] = t3; + } else { + t3 = $[4]; + } + return t3; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ a: 42 }], + isComponent: true, +}; + +``` + +### Eval output +(kind: ok)
42
\ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js new file mode 100644 index 00000000000..a5731f2f09d --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/useMemo-simple-preserved.js @@ -0,0 +1,13 @@ +// @enablePreserveExistingManualUseMemo +import { useMemo } from "react"; + +function Component({ a }) { + let x = useMemo(() => [a], []); + return
{x}
; +} + +export const FIXTURE_ENTRYPOINT = { + fn: Component, + params: [{ a: 42 }], + isComponent: true, +}; diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts index ab2cf5cef88..c949fa7fbe1 100644 --- a/compiler/packages/snap/src/compiler.ts +++ b/compiler/packages/snap/src/compiler.ts @@ -43,6 +43,7 @@ function makePluginOptions( let hookPattern: string | null = null; // TODO(@mofeiZ) rewrite snap fixtures to @validatePreserveExistingMemo:false let validatePreserveExistingMemoizationGuarantees = false; + let enablePreserveExistingManualUseMemo = false; if (firstLine.indexOf("@compilationMode(annotation)") !== -1) { assert( @@ -120,6 +121,9 @@ function makePluginOptions( validatePreserveExistingMemoizationGuarantees = true; } + if (firstLine.includes("@enablePreserveExistingManualUseMemo")) { + enablePreserveExistingManualUseMemo = true; + } const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine); if ( hookPatternMatch && @@ -173,6 +177,7 @@ function makePluginOptions( enableSharedRuntime__testonly: true, hookPattern, validatePreserveExistingMemoizationGuarantees, + enablePreserveExistingManualUseMemo, }, compilationMode, logger: null,