Option to assume hooks follow the rules

Adds a new feature flag which tells the compiler to assume that hooks follow the 
Rules of React. Specifically, the idea that since any hook could be wrapped in a 
giant `useMemo()` call, all arguments to hooks have to be treated as if they're 
owned by React — and therefore become immutable — and that the return value of 
the hook is immutable. 

Our default is to assume that hooks break the rules, but in practice nearly 
every component follows them.

compiler/forget/packages/snap/src/compiler-worker.ts

@@ -93,6 +93,7 @@ export async function compile(
     let instrumentForget = null;
     let panicOnBailout = true;
     let memoizeJsxElements = true;
+    let enableAssumeHooksFollowRulesOfReact = false;
     if (firstLine.indexOf("@forgetDirective") !== -1) {
       enableOnlyOnUseForgetDirective = true;
     }
@@ -120,6 +121,9 @@ export async function compile(
     if (firstLine.indexOf("@memoizeJsxElements false") !== -1) {
       memoizeJsxElements = false;
     }
+    if (firstLine.indexOf("@enableAssumeHooksFollowRulesOfReact true") !== -1) {
+      enableAssumeHooksFollowRulesOfReact = true;
+    }
 
     const language = parseLanguage(firstLine);
 
@@ -139,6 +143,7 @@ export async function compile(
         ]),
         validateHooksUsage: true,
         enableFunctionCallSignatureOptimizations: true,
+        enableAssumeHooksFollowRulesOfReact,
         inlineUseMemo: true,
         memoizeJsxElements,
       },

compiler/forget/src/HIR/Environment.ts

@@ -17,12 +17,10 @@ import {
 import {
   BlockId,
   BuiltInType,
-  Effect,
   FunctionType,
   IdentifierId,
   ObjectType,
   PolyType,
-  ValueKind,
   makeBlockId,
   makeIdentifierId,
 } from "./HIR";
@@ -38,10 +36,50 @@ import { FunctionSignature, ShapeRegistry } from "./ObjectShape";
 //   missing some recursive Object / Function shapeIds
 export type EnvironmentConfig = Partial<{
   customHooks: Map<string, Hook>;
+
+  /**
+   * Enable memoization of JSX elements in addition to other types of values. When disabled,
+   * other types (objects, arrays, call expressions, etc) are memoized, but not known JSX
+   * values.
+   *
+   * Defaults to true
+   */
   memoizeJsxElements: boolean;
+
+  /**
+   * Enable validation of hooks to partially check that the component honors the rules of hooks.
+   * When disabled, the component is assumed to follow the rules (though the Babel plugin looks
+   * for suppressions of the lint rule).
+   *
+   * Defaults to false
+   */
   validateHooksUsage: boolean;
+
+  /**
+   * Enable inlining of `useMemo()` function expressions so that they can be more optimally
+   * compiled.
+   *
+   * Defaults to false
+   */
   inlineUseMemo: boolean;
+
+  /**
+   * Enable optimizations based on the signature of (non-method) built-in function calls.
+   *
+   * Defaults to false
+   */
   enableFunctionCallSignatureOptimizations: boolean;
+
+  /**
+   * When enabled, the compiler assumes that hooks follow the Rules of React:
+   * - Hooks may memoize computation based on any of their parameters, thus
+   *   any arguments to a hook are assumed frozen after calling the hook.
+   * - Hooks may memoize the result they return, thus the return value is
+   *   assumed frozen.
+
+   * Defaults to false
+   */
+  enableAssumeHooksFollowRulesOfReact: boolean;
 }>;
 
 export class Environment {
@@ -51,6 +89,7 @@ export class Environment {
   #nextBlock: number = 0;
   validateHooksUsage: boolean;
   enableFunctionCallSignatureOptimizations: boolean;
+  enableAssumeHooksFollowRulesOfReact: boolean;
   #contextIdentifiers: Set<t.Identifier>;
 
   constructor(
@@ -77,6 +116,8 @@ export class Environment {
     this.validateHooksUsage = config?.validateHooksUsage ?? false;
     this.enableFunctionCallSignatureOptimizations =
       config?.enableFunctionCallSignatureOptimizations ?? false;
+    this.enableAssumeHooksFollowRulesOfReact =
+      config?.enableAssumeHooksFollowRulesOfReact ?? false;
     this.#contextIdentifiers = contextIdentifiers;
   }
 
@@ -98,12 +139,7 @@ export class Environment {
       if (isHookName(name)) {
         return {
           kind: "Hook",
-          definition: {
-            kind: "Custom",
-            name,
-            effectKind: Effect.Mutate,
-            valueKind: ValueKind.Mutable,
-          },
+          definition: null,
         };
       } else {
         log(() => `Undefined global '${name}'`);

compiler/forget/src/HIR/Types.ts

@@ -20,7 +20,7 @@ export type Type =
 export type PrimitiveType = { kind: "Primitive" };
 export type HookType = {
   kind: "Hook";
-  definition: Hook;
+  definition: Hook | null;
 };
 
 /**

compiler/forget/src/Inference/DropMemoCalls.ts

@@ -14,7 +14,7 @@ export default function (func: HIRFunction): void {
         case "CallExpression": {
           if (isHookType(instr.value.callee.identifier)) {
             const name = (instr.value.callee.identifier.type as HookType)
-              .definition.name;
+              .definition?.name;
             if (name === "useMemo") {
               const [fn] = instr.value.args;
 

compiler/forget/src/Inference/InferReferenceEffects.ts

@@ -708,13 +708,18 @@ function inferBlock(
         continue;
       }
       case "CallExpression": {
-        const hook =
-          instrValue.callee.identifier.type.kind === "Hook"
-            ? instrValue.callee.identifier.type.definition
-            : null;
-        if (hook !== null) {
-          effectKind = hook.effectKind;
-          valueKind = hook.valueKind;
+        if (instrValue.callee.identifier.type.kind === "Hook") {
+          const definition = instrValue.callee.identifier.type.definition;
+          if (definition !== null) {
+            effectKind = definition.effectKind;
+            valueKind = definition.valueKind;
+          } else if (env.enableAssumeHooksFollowRulesOfReact) {
+            effectKind = Effect.Freeze;
+            valueKind = ValueKind.Frozen;
+          } else {
+            effectKind = Effect.Mutate;
+            valueKind = ValueKind.Mutable;
+          }
           break;
         }
 

compiler/forget/src/__tests__/compiler-test.ts

@@ -60,6 +60,9 @@ describe("React Forget", () => {
               validateHooksUsage: true,
               inlineUseMemo: options.environment?.inlineUseMemo ?? false,
               enableFunctionCallSignatureOptimizations: true,
+              enableAssumeHooksFollowRulesOfReact:
+                options.environment?.enableAssumeHooksFollowRulesOfReact ??
+                false,
             },
             logger: null,
             gating: options.gating,

compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.expect.md

@@ -0,0 +1,48 @@
+
+## Input
+
+```javascript
+// @enableAssumeHooksFollowRulesOfReact true
+function Component(props) {
+  const x = {};
+  // In enableAssumeHooksFollowRulesOfReact mode hooks freeze their inputs and return frozen values
+  const y = useFoo(x);
+  // Thus both x and y are frozen here, and x can be independently memoized
+  bar(x, y);
+  return [x, y];
+}
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react"; // @enableAssumeHooksFollowRulesOfReact true
+function Component(props) {
+  const $ = useMemoCache(3);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = {};
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const x = t0;
+
+  const y = useFoo(x);
+
+  bar(x, y);
+  const c_1 = $[1] !== y;
+  let t1;
+  if (c_1) {
+    t1 = [x, y];
+    $[1] = y;
+    $[2] = t1;
+  } else {
+    t1 = $[2];
+  }
+  return t1;
+}
+
+```
+      

compiler/forget/src/__tests__/fixtures/compiler/immutable-hooks.js

@@ -0,0 +1,9 @@
+// @enableAssumeHooksFollowRulesOfReact true
+function Component(props) {
+  const x = {};
+  // In enableAssumeHooksFollowRulesOfReact mode hooks freeze their inputs and return frozen values
+  const y = useFoo(x);
+  // Thus both x and y are frozen here, and x can be independently memoized
+  bar(x, y);
+  return [x, y];
+}

compiler/forget/src/__tests__/test-utils/generateTestsFromFixtures.ts

@@ -102,6 +102,7 @@ export default function generateTestsFromFixtures(
         let inlineUseMemo = true;
         let panicOnBailout = true;
         let memoizeJsxElements = true;
+        let enableAssumeHooksFollowRulesOfReact = false;
 
         if (inputFile != null) {
           input = fs.readFileSync(inputFile, "utf8");
@@ -145,13 +146,23 @@ export default function generateTestsFromFixtures(
           if (lines[0]!.indexOf("@memoizeJsxElements false") !== -1) {
             memoizeJsxElements = false;
           }
+          if (
+            lines[0]!.indexOf("@enableAssumeHooksFollowRulesOfReact true") !==
+            -1
+          ) {
+            enableAssumeHooksFollowRulesOfReact = true;
+          }
         }
 
         testCommand(basename, () => {
           let receivedOutput;
           if (input !== null) {
             receivedOutput = transform(input, basename, {
-              environment: { inlineUseMemo, memoizeJsxElements },
+              environment: {
+                inlineUseMemo,
+                memoizeJsxElements,
+                enableAssumeHooksFollowRulesOfReact,
+              },
               logger: null,
               debug,
               enableOnlyOnUseForgetDirective,