[validation] Runtime validation for hook calls

--- 

I modeled guards as try-finally blocks to be extremely explicit. An alternative 
implementation could flatten all nested hooks and only set / restore hook guards 
when entering / exiting a React function (i.e. hook or component) -- this 
alternative approach would be the easiest to represent as a separate pass 

```js 

// source 

function Foo() { 

const result = useHook(useContext(Context)); 

... 

} 

// current output 

function Foo() { 

try { 

pushHookGuard(); 

const result = (() => { 

try { 

pushEnableHook(); 

return useHook((() => { 

try { 

pushEnableHook(); 

return useContext(Context); 

} finally { 

popEnableHook(); 

} 

})()); 

} finally { 

popEnableHook(); 

}; 

})(); 

// ... 

} finally { 

popHookGuard(); 

} 

} 

// alternative output 

function Foo() { 

try { 

// check current is not lazyDispatcher; 

// save originalDispatcher, set lazyDispatcher 

pushHookGuard(); 

allowHook(); // always set originalDispatcher 

const t0 = useContext(Context); 

disallowHook(); // always set LazyDispatcher 

allowHook(); // always set originalDispatcher 

const result = useHook(t0); 

disallowHook(); // always set LazyDispatcher 

// ... 

} finally { 

popHookGuard(); // restore originalDispatcher 

} 

} 

``` 

Checked that IG Web works as expected 

Unless I add a sneaky useState: 

<img width="705" alt="Screenshot 2023-12-05 at 6 44 59 PM" 
src="https://github.com/facebook/react-forget/assets/34200447/3790bd76-7d71-44b5-a62e-f53256fb5736">

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

@@ -312,6 +312,13 @@ export function compileProgram(
       );
       externalFunctions.push(enableEmitFreeze);
     }
+
+    if (options.environment?.enableEmitHookGuards != null) {
+      const enableEmitHookGuards = tryParseExternalFunction(
+        options.environment.enableEmitHookGuards
+      );
+      externalFunctions.push(enableEmitHookGuards);
+    }
   } catch (err) {
     handleError(err, pass, null);
     return;

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

@@ -191,6 +191,8 @@ const EnvironmentConfigSchema = z.object({
    */
   enableEmitFreeze: ExternalFunctionSchema.nullish(),
 
+  enableEmitHookGuards: ExternalFunctionSchema.nullish(),
+
   /*
    * Enables instrumentation codegen. This emits a dev-mode only call to an
    * instrumentation function, for components and hooks that Forget compiles.

compiler/packages/babel-plugin-react-forget/src/ReactiveScopes/CodegenReactiveFunction.ts

@@ -8,7 +8,7 @@
 import * as t from "@babel/types";
 import { pruneUnusedLValues, pruneUnusedLabels, renameVariables } from ".";
 import { CompilerError, ErrorSeverity } from "../CompilerError";
-import { Environment } from "../HIR";
+import { Environment, EnvironmentConfig, ExternalFunction } from "../HIR";
 import {
   BlockId,
   GeneratedSource,
@@ -30,10 +30,12 @@ import {
   ReactiveValue,
   SourceLocation,
   SpreadPattern,
+  getHookKind,
 } from "../HIR/HIR";
 import { printPlace } from "../HIR/PrintHIR";
 import { eachPatternOperand } from "../HIR/visitors";
 import { Err, Ok, Result } from "../Utils/Result";
+import { GuardKind } from "../Utils/RuntimeDiagnosticConstants";
 import { assertExhaustive } from "../Utils/utils";
 import { buildReactiveFunction } from "./BuildReactiveFunction";
 import { SINGLE_CHILD_FBT_TAGS } from "./MemoizeFbtOperandsInSameScope";
@@ -86,6 +88,18 @@ export function codegenFunction(
     );
     compiled.body.body.unshift(test);
   }
+
+  const hookGuard = fn.env.config.enableEmitHookGuards;
+  if (hookGuard != null) {
+    compiled.body = t.blockStatement([
+      createHookGuard(
+        hookGuard,
+        compiled.body.body,
+        GuardKind.PushHookGuard,
+        GuardKind.PopHookGuard
+      ),
+    ]);
+  }
   return compileResult;
 }
 
@@ -970,7 +984,6 @@ function withLoc<T extends (...args: any[]) => t.Node>(
 }
 
 const createBinaryExpression = withLoc(t.binaryExpression);
-const createCallExpression = withLoc(t.callExpression);
 const createExpressionStatement = withLoc(t.expressionStatement);
 const _createLabelledStatement = withLoc(t.labeledStatement);
 const createVariableDeclaration = withLoc(t.variableDeclaration);
@@ -989,6 +1002,77 @@ const createJsxText = withLoc(t.jsxText);
 const createJsxClosingElement = withLoc(t.jsxClosingElement);
 const createStringLiteral = withLoc(t.stringLiteral);
 
+function createHookGuard(
+  guard: ExternalFunction,
+  stmts: t.Statement[],
+  before: GuardKind,
+  after: GuardKind
+): t.TryStatement {
+  function createHookGuardImpl(kind: number): t.ExpressionStatement {
+    return t.expressionStatement(
+      t.callExpression(t.identifier(guard.importSpecifierName), [
+        t.numericLiteral(kind),
+      ])
+    );
+  }
+
+  return t.tryStatement(
+    t.blockStatement([createHookGuardImpl(before), ...stmts]),
+    null,
+    t.blockStatement([createHookGuardImpl(after)])
+  );
+}
+
+/**
+ * Create a call expression.
+ * If enableEmitHookGuards is set and the callExpression is a hook call,
+ * the following transform will be made.
+ * ```js
+ * // source
+ * useHook(arg1, arg2)
+ *
+ * // codegen
+ * (() => {
+ *   try {
+ *     $dispatcherGuard(PUSH_EXPECT_HOOK);
+ *     return useHook(arg1, arg2);
+ *   } finally {
+ *     $dispatcherGuard(POP_EXPECT_HOOK);
+ *   }
+ * })()
+ * ```
+ */
+function createCallExpression(
+  config: EnvironmentConfig,
+  callee: t.Expression,
+  args: Array<t.Expression | t.SpreadElement>,
+  loc: SourceLocation | null,
+  isHook: boolean
+): t.CallExpression {
+  const callExpr = t.callExpression(callee, args);
+  if (loc != null && loc != GeneratedSource) {
+    callExpr.loc = loc;
+  }
+
+  const hookGuard = config.enableEmitHookGuards;
+  if (hookGuard != null && isHook) {
+    const iife = t.arrowFunctionExpression(
+      [],
+      t.blockStatement([
+        createHookGuard(
+          hookGuard,
+          [t.returnStatement(callExpr)],
+          GuardKind.AllowHook,
+          GuardKind.DisallowHook
+        ),
+      ])
+    );
+    return t.callExpression(iife, []);
+  } else {
+    return callExpr;
+  }
+}
+
 type Temporaries = Map<IdentifierId, t.Expression | t.JSXText | null>;
 
 function codegenLabel(id: BlockId): string {
@@ -1091,9 +1175,16 @@ function codegenInstructionValue(
       break;
     }
     case "CallExpression": {
+      const isHook = getHookKind(cx.env, instrValue.callee.identifier) != null;
       const callee = codegenPlaceToExpression(cx, instrValue.callee);
       const args = instrValue.args.map((arg) => codegenArgument(cx, arg));
-      value = createCallExpression(instrValue.loc, callee, args);
+      value = createCallExpression(
+        cx.env.config,
+        callee,
+        args,
+        instrValue.loc,
+        isHook
+      );
       break;
     }
     case "OptionalExpression": {
@@ -1147,6 +1238,8 @@ function codegenInstructionValue(
       break;
     }
     case "MethodCall": {
+      const isHook =
+        getHookKind(cx.env, instrValue.property.identifier) != null;
       const memberExpr = codegenPlaceToExpression(cx, instrValue.property);
       CompilerError.invariant(
         t.isMemberExpression(memberExpr) ||
@@ -1175,7 +1268,13 @@ function codegenInstructionValue(
         }
       );
       const args = instrValue.args.map((arg) => codegenArgument(cx, arg));
-      value = createCallExpression(instrValue.loc, memberExpr, args);
+      value = createCallExpression(
+        cx.env.config,
+        memberExpr,
+        args,
+        instrValue.loc,
+        isHook
+      );
       break;
     }
     case "NewExpression": {

compiler/packages/babel-plugin-react-forget/src/Utils/RuntimeDiagnosticConstants.ts

@@ -0,0 +1,14 @@
+/*
+ * 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.
+ */
+
+// WARNING: ensure this is synced with enum values in react-forget-runtime:GuardKind
+export enum GuardKind {
+  PushHookGuard = 0,
+  PopHookGuard = 1,
+  AllowHook = 2,
+  DisallowHook = 3,
+}

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/flag-enable-emit-hook-guards.expect.md

@@ -0,0 +1,132 @@
+
+## Input
+
+```javascript
+// @enableEmitHookGuards
+import { createContext, useContext, useEffect, useState } from "react";
+import {
+  CONST_STRING0,
+  ObjectWithHooks,
+  getNumber,
+  identity,
+  print,
+} from "shared-runtime";
+
+const MyContext = createContext("my context value");
+function Component({ value }) {
+  print(identity(CONST_STRING0));
+  const [state, setState] = useState(getNumber());
+  print(value, state);
+  useEffect(() => {
+    if (state === 4) {
+      setState(5);
+    }
+  }, [state]);
+  print(identity(value + state));
+  return ObjectWithHooks.useIdentity(useContext(MyContext));
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  args: [{ value: 0 }],
+};
+
+```
+
+## Code
+
+```javascript
+import { $dispatcherGuard } from "react-forget-runtime"; // @enableEmitHookGuards
+import {
+  createContext,
+  useContext,
+  useEffect,
+  useState,
+  unstable_useMemoCache as useMemoCache,
+} from "react";
+import {
+  CONST_STRING0,
+  ObjectWithHooks,
+  getNumber,
+  identity,
+  print,
+} from "shared-runtime";
+
+const MyContext = createContext("my context value");
+function Component(t47) {
+  try {
+    $dispatcherGuard(0);
+    const $ = useMemoCache(4);
+    const { value } = t47;
+    print(identity(CONST_STRING0));
+    let t0;
+    if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+      t0 = getNumber();
+      $[0] = t0;
+    } else {
+      t0 = $[0];
+    }
+    const [state, setState] = (() => {
+      try {
+        $dispatcherGuard(2);
+        return useState(t0);
+      } finally {
+        $dispatcherGuard(3);
+      }
+    })();
+    print(value, state);
+    let t1;
+    let t2;
+    if ($[1] !== state) {
+      t1 = () => {
+        if (state === 4) {
+          setState(5);
+        }
+      };
+
+      t2 = [state];
+      $[1] = state;
+      $[2] = t1;
+      $[3] = t2;
+    } else {
+      t1 = $[2];
+      t2 = $[3];
+    }
+    (() => {
+      try {
+        $dispatcherGuard(2);
+        return useEffect(t1, t2);
+      } finally {
+        $dispatcherGuard(3);
+      }
+    })();
+    print(identity(value + state));
+    return (() => {
+      try {
+        $dispatcherGuard(2);
+        return ObjectWithHooks.useIdentity(
+          (() => {
+            try {
+              $dispatcherGuard(2);
+              return useContext(MyContext);
+            } finally {
+              $dispatcherGuard(3);
+            }
+          })()
+        );
+      } finally {
+        $dispatcherGuard(3);
+      }
+    })();
+  } finally {
+    $dispatcherGuard(1);
+  }
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  args: [{ value: 0 }],
+};
+
+```
+      

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/flag-enable-emit-hook-guards.ts

@@ -0,0 +1,28 @@
+// @enableEmitHookGuards
+import { createContext, useContext, useEffect, useState } from "react";
+import {
+  CONST_STRING0,
+  ObjectWithHooks,
+  getNumber,
+  identity,
+  print,
+} from "shared-runtime";
+
+const MyContext = createContext("my context value");
+function Component({ value }) {
+  print(identity(CONST_STRING0));
+  const [state, setState] = useState(getNumber());
+  print(value, state);
+  useEffect(() => {
+    if (state === 4) {
+      setState(5);
+    }
+  }, [state]);
+  print(identity(value + state));
+  return ObjectWithHooks.useIdentity(useContext(MyContext));
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  args: [{ value: 0 }],
+};

compiler/packages/fixture-test-utils/src/compiler-utils.ts

@@ -33,6 +33,7 @@ export function transformFixtureInput(
   let gating = null;
   let enableEmitInstrumentForget = null;
   let enableEmitFreeze = null;
+  let enableEmitHookGuards = null;
   let compilationMode: CompilationMode = "all";
   let enableUseMemoCachePolyfill = false;
   let panicThreshold: PanicThresholdOptions = "ALL_ERRORS";
@@ -70,6 +71,12 @@ export function transformFixtureInput(
       importSpecifierName: "makeReadOnly",
     };
   }
+  if (firstLine.includes("@enableEmitHookGuards")) {
+    enableEmitHookGuards = {
+      source: "react-forget-runtime",
+      importSpecifierName: "$dispatcherGuard",
+    };
+  }
   if (firstLine.includes("@enableUseMemoCachePolyfill")) {
     enableUseMemoCachePolyfill = true;
   }
@@ -116,6 +123,7 @@ export function transformFixtureInput(
         ]),
         enableEmitFreeze,
         enableEmitInstrumentForget,
+        enableEmitHookGuards,
         assertValidMutableRanges: true,
       },
       compilationMode,

compiler/packages/react-forget-runtime/src/index.ts

@@ -5,6 +5,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+import invariant from "invariant";
 import * as React from "react";
 
 const {
@@ -41,7 +42,7 @@ export function $read(memoCache: MemoCache, index: number) {
   return value;
 }
 
-const LazyGuardDispatcher: { [key: string]: () => never } = {};
+const LazyGuardDispatcher: { [key: string]: (...args: Array<any>) => any } = {};
 [
   "readContext",
   "useCallback",
@@ -66,26 +67,121 @@ const LazyGuardDispatcher: { [key: string]: () => never } = {};
   "useCacheRefresh",
 ].forEach((name) => {
   LazyGuardDispatcher[name] = () => {
-    throw new Error(`Cannot call ${name} within ReactForget lazy block.`);
+    throw new Error(
+      `[React] Unexpected React hook call (${name}) from a React Forget compiled function. ` +
+        "Check that all hooks are called directly and named according to convention ('use[A-Z]') "
+    );
   };
 });
 
 let originalDispatcher: unknown = null;
 
-export function $startLazy() {
-  if (originalDispatcher !== null) {
-    throw new Error("unexpected startLazy with dispatcher set");
+// Allow guards are not emitted for useMemoCache
+LazyGuardDispatcher["useMemoCache"] = (count: number) => {
+  if (originalDispatcher == null) {
+    throw new Error(
+      "React Forget internal invariant violation: unexpected null dispatcher"
+    );
+  } else {
+    return (originalDispatcher as any).useMemoCache(count);
   }
-  originalDispatcher = ReactCurrentDispatcher.current;
-  ReactCurrentDispatcher.current = LazyGuardDispatcher;
+};
+
+enum GuardKind {
+  PushGuardContext = 0,
+  PopGuardContext = 1,
+  PushExpectHook = 2,
+  PopExpectHook = 3,
 }
 
-export function $endLazy() {
-  if (originalDispatcher === null) {
-    throw new Error("unexpected endLazy with dispatcher not set");
+function setCurrent(newDispatcher: any) {
+  ReactCurrentDispatcher.current = newDispatcher;
+  return ReactCurrentDispatcher.current;
+}
+
+const guardFrames: Array<unknown> = [];
+
+/**
+ * When `enableEmitHookGuards` is set, this does runtime validation
+ * of the no-conditional-hook-calls rule.
+ * As Forget needs to statically understand which calls to move out of
+ * conditional branches (i.e. Forget cannot memoize the results of hook
+ * calls), its understanding of "the rules of React" are more restrictive.
+ * This validation throws on unsound inputs at runtime.
+ *
+ * Components should only be invoked through React as Forget could memoize
+ * the call to AnotherComponent, introducing conditional hook calls in its
+ * compiled output.
+ * ```js
+ * function Invalid(props) {
+ *  const myJsx = AnotherComponent(props);
+ *  return <div> { myJsx } </div>;
+ * }
+ *
+ * Hooks must be named as hooks.
+ * ```js
+ * const renamedHook = useState;
+ * function Invalid() {
+ *   const [state, setState] = renamedHook(0);
+ * }
+ * ```
+ *
+ * Hooks must be directly called.
+ * ```
+ * function call(fn) {
+ *  return fn();
+ * }
+ * function Invalid() {
+ *   const result = call(useMyHook);
+ * }
+ * ```
+ */
+export function $dispatcherGuard(kind: GuardKind) {
+  const curr = ReactCurrentDispatcher.current;
+  if (kind === GuardKind.PushGuardContext) {
+    // Push before checking invariant or errors
+    guardFrames.push(curr);
+
+    if (guardFrames.length === 1) {
+      // save if we're the first guard on the stack
+      originalDispatcher = curr;
+    }
+
+    if (curr === LazyGuardDispatcher) {
+      throw new Error(
+        `[React] Unexpected call to custom hook or component from a React Forget compiled function. ` +
+          "Check that (1) all hooks are called directly and named according to convention ('use[A-Z]') " +
+          "and (2) components are returned as JSX instead of being directly invoked."
+      );
+    }
+    setCurrent(LazyGuardDispatcher);
+  } else if (kind === GuardKind.PopGuardContext) {
+    // Pop before checking invariant or errors
+    const lastFrame = guardFrames.pop();
+
+    invariant(
+      lastFrame != null,
+      "React Forget internal error: unexpected null in guard stack"
+    );
+    if (guardFrames.length === 0) {
+      originalDispatcher = null;
+    }
+    setCurrent(lastFrame);
+  } else if (kind === GuardKind.PushExpectHook) {
+    // ExpectHooks could be nested, so we save the current dispatcher
+    // for the matching PopExpectHook to restore.
+    guardFrames.push(curr);
+    setCurrent(originalDispatcher);
+  } else if (kind === GuardKind.PopExpectHook) {
+    const lastFrame = guardFrames.pop();
+    invariant(
+      lastFrame != null,
+      "React Forget internal error: unexpected null in guard stack"
+    );
+    setCurrent(lastFrame);
+  } else {
+    invariant(false, "Forget internal error: unreachable block" + kind);
   }
-  ReactCurrentDispatcher.current = originalDispatcher;
-  originalDispatcher = null;
 }
 
 export function $reset($: MemoCache) {

compiler/packages/sprout/src/SproutTodoFilter.ts

@@ -515,6 +515,9 @@ const skipFilter = new Set([
   "bug-jsx-memberexpr-tag-in-lambda",
   "bug-invalid-code-when-bailout",
   "component-syntax-ref-gating.flow",
+  
+  // 'react-forget-runtime' not yet supported
+  "flag-enable-emit-hook-guards",
 ]);
 
 export default skipFilter;

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

@@ -229,4 +229,7 @@ export const ObjectWithHooks = {
   useMakeArray(): Array<number> {
     return [1, 2, 3];
   },
+  useIdentity<T>(arg: T): T {
+    return arg;
+  }
 };