From c77acb3ac664a85d303f898220eb35dc8d2c0bca Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Fri, 15 Dec 2023 17:12:06 -0800
Subject: [PATCH] Separate mode to validate preserving manual memoization
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Adds a new mode which validates that existing manual memoization is preserved
_without_ using information from the manual memoization to affect compilation.
This gives us a way to try out the more aggressive version of Forget — ignoring
manual memoization — first and see how much code bails out and what patterns
cause this.

We can then proceed to enable the mode to actually _preserve_ existing memo
guarantees only where necessary.
---
 .../src/Entrypoint/Pipeline.ts                    |  5 ++++-
 .../src/HIR/Environment.ts                        | 15 +++++++++++++++
 .../src/Inference/DropManualMemoization.ts        | 10 ++++++----
 .../src/Inference/InferReferenceEffects.ts        |  6 +++++-
 4 files changed, 30 insertions(+), 6 deletions(-)

diff --git a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
index eded1c500d..b6f64c7279 100644
--- a/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts
@@ -349,7 +349,10 @@ function* runWithEnvironment(
     validateMemoizedEffectDependencies(reactiveFunction);
   }
 
-  if (env.config.enablePreserveExistingMemoizationGuarantees) {
+  if (
+    env.config.enablePreserveExistingMemoizationGuarantees ||
+    env.config.validatePreserveExistingMemoizationGuarantees
+  ) {
     validatePreservedManualMemoization(reactiveFunction);
   }
 
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 134f0cb3bc..bd35665e26 100644
--- a/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/HIR/Environment.ts
@@ -126,6 +126,21 @@ const EnvironmentConfigSchema = z.object({
    */
   enablePreserveExistingMemoizationGuarantees: z.boolean().default(false),
 
+  /**
+   * Validates that all useMemo/useCallback values are also memoized by Forget. This mode can be
+   * used with or without @enablePreserveExistingMemoizationGuarantees.
+   *
+   * With enablePreserveExistingMemoizationGuarantees, this validation enables automatically and
+   * verifies that Forget was able to preserve manual memoization semantics under that mode's
+   * additional assumptions about the input.
+   *
+   * With enablePreserveExistingMemoizationGuarantees off, this validation ignores manual memoization
+   * when determining program behavior, and only uses information from useMemo/useCallback to check
+   * that the memoization was preserved. This can be useful for determining where referential equalities
+   * may change under Forget.
+   */
+  validatePreserveExistingMemoizationGuarantees: z.boolean().default(false),
+
   // 🌲
   enableForest: z.boolean().default(false),
   // <🌲>
diff --git a/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts b/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts
index ccfcc39b51..4697f48842 100644
--- a/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts
@@ -16,7 +16,7 @@ import {
   SpreadPattern,
   makeInstructionId,
 } from "../HIR";
-import { createTemporaryPlace } from "../HIR/HIRBuilder";
+import { createTemporaryPlace, markInstructionIds } from "../HIR/HIRBuilder";
 import { HookKind } from "../HIR/ObjectShape";
 import { eachInstructionValueOperand } from "../HIR/visitors";
 
@@ -114,7 +114,8 @@ export function dropManualMemoization(func: HIRFunction): void {
                 };
 
                 if (
-                  func.env.config.enablePreserveExistingMemoizationGuarantees
+                  func.env.config.enablePreserveExistingMemoizationGuarantees ||
+                  func.env.config.validatePreserveExistingMemoizationGuarantees
                 ) {
                   /**
                    * When this flag is enabled we also compile in a 'Memoize' instruction
@@ -215,7 +216,8 @@ export function dropManualMemoization(func: HIRFunction): void {
                   loc: instr.value.loc,
                 };
                 if (
-                  func.env.config.enablePreserveExistingMemoizationGuarantees
+                  func.env.config.enablePreserveExistingMemoizationGuarantees ||
+                  func.env.config.validatePreserveExistingMemoizationGuarantees
                 ) {
                   nextInstructions =
                     nextInstructions ?? block.instructions.slice(0, i);
@@ -295,6 +297,6 @@ export function dropManualMemoization(func: HIRFunction): void {
     }
   }
   if (hasChanges) {
-    // markInstructionIds(func.body);
+    markInstructionIds(func.body);
   }
 }
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 a7227a5497..ca8d4f1625 100644
--- a/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
+++ b/compiler/packages/babel-plugin-react-forget/src/Inference/InferReferenceEffects.ts
@@ -1163,7 +1163,11 @@ function inferBlock(
         continue;
       }
       case "Memoize": {
-        state.reference(instrValue.value, Effect.Freeze, ValueReason.Other);
+        if (env.config.enablePreserveExistingMemoizationGuarantees) {
+          state.reference(instrValue.value, Effect.Freeze, ValueReason.Other);
+        } else {
+          state.reference(instrValue.value, Effect.Read, ValueReason.Other);
+        }
         const lvalue = instr.lvalue;
         lvalue.effect = Effect.ConditionallyMutate;
         state.initialize(instrValue, {