From f504eaa16ecbfc1000d9fc9fc007b2ec313f7dda Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Mon, 18 Dec 2023 15:33:02 -0800
Subject: [PATCH] Add back transitive freeze functions option

Adds back a mode to transitively freeze function expressions, independently from
the mode to preserve existing manual memoization. This lets us experiment with a
few variants:

* Preserve existing memoization

* Validate existing memoization with:

* `enableAssumeHooksFollowRulesOfReact` &&
`enableTransitivelyFreezeFunctionExpressions`

* `enableAssumeHooksFollowRulesOfReact` only

* neither of those flags

Note that `enableTransitivelyFreezeFunctionExpressions` alone probably doesn't
make sense, it's more aggressive than

`enableAssumeHooksFollowRulesOfReact` so we might as well try them together.
---
 .../babel-plugin-react-forget/src/HIR/Environment.ts      | 8 ++++++++
 .../src/Inference/InferReferenceEffects.ts                | 5 ++++-
 .../transitive-freeze-function-expressions.expect.md      | 4 ++--
 .../compiler/transitive-freeze-function-expressions.js    | 2 +-
 4 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
index bd35665e26..fda05030d5 100644
--- a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
@@ -210,6 +210,14 @@ const EnvironmentConfigSchema = z.object({
    */
   enableAssumeHooksFollowRulesOfReact: z.boolean().default(false),
 
+  /**
+   * When enabled, the compiler assumes that any values are not subsequently
+   * modified after they are captured by a function passed to React. For example,
+   * if a value `x` is referenced inside a function expression passed to `useEffect`,
+   * then this flag will assume that `x` is not subusequently modified.
+   */
+  enableTransitivelyFreezeFunctionExpressions: z.boolean().default(false),
+
   /*
    * When enabled, removes *all* memoization from the function: this includes
    * removing manually added useMemo/useCallback as well as not adding Forget's
diff --git a/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts b/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
index ca8d4f1625..33e5c21a77 100644
--- a/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
@@ -354,7 +354,10 @@ class InferenceState {
               reason: reasonSet,
             });
 
-            if (this.#env.config.enablePreserveExistingMemoizationGuarantees) {
+            if (
+              this.#env.config.enablePreserveExistingMemoizationGuarantees ||
+              this.#env.config.enableTransitivelyFreezeFunctionExpressions
+            ) {
               if (value.kind === "FunctionExpression") {
                 for (const operand of eachInstructionValueOperand(value)) {
                   this.reference(operand, Effect.Freeze, ValueReason.Other);
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/transitive-freeze-function-expressions.expect.md b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/transitive-freeze-function-expressions.expect.md
index 8478234590..009fe6d34d 100644
--- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/transitive-freeze-function-expressions.expect.md
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/transitive-freeze-function-expressions.expect.md
@@ -2,7 +2,7 @@
 ## Input
 
 ```javascript
-// @enablePreserveExistingMemoizationGuarantees
+// @enableTransitivelyFreezeFunctionExpressions
 function Component(props) {
   const { data, loadNext, isLoadingNext } =
     usePaginationFragment(props.key).items ?? [];
@@ -31,7 +31,7 @@ function Component(props) {
 ## Code
 
 ```javascript
-import { unstable_useMemoCache as useMemoCache } from "react"; // @enablePreserveExistingMemoizationGuarantees
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableTransitivelyFreezeFunctionExpressions
 function Component(props) {
   const $ = useMemoCache(10);
   const { data, loadNext, isLoadingNext } =
diff --git a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/transitive-freeze-function-expressions.js b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/transitive-freeze-function-expressions.js
index bd6897cd32..7a3dd5a5e5 100644
--- a/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/transitive-freeze-function-expressions.js
+++ b/compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/transitive-freeze-function-expressions.js
@@ -1,4 +1,4 @@
-// @enablePreserveExistingMemoizationGuarantees
+// @enableTransitivelyFreezeFunctionExpressions
 function Component(props) {
   const { data, loadNext, isLoadingNext } =
     usePaginationFragment(props.key).items ?? [];