From 8bf1ec97a40e04a95dcc9bb1a576e707d8c585db Mon Sep 17 00:00:00 2001
From: mofeiZ <34200447+mofeiZ@users.noreply.github.com>
Date: Mon, 15 May 2023 14:27:26 -0400
Subject: [PATCH] Add flag back for inlineUseMemo

Reverts #1502, but flips test flags (e.g. `inlineUseMemo` by default, unless a
test specifies `@inlineUseMemo false`. I figured this add less thrash for test
fixtures, but happy to just do a clean revert (or remove the pragma altogether
and always pass `inlineUseMemo: true`)
---
 compiler/forget/src/CompilerPipeline.ts                     | 6 ++++--
 compiler/forget/src/HIR/Environment.ts                      | 1 +
 compiler/forget/src/__tests__/compiler-test.ts              | 1 +
 .../forget/src/__tests__/disableMemoizeJsxElements-test.ts  | 1 +
 .../src/__tests__/test-utils/generateTestsFromFixtures.ts   | 6 +++++-
 5 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/compiler/forget/src/CompilerPipeline.ts b/compiler/forget/src/CompilerPipeline.ts
index 5f5eafb74b..983bc6bf19 100644
--- a/compiler/forget/src/CompilerPipeline.ts
+++ b/compiler/forget/src/CompilerPipeline.ts
@@ -66,8 +66,10 @@ export function* run(
   const hir = lower(func, env).unwrap();
   yield log({ kind: "hir", name: "HIR", value: hir });
 
-  inlineUseMemo(hir);
-  yield log({ kind: "hir", name: "RewriteUseMemo", value: hir });
+  if (config?.inlineUseMemo) {
+    inlineUseMemo(hir);
+    yield log({ kind: "hir", name: "RewriteUseMemo", value: hir });
+  }
 
   mergeConsecutiveBlocks(hir);
   yield log({ kind: "hir", name: "MergeConsecutiveBlocks", value: hir });
diff --git a/compiler/forget/src/HIR/Environment.ts b/compiler/forget/src/HIR/Environment.ts
index ec23034ec9..df5cf3c8ee 100644
--- a/compiler/forget/src/HIR/Environment.ts
+++ b/compiler/forget/src/HIR/Environment.ts
@@ -40,6 +40,7 @@ export type EnvironmentConfig = Partial<{
   customHooks: Map<string, Hook>;
   memoizeJsxElements: boolean;
   validateHooksUsage: boolean;
+  inlineUseMemo: boolean;
 }>;
 
 export class Environment {
diff --git a/compiler/forget/src/__tests__/compiler-test.ts b/compiler/forget/src/__tests__/compiler-test.ts
index ffe40a03ee..cda5f2d575 100644
--- a/compiler/forget/src/__tests__/compiler-test.ts
+++ b/compiler/forget/src/__tests__/compiler-test.ts
@@ -56,6 +56,7 @@ describe("React Forget", () => {
                 ],
               ]),
               validateHooksUsage: true,
+              inlineUseMemo: options.environment?.inlineUseMemo ?? false,
             },
             logger: null,
             gating: options.gating,
diff --git a/compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts b/compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts
index 2a9408a152..1148bc2aed 100644
--- a/compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts
+++ b/compiler/forget/src/__tests__/disableMemoizeJsxElements-test.ts
@@ -50,6 +50,7 @@ describe("React Forget (Disable memoization of JSX elements)", () => {
                   },
                 ],
               ]),
+              inlineUseMemo: options.environment?.inlineUseMemo ?? false,
             },
             logger: null,
             gating: options.gating,
diff --git a/compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts b/compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts
index 149f4cd65f..7590ec5fa7 100644
--- a/compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts
+++ b/compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts
@@ -94,6 +94,7 @@ export default function generateTestsFromFixtures(
         let debug = false;
         let enableOnlyOnUseForgetDirective = false;
         let gating: GatingOptions | null = null;
+        let inlineUseMemo = true;
         let panicOnBailout = true;
 
         if (inputFile != null) {
@@ -117,6 +118,9 @@ export default function generateTestsFromFixtures(
               importSpecifierName: "isForgetEnabled_Fixtures",
             };
           }
+          if (lines[0]!.indexOf("@inlineUseMemo false") !== -1) {
+            inlineUseMemo = false;
+          }
           if (lines[0]!.indexOf("@panicOnBailout false") !== -1) {
             panicOnBailout = false;
           }
@@ -126,7 +130,7 @@ export default function generateTestsFromFixtures(
           let receivedOutput;
           if (input !== null) {
             receivedOutput = transform(input, basename, {
-              environment: {},
+              environment: { inlineUseMemo },
               logger: null,
               debug,
               enableOnlyOnUseForgetDirective,