ValidateHooksUsage

Adds a validation pass to check that the only thing you can do with hooks is 
call them. A follow-up PR (still early WIP) will check the other aspect of the 
rules of hooks, that they are not called conditionally. That's a more involved 
algorithm.

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

@@ -113,7 +113,7 @@ export async function compile(
             },
           ],
         ]),
-        inlineUseMemo: true,
+        validateHooksUsage: true,
       },
       logger: null,
       gating,

compiler/forget/src/CompilerPipeline.ts

@@ -13,6 +13,7 @@ import {
   mergeConsecutiveBlocks,
   ReactiveFunction,
   validateConsistentIdentifiers,
+  validateHooksUsage,
   validateTerminalSuccessors,
 } from "./HIR";
 import { Environment, EnvironmentConfig } from "./HIR/Environment";
@@ -73,17 +74,21 @@ export function* run(
   enterSSA(hir);
   yield log({ kind: "hir", name: "SSA", value: hir });
 
-  validateConsistentIdentifiers(hir);
-
   eliminateRedundantPhi(hir);
   yield log({ kind: "hir", name: "EliminateRedundantPhi", value: hir });
 
+  validateConsistentIdentifiers(hir);
+
   constantPropagation(hir);
   yield log({ kind: "hir", name: "ConstantPropagation", value: hir });
 
   inferTypes(hir);
   yield log({ kind: "hir", name: "InferTypes", value: hir });
 
+  if (env.validateHooksUsage) {
+    validateHooksUsage(hir);
+  }
+
   dropMemoCalls(hir);
   yield log({ kind: "hir", name: "DropMemoCalls", value: hir });
 

compiler/forget/src/HIR/Environment.ts

@@ -19,11 +19,11 @@ import {
   Effect,
   FunctionType,
   IdentifierId,
-  makeBlockId,
-  makeIdentifierId,
   ObjectType,
   PolyType,
   ValueKind,
+  makeBlockId,
+  makeIdentifierId,
 } from "./HIR";
 import { Hook } from "./Hooks";
 import { FunctionSignature, ShapeRegistry } from "./ObjectShape";
@@ -40,6 +40,7 @@ const HOOK_PATTERN = /^_?use/;
 export type EnvironmentConfig = Partial<{
   customHooks: Map<string, Hook>;
   memoizeJsxElements: boolean;
+  validateHooksUsage: boolean;
 }>;
 
 export class Environment {
@@ -47,6 +48,7 @@ export class Environment {
   #shapes: ShapeRegistry;
   #nextIdentifer: number = 0;
   #nextBlock: number = 0;
+  validateHooksUsage: boolean;
 
   constructor(config: EnvironmentConfig | null) {
     this.#shapes = DEFAULT_SHAPES;
@@ -66,6 +68,7 @@ export class Environment {
     } else {
       this.#globals = DEFAULT_GLOBALS;
     }
+    this.validateHooksUsage = config?.validateHooksUsage ?? false;
   }
 
   get nextIdentifierId(): IdentifierId {

compiler/forget/src/HIR/ValidateHooksUsage.ts

@@ -0,0 +1,90 @@
+/**
+ * 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,
+  CompilerErrorDetail,
+  ErrorSeverity,
+} from "../CompilerError";
+import { hasBackEdge } from "../Optimization/DeadCodeElimination";
+import { HIRFunction, IdentifierId, Place, isHookType } from "./HIR";
+import { eachInstructionValueOperand, eachTerminalOperand } from "./visitors";
+
+/**
+ * Validates that the function honors the [Rules of Hooks](https://react.dev/warnings/invalid-hook-call-warning)
+ * rule that hooks may only be called and not otherwise referenced as first-class values.
+ */
+export function validateHooksUsage(fn: HIRFunction): void {
+  const errors = new CompilerError();
+  const pushError = (place: Place): void => {
+    errors.pushErrorDetail(
+      new CompilerErrorDetail({
+        codeframe: null,
+        description: null,
+        reason:
+          "Hooks may not be referenced as normal values, they must be called. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning)",
+        loc: typeof place.loc !== "symbol" ? place.loc : null,
+        severity: ErrorSeverity.InvalidInput,
+      })
+    );
+  };
+
+  const hooks: Set<IdentifierId> = new Set();
+  const hasLoop = hasBackEdge(fn);
+
+  let size = hooks.size;
+  do {
+    size = hooks.size;
+    for (const [, block] of fn.body.blocks) {
+      for (const phi of block.phis) {
+        let possibleHook = false;
+        for (const [, predecessor] of phi.operands) {
+          if (hooks.has(predecessor.id)) {
+            possibleHook = true;
+            break;
+          }
+        }
+        if (possibleHook) {
+          hooks.add(phi.id.id);
+        }
+      }
+
+      for (const instr of block.instructions) {
+        if (
+          instr.value.kind === "LoadGlobal" &&
+          isHookType(instr.lvalue.identifier)
+        ) {
+          hooks.add(instr.lvalue.identifier.id);
+        } else if (instr.value.kind === "CallExpression") {
+          for (const operand of eachInstructionValueOperand(instr.value)) {
+            if (operand === instr.value.callee) {
+              continue;
+            }
+            if (hooks.has(operand.identifier.id)) {
+              pushError(operand);
+            }
+          }
+        } else {
+          for (const operand of eachInstructionValueOperand(instr.value)) {
+            if (hooks.has(operand.identifier.id)) {
+              pushError(operand);
+            }
+          }
+        }
+      }
+      for (const operand of eachTerminalOperand(block.terminal)) {
+        if (hooks.has(operand.identifier.id)) {
+          pushError(operand);
+        }
+      }
+    }
+  } while (hooks.size > size && hasLoop);
+
+  if (errors.hasErrors()) {
+    throw errors;
+  }
+}

compiler/forget/src/HIR/index.ts

@@ -19,4 +19,5 @@ export { Hook } from "./Hooks";
 export { mergeConsecutiveBlocks } from "./MergeConsecutiveBlocks";
 export { printFunction, printHIR } from "./PrintHIR";
 export { validateConsistentIdentifiers } from "./ValidateConsistentIdentifiers";
+export { validateHooksUsage } from "./ValidateHooksUsage";
 export { validateTerminalSuccessors } from "./ValidateTerminalSuccessors";

compiler/forget/src/Optimization/DeadCodeElimination.ts

@@ -247,7 +247,7 @@ function pruneableValue(value: InstructionValue, state: State): boolean {
   }
 }
 
-function hasBackEdge(fn: HIRFunction): boolean {
+export function hasBackEdge(fn: HIRFunction): boolean {
   const visited = new Set<BlockId>();
   for (const [blockId, block] of fn.body.blocks) {
     for (const predId of block.preds) {

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

@@ -55,6 +55,7 @@ describe("React Forget", () => {
                   },
                 ],
               ]),
+              validateHooksUsage: true,
             },
             logger: null,
             gating: options.gating,

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-assign-hook-to-local.expect.md

@@ -0,0 +1,20 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const x = useRef;
+  const ref = x(null);
+  return ref.current;
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidInput: Hooks may not be referenced as normal values, they must be called. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning) (2:2)
+```
+          
+      

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-pass-hook-as-call-arg.expect.md

@@ -0,0 +1,18 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  return foo(useFoo);
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidInput: Hooks may not be referenced as normal values, they must be called. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning) (2:2)
+```
+          
+      

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-pass-hook-as-call-arg.js

@@ -0,0 +1,3 @@
+function Component(props) {
+  return foo(useFoo);
+}

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-pass-hook-as-prop.expect.md

@@ -0,0 +1,18 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  return <Child foo={useFoo} />;
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidInput: Hooks may not be referenced as normal values, they must be called. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning) (2:2)
+```
+          
+      

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-pass-hook-as-prop.js

@@ -0,0 +1,3 @@
+function Component(props) {
+  return <Child foo={useFoo} />;
+}

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-ternary-with-hook-values.expect.md

@@ -0,0 +1,21 @@
+
+## Input
+
+```javascript
+function Component(props) {
+  const x = props.cond ? useA : useB;
+  return x();
+}
+
+```
+
+
+## Error
+
+```
+[ReactForget] InvalidInput: Hooks may not be referenced as normal values, they must be called. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning) (2:2)
+
+[ReactForget] InvalidInput: Hooks may not be referenced as normal values, they must be called. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning) (2:2)
+```
+          
+      

compiler/forget/src/__tests__/fixtures/compiler/error.invalid-ternary-with-hook-values.js

@@ -0,0 +1,4 @@
+function Component(props) {
+  const x = props.cond ? useA : useB;
+  return x();
+}

compiler/forget/src/__tests__/fixtures/compiler/useRef-rename-mutable.expect.md

@@ -1,22 +0,0 @@
-
-## Input
-
-```javascript
-function Component(props) {
-  const x = useRef;
-  const ref = x(null);
-  return ref.current;
-}
-
-```
-
-## Code
-
-```javascript
-function Component(props) {
-  const ref = useRef(null);
-  return ref.current;
-}
-
-```
-