enablePreserveMemo treats memo deps as frozen

See discussion on #2448 for full context. In the new 
`@enablePreserveExistingMemoizationGuarantees` mode, the goal is to preserve the 
existing referential equality guarantees from the original code. #2448 lays the 
groundwork by explicitly marking the _output_ of each useMemo block as memoized, 
hinting to the compiler that the value cannot subsequently change. This ensures 
the mutable range doesn't extend _later_, possibly overlapping a hook call and 
causing memoization to gett pruned. 

This PR fixes the other direction. There are cases where free variables 
referenced in the useMemo block could have been inferred as mutated, which could 
then extend the _start_ of the range earlier past a hook: 

```javascript 

const foo = createObject(); 

useBar(); 

const baz = useMemo(() => { 

const baz = createObject(); 

maybeMutate(foo, baz); 

return baz; 

}, [foo]); 

``` 

Here the compiler would infer that both `baz` and `foo` are mutable at the 
`maybeMutate()` call, grouping them in the same scope. But that scope would span 
the `useBar()` call, and be pruned, meaning that `baz` went unmemoized. 

However, useMemo blocks shouldn't be mutating free variables. Only variables 
newly created within the useMemo block should be mutable. So this PR extends the 
feature to treat all free variables referenced in a useMemo block as frozen as 
of the block itself.

compiler/packages/babel-plugin-react-forget/src/HIR/PrintHIR.ts

@@ -494,6 +494,8 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
     }
     case "ObjectMethod":
     case "FunctionExpression": {
+      const kind =
+        instrValue.kind === "FunctionExpression" ? "Function" : "ObjectMethod";
       const name = getFunctionName(instrValue, "");
       const fn = printFunction(instrValue.loweredFunc.func)
         .split("\n")
@@ -505,7 +507,7 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       const context = instrValue.loweredFunc.func.context
         .map((dep) => printPlace(dep))
         .join(",");
-      value = `Function ${name} @deps[${deps}] @context[${context}]:\n${fn}`;
+      value = `${kind} ${name} @deps[${deps}] @context[${context}]:\n${fn}`;
       break;
     }
     case "TaggedTemplateExpression": {

compiler/packages/babel-plugin-react-forget/src/Inference/DropManualMemoization.ts

@@ -8,16 +8,17 @@
 import { CompilerError } from "..";
 import {
   Effect,
+  FunctionExpression,
   HIRFunction,
   IdentifierId,
   Instruction,
   Place,
   SpreadPattern,
   makeInstructionId,
-  markInstructionIds,
 } from "../HIR";
 import { createTemporaryPlace } from "../HIR/HIRBuilder";
 import { HookKind } from "../HIR/ObjectShape";
+import { eachInstructionValueOperand } from "../HIR/visitors";
 
 /*
  * Removes manual memoization using the `useMemo` and `useCallback` APIs. This pass is designed
@@ -28,6 +29,7 @@ import { HookKind } from "../HIR/ObjectShape";
  * eg `React.useMemo()`.
  */
 export function dropManualMemoization(func: HIRFunction): void {
+  const functions = new Map<IdentifierId, FunctionExpression>();
   const hooks = new Map<IdentifierId, HookKind>();
   const react = new Set<IdentifierId>();
   let hasChanges = false;
@@ -35,10 +37,11 @@ export function dropManualMemoization(func: HIRFunction): void {
     let nextInstructions: Array<Instruction> | null = null;
     for (let i = 0; i < block.instructions.length; i++) {
       const instr = block.instructions[i]!;
-      if (nextInstructions !== null) {
-        nextInstructions.push(instr);
-      }
       switch (instr.value.kind) {
+        case "FunctionExpression": {
+          functions.set(instr.lvalue.identifier.id, instr.value);
+          break;
+        }
         case "LoadGlobal": {
           if (
             instr.value.name === "useMemo" ||
@@ -109,6 +112,7 @@ export function dropManualMemoization(func: HIRFunction): void {
                   args: [],
                   loc: instr.value.loc,
                 };
+
                 if (
                   func.env.config.enablePreserveExistingMemoizationGuarantees
                 ) {
@@ -137,7 +141,29 @@ export function dropManualMemoization(func: HIRFunction): void {
                   const temp = createTemporaryPlace(func.env);
                   instr.lvalue = { ...temp };
                   nextInstructions =
-                    nextInstructions ?? block.instructions.slice(0, i + 1);
+                    nextInstructions ?? block.instructions.slice(0, i);
+
+                  const functionExpression = functions.get(fn.identifier.id);
+                  if (functionExpression !== undefined) {
+                    for (const operand of eachInstructionValueOperand(
+                      functionExpression
+                    )) {
+                      const operandLValue = createTemporaryPlace(func.env);
+                      nextInstructions.push({
+                        id: makeInstructionId(0),
+                        lvalue: operandLValue,
+                        value: {
+                          kind: "Memoize",
+                          value: { ...operand },
+                          loc: instr.loc,
+                        },
+                        loc: instr.loc,
+                      });
+                    }
+                  }
+
+                  nextInstructions.push(instr);
+
                   nextInstructions.push({
                     id: makeInstructionId(0),
                     lvalue,
@@ -148,7 +174,12 @@ export function dropManualMemoization(func: HIRFunction): void {
                     },
                     loc: instr.loc,
                   });
+                } else {
+                  if (nextInstructions !== null) {
+                    nextInstructions.push(instr);
+                  }
                 }
+                continue;
               }
             } else if (hookKind === "useCallback") {
               const [fn] = instr.value.args as Array<
@@ -229,6 +260,9 @@ export function dropManualMemoization(func: HIRFunction): void {
           break;
         }
       }
+      if (nextInstructions !== null) {
+        nextInstructions.push(instr);
+      }
     }
     if (nextInstructions !== null) {
       block.instructions = nextInstructions;
@@ -236,6 +270,6 @@ export function dropManualMemoization(func: HIRFunction): void {
     }
   }
   if (hasChanges) {
-    markInstructionIds(func.body);
+    // markInstructionIds(func.body);
   }
 }

compiler/packages/babel-plugin-react-forget/src/Optimization/DeadCodeElimination.ts

@@ -313,7 +313,14 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
     case "StoreContext": {
       return false;
     }
-    case "Memoize":
+    case "Memoize": {
+      /**
+       * This instruction is used by the @enablePreserveExistingMemoizationGuarantees feature
+       * to preserve information about memoization semantics in the original code. We can't
+       * DCE without losing the memoization guarantees.
+       */
+      return false;
+    }
     case "RegExpLiteral":
     case "LoadGlobal":
     case "ArrayExpression":

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.expect.md

@@ -0,0 +1,74 @@
+
+## Input
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees:false
+import { useMemo } from "react";
+import {
+  identity,
+  makeObject_Primitives,
+  mutate,
+  useHook,
+} from "shared-runtime";
+
+function Component(props) {
+  // With the feature disabled these variables are inferred as being mutated inside the useMemo block
+  const free = makeObject_Primitives();
+  const free2 = makeObject_Primitives();
+  const part = free2.part;
+
+  // This causes their range to extend to include this hook call, and in turn for the memoization to be pruned
+  useHook();
+  const object = useMemo(() => {
+    const x = makeObject_Primitives();
+    x.value = props.value;
+    mutate(x, free, part);
+    return x;
+  }, [props.value]);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};
+
+```
+
+## Code
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees:false
+import { useMemo } from "react";
+import {
+  identity,
+  makeObject_Primitives,
+  mutate,
+  useHook,
+} from "shared-runtime";
+
+function Component(props) {
+  const free = makeObject_Primitives();
+  const free2 = makeObject_Primitives();
+  const part = free2.part;
+
+  useHook();
+  let t39;
+
+  const x = makeObject_Primitives();
+  x.value = props.value;
+  mutate(x, free, part);
+  t39 = x;
+  const object = t39;
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) {"a":0,"b":"value1","c":true,"value":42,"wat0":"joe"}

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-dont-preserve-memoization-guarantees.js

@@ -0,0 +1,30 @@
+// @enablePreserveExistingMemoizationGuarantees:false
+import { useMemo } from "react";
+import {
+  identity,
+  makeObject_Primitives,
+  mutate,
+  useHook,
+} from "shared-runtime";
+
+function Component(props) {
+  // With the feature disabled these variables are inferred as being mutated inside the useMemo block
+  const free = makeObject_Primitives();
+  const free2 = makeObject_Primitives();
+  const part = free2.part;
+
+  // This causes their range to extend to include this hook call, and in turn for the memoization to be pruned
+  useHook();
+  const object = useMemo(() => {
+    const x = makeObject_Primitives();
+    x.value = props.value;
+    mutate(x, free, part);
+    return x;
+  }, [props.value]);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.expect.md

@@ -0,0 +1,95 @@
+
+## Input
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees
+import { useMemo } from "react";
+import {
+  identity,
+  makeObject_Primitives,
+  mutate,
+  useHook,
+} from "shared-runtime";
+
+function Component(props) {
+  const free = makeObject_Primitives();
+  const free2 = makeObject_Primitives();
+  const part = free2.part;
+  useHook();
+  const object = useMemo(() => {
+    const x = makeObject_Primitives();
+    x.value = props.value;
+    mutate(x, free, part);
+    return x;
+  }, [props.value]);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};
+
+```
+
+## Code
+
+```javascript
+// @enablePreserveExistingMemoizationGuarantees
+import { useMemo, unstable_useMemoCache as useMemoCache } from "react";
+import {
+  identity,
+  makeObject_Primitives,
+  mutate,
+  useHook,
+} from "shared-runtime";
+
+function Component(props) {
+  const $ = useMemoCache(4);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = makeObject_Primitives();
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const free = t0;
+  let t1;
+  if ($[1] === Symbol.for("react.memo_cache_sentinel")) {
+    t1 = makeObject_Primitives();
+    $[1] = t1;
+  } else {
+    t1 = $[1];
+  }
+  const free2 = t1;
+  const part = free2.part;
+  useHook();
+
+  props.value;
+  free;
+  part;
+  let t44;
+  let x;
+  if ($[2] !== props.value) {
+    x = makeObject_Primitives();
+    x.value = props.value;
+    mutate(x, free, part);
+    $[2] = props.value;
+    $[3] = x;
+  } else {
+    x = $[3];
+  }
+  t44 = x;
+  const object = t44;
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};
+
+```
+      
+### Eval output
+(kind: ok) {"a":0,"b":"value1","c":true,"value":42,"wat0":"joe"}

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/useMemo-mabye-modified-free-variable-preserve-memoization-guarantees.js

@@ -0,0 +1,27 @@
+// @enablePreserveExistingMemoizationGuarantees
+import { useMemo } from "react";
+import {
+  identity,
+  makeObject_Primitives,
+  mutate,
+  useHook,
+} from "shared-runtime";
+
+function Component(props) {
+  const free = makeObject_Primitives();
+  const free2 = makeObject_Primitives();
+  const part = free2.part;
+  useHook();
+  const object = useMemo(() => {
+    const x = makeObject_Primitives();
+    x.value = props.value;
+    mutate(x, free, part);
+    return x;
+  }, [props.value]);
+  return object;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ value: 42 }],
+};

compiler/packages/sprout/src/shared-runtime.ts

@@ -153,6 +153,10 @@ export function throwInput(x: Object): never {
   throw x;
 }
 
+export function useHook(): Object {
+  return makeObject_Primitives();
+}
+
 const noAliasObject = Object.freeze({});
 export function useNoAlias(...args: Array<any>): object {
   return noAliasObject;