Validation that useMemo/useCallback is preserved in the output

Extends `@enablePreserveExistingMemoization` to validate that all of the 
original values were actually memoized. This works nearly identically to how we 
validate effect deps are memoized. We look for Memoize instructions whose values 
need memoization but whose range extends past the memoize instruction, or where 
the value isn't memoized at all.

compiler/packages/babel-plugin-react-forget/src/Entrypoint/Pipeline.ts

@@ -74,6 +74,7 @@ import {
   validateMemoizedEffectDependencies,
   validateNoRefAccessInRender,
   validateNoSetStateInRender,
+  validatePreservedManualMemoization,
   validateUseMemo,
 } from "../Validation";
 
@@ -348,6 +349,10 @@ function* runWithEnvironment(
     validateMemoizedEffectDependencies(reactiveFunction);
   }
 
+  if (env.config.enablePreserveExistingMemoizationGuarantees) {
+    validatePreservedManualMemoization(reactiveFunction);
+  }
+
   if (env.config.enableForest) {
     yield* lowerToForest(reactiveFunction);
   }

compiler/packages/babel-plugin-react-forget/src/Validation/ValidatePreservedManualMemoization.ts

@@ -0,0 +1,95 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import { CompilerError, ErrorSeverity } from "..";
+import {
+  Identifier,
+  Instruction,
+  ReactiveFunction,
+  ReactiveInstruction,
+  ReactiveScopeBlock,
+  ScopeId,
+} from "../HIR";
+import { isMutable } from "../ReactiveScopes/InferReactiveScopeVariables";
+import {
+  ReactiveFunctionVisitor,
+  visitReactiveFunction,
+} from "../ReactiveScopes/visitors";
+
+/**
+ * Validates that all explicit manual memoization (useMemo/useCallback) was accurately
+ * preserved, and that no originally memoized values became unmemoized in the output.
+ *
+ * This can occur if a value's mutable range somehow extended to include a hook and
+ * was pruned.
+ */
+export function validatePreservedManualMemoization(fn: ReactiveFunction): void {
+  const errors = new CompilerError();
+  visitReactiveFunction(fn, new Visitor(), errors);
+  if (errors.hasErrors()) {
+    throw errors;
+  }
+}
+
+class Visitor extends ReactiveFunctionVisitor<CompilerError> {
+  scopes: Set<ScopeId> = new Set();
+
+  override visitScope(
+    scopeBlock: ReactiveScopeBlock,
+    state: CompilerError
+  ): void {
+    this.traverseScope(scopeBlock, state);
+
+    /*
+     * Record scopes that exist in the AST so we can later check to see if
+     * effect dependencies which should be memoized (have a scope assigned)
+     * actually are memoized (that scope exists).
+     * However, we only record scopes if *their* dependencies are also
+     * memoized, allowing a transitive memoization check.
+     */
+    let areDependenciesMemoized = true;
+    for (const dep of scopeBlock.scope.dependencies) {
+      if (isUnmemoized(dep.identifier, this.scopes)) {
+        areDependenciesMemoized = false;
+        break;
+      }
+    }
+    if (areDependenciesMemoized) {
+      this.scopes.add(scopeBlock.scope.id);
+      for (const id of scopeBlock.scope.merged) {
+        this.scopes.add(id);
+      }
+    }
+  }
+
+  override visitInstruction(
+    instruction: ReactiveInstruction,
+    state: CompilerError
+  ): void {
+    this.traverseInstruction(instruction, state);
+    if (instruction.value.kind === "Memoize") {
+      const value = instruction.value.value;
+      if (
+        isMutable(instruction as Instruction, value) ||
+        isUnmemoized(value.identifier, this.scopes)
+      ) {
+        state.push({
+          reason:
+            "This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized",
+          description: null,
+          severity: ErrorSeverity.InvalidReact,
+          loc: typeof instruction.loc !== "symbol" ? instruction.loc : null,
+          suggestions: null,
+        });
+      }
+    }
+  }
+}
+
+function isUnmemoized(operand: Identifier, scopes: Set<ScopeId>): boolean {
+  return operand.scope != null && !scopes.has(operand.scope.id);
+}

compiler/packages/babel-plugin-react-forget/src/Validation/index.ts

@@ -10,4 +10,5 @@ export { validateHooksUsage } from "./ValidateHooksUsage";
 export { validateMemoizedEffectDependencies } from "./ValidateMemoizedEffectDependencies";
 export { validateNoRefAccessInRender } from "./ValidateNoRefAccesInRender";
 export { validateNoSetStateInRender } from "./ValidateNoSetStateInRender";
+export { validatePreservedManualMemoization } from "./ValidatePreservedManualMemoization";
 export { validateUseMemo } from "./ValidateUseMemo";

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/error.todo-useCallback-set-ref-nested-property-preserve-memoization.expect.md

@@ -0,0 +1,38 @@
+
+## Input
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees
+import { useCallback, useRef } from "react";
+
+function Component(props) {
+  const ref = useRef({ inner: null });
+
+  const onChange = useCallback((event) => {
+    // The ref should still be mutable here even though function deps are frozen in
+    // @enablePreserveExistingMemoizationGuarantees mode
+    ref.current.inner = event.target.value;
+  });
+
+  ref.current.inner = null;
+
+  return <input onChange={onChange} />;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidReact: This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized (7:11)
+
+[ReactForget] InvalidReact: This value was manually memoized, but cannot be memoized under Forget because it may be mutated after it is memoized (7:11)
+```
+          
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useCallback-set-ref-nested-property-dont-preserve-memoization.expect.md

@@ -62,6 +62,4 @@ export const FIXTURE_ENTRYPOINT = {
 };
 
 ```
-      
-### Eval output
-(kind: ok) <input>
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useCallback-set-ref-nested-property-preserve-memoization.expect.md

@@ -1,74 +0,0 @@
-
-## Input
-
-```javascript
-// @enablePreserveExistingMemoizationGuarantees
-import { useCallback, useRef } from "react";
-
-function Component(props) {
-  const ref = useRef({ inner: null });
-
-  const onChange = useCallback((event) => {
-    // The ref should still be mutable here even though function deps are frozen in
-    // @enablePreserveExistingMemoizationGuarantees mode
-    ref.current.inner = event.target.value;
-  });
-
-  ref.current.inner = null;
-
-  return <input onChange={onChange} />;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
-  params: [{}],
-};
-
-```
-
-## Code
-
-```javascript
-// @enablePreserveExistingMemoizationGuarantees
-import {
-  useCallback,
-  useRef,
-  unstable_useMemoCache as useMemoCache,
-} from "react";
-
-function Component(props) {
-  const $ = useMemoCache(3);
-  let t0;
-  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = { inner: null };
-    $[0] = t0;
-  } else {
-    t0 = $[0];
-  }
-  const ref = useRef(t0);
-
-  const onChange = (event) => {
-    ref.current.inner = event.target.value;
-  };
-
-  ref.current.inner = null;
-  let t1;
-  if ($[1] !== onChange) {
-    t1 = <input onChange={onChange} />;
-    $[1] = onChange;
-    $[2] = t1;
-  } else {
-    t1 = $[2];
-  }
-  return t1;
-}
-
-export const FIXTURE_ENTRYPOINT = {
-  fn: Component,
-  params: [{}],
-};
-
-```
-      
-### Eval output
-(kind: ok) <input>