Allow prefixed hooks for compiling bundled code

We're doing some internal benchmarking using a lightweight bundler that @pieterv 
wrote for experimentation purposes. It's designed to fully preserve Flow type 
annotations so we can experiment with type-driven compilation and test out what 
benefits we might get from "cross-module" compilation more easily (ie by just 
bundling together a few modules so we can see them all as one). 

However, the bundler renames local variables and imports, so that a reference to 
`useMemo()` might end up as `React$useMemo()` or similar. This PR adds a flag to 
tell the compiler that builtin hooks might be prefixed and resolve them 
appropriately.

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

@@ -376,11 +376,12 @@ function shouldVisitNode(fn: BabelFn, pass: CompilerPass): boolean {
       return false;
     }
     case "infer": {
+      const hookPattern = pass.opts.environment?.hookPattern ?? null;
       return (
         // Component declarations are known components
         (fn.isFunctionDeclaration() && isComponentDeclaration(fn.node)) ||
         // Otherwise check if this is a component or hook-like function
-        isComponentOrHookLike(fn)
+        isComponentOrHookLike(fn, hookPattern)
       );
     }
     case "all": {
@@ -414,7 +415,10 @@ function hasUseMemoCacheCall(
   return hasUseMemoCache;
 }
 
-function isHookName(s: string): boolean {
+function isHookName(s: string, hookPattern: string | null): boolean {
+  if (hookPattern !== null) {
+    return new RegExp(hookPattern).test(s);
+  }
   return /^use[A-Z0-9]/.test(s);
 }
 
@@ -423,13 +427,16 @@ function isHookName(s: string): boolean {
  * containing a hook name.
  */
 
-function isHook(path: NodePath<t.Expression | t.PrivateName>): boolean {
+function isHook(
+  path: NodePath<t.Expression | t.PrivateName>,
+  hookPattern: string | null
+): boolean {
   if (path.isIdentifier()) {
-    return isHookName(path.node.name);
+    return isHookName(path.node.name, hookPattern);
   } else if (
     path.isMemberExpression() &&
     !path.node.computed &&
-    isHook(path.get("property"))
+    isHook(path.get("property"), hookPattern)
   ) {
     const obj = path.get("object").node;
     const isPascalCaseNameSpace = /^[A-Z].*/;
@@ -518,18 +525,20 @@ function isValidComponentParams(
 function isComponentOrHookLike(
   node: NodePath<
     t.FunctionDeclaration | t.ArrowFunctionExpression | t.FunctionExpression
-  >
+  >,
+  hookPattern: string | null
 ): boolean {
   const functionName = getFunctionName(node);
   // Check if the name is component or hook like:
   if (functionName !== null && isComponentName(functionName)) {
     return (
       // As an added check we also look for hook invocations or JSX
-      callsHooksOrCreatesJsx(node) && isValidComponentParams(node.get("params"))
+      callsHooksOrCreatesJsx(node, hookPattern) &&
+      isValidComponentParams(node.get("params"))
     );
-  } else if (functionName !== null && isHook(functionName)) {
+  } else if (functionName !== null && isHook(functionName, hookPattern)) {
     // Hooks have hook invocations or JSX, but can take any # of arguments
-    return callsHooksOrCreatesJsx(node);
+    return callsHooksOrCreatesJsx(node, hookPattern);
   }
 
   /*
@@ -539,7 +548,7 @@ function isComponentOrHookLike(
   if (node.isFunctionExpression() || node.isArrowFunctionExpression()) {
     if (isForwardRefCallback(node) || isMemoCallback(node)) {
       // As an added check we also look for hook invocations or JSX
-      return callsHooksOrCreatesJsx(node);
+      return callsHooksOrCreatesJsx(node, hookPattern);
     } else {
       return false;
     }
@@ -547,7 +556,10 @@ function isComponentOrHookLike(
   return false;
 }
 
-function callsHooksOrCreatesJsx(node: NodePath<t.Node>): boolean {
+function callsHooksOrCreatesJsx(
+  node: NodePath<t.Node>,
+  hookPattern: string | null
+): boolean {
   let invokesHooks = false;
   let createsJsx = false;
   node.traverse({
@@ -556,7 +568,7 @@ function callsHooksOrCreatesJsx(node: NodePath<t.Node>): boolean {
     },
     CallExpression(call) {
       const callee = call.get("callee");
-      if (callee.isExpression() && isHook(callee)) {
+      if (callee.isExpression() && isHook(callee, hookPattern)) {
         invokesHooks = true;
       }
     },

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

@@ -344,6 +344,19 @@ const EnvironmentConfigSchema = z.object({
    * non-ideal.
    */
   enableTreatFunctionDepsAsConditional: z.boolean().default(false),
+
+  /**
+   * If specified, this value is used as a pattern for determing which global values should be
+   * treated as hooks. The pattern should have a single capture group, which will be used as
+   * the hook name for the purposes of resolving hook definitions (for builtin hooks)_.
+   *
+   * For example, by default `React$useState` would not be treated as a hook. By specifying
+   * `hookPattern: 'React$(\w+)'`, the compiler will treat this value equivalently to `useState()`.
+   *
+   * This setting is intended for cases where Forget is compiling code that has been prebundled
+   * and identifiers have been changed.
+   */
+  hookPattern: z.string().nullable().default(null),
 });
 
 export type EnvironmentConfig = z.infer<typeof EnvironmentConfigSchema>;
@@ -447,7 +460,20 @@ export class Environment {
   }
 
   getGlobalDeclaration(name: string): Global | null {
-    let resolvedGlobal: Global | null = this.#globals.get(name) ?? null;
+    let resolvedName = name;
+
+    if (this.config.hookPattern != null) {
+      const match = new RegExp(this.config.hookPattern).exec(name);
+      if (
+        match != null &&
+        typeof match[1] === "string" &&
+        isHookName(match[1])
+      ) {
+        resolvedName = match[1];
+      }
+    }
+
+    let resolvedGlobal: Global | null = this.#globals.get(resolvedName) ?? null;
     if (resolvedGlobal === null) {
       // Hack, since we don't track module level declarations and imports
       if (isHookName(name)) {

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

@@ -1233,9 +1233,15 @@ export function isUseInsertionEffectHookType(id: Identifier): boolean {
 }
 
 export function getHookKind(env: Environment, id: Identifier): HookKind | null {
-  const idType = id.type;
-  if (idType.kind === "Function") {
-    const signature = env.getFunctionSignature(idType);
+  return getHookKindForType(env, id.type);
+}
+
+export function getHookKindForType(
+  env: Environment,
+  type: Type
+): HookKind | null {
+  if (type.kind === "Function") {
+    const signature = env.getFunctionSignature(type);
     return signature?.hookKind ?? null;
   }
   return null;

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

@@ -14,6 +14,7 @@ import {
   Instruction,
   Place,
   SpreadPattern,
+  getHookKindForType,
   makeInstructionId,
 } from "../HIR";
 import { createTemporaryPlace, markInstructionIds } from "../HIR/HIRBuilder";
@@ -43,11 +44,11 @@ export function dropManualMemoization(func: HIRFunction): void {
           break;
         }
         case "LoadGlobal": {
-          if (
-            instr.value.name === "useMemo" ||
-            instr.value.name === "useCallback"
-          ) {
-            hooks.set(instr.lvalue.identifier.id, instr.value.name);
+          const global = func.env.getGlobalDeclaration(instr.value.name);
+          const hookKind =
+            global !== null ? getHookKindForType(func.env, global) : null;
+          if (hookKind === "useMemo" || hookKind === "useCallback") {
+            hooks.set(instr.lvalue.identifier.id, hookKind);
           } else if (instr.value.name === "React") {
             react.add(instr.lvalue.identifier.id);
           }

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.expect.md

@@ -0,0 +1,75 @@
+
+## Input
+
+```javascript
+// @hookPattern:"React\$(\w+)"
+
+import * as React from "react";
+import { makeArray } from "shared-runtime";
+
+const React$useState = React.useState;
+const React$useMemo = React.useMemo;
+
+function Component() {
+  const [state, setState] = React$useState(0);
+  const doubledArray = React$useMemo(() => {
+    return makeArray(state);
+  }, [state]);
+  return <div>{doubledArray.join("")}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+
+## Code
+
+```javascript
+import { unstable_useMemoCache as useMemoCache } from "react"; // @hookPattern:"React\$(\w+)"
+
+import * as React from "react";
+import { makeArray } from "shared-runtime";
+
+const React$useState = React.useState;
+const React$useMemo = React.useMemo;
+
+function Component() {
+  const $ = useMemoCache(5);
+  const [state] = React$useState(0);
+  let t15;
+  let t0;
+  if ($[0] !== state) {
+    t15 = makeArray(state);
+    const doubledArray = t15;
+
+    t0 = doubledArray.join("");
+    $[0] = state;
+    $[1] = t0;
+    $[2] = t15;
+  } else {
+    t0 = $[1];
+    t15 = $[2];
+  }
+  let t1;
+  if ($[3] !== t0) {
+    t1 = <div>{t0}</div>;
+    $[3] = t0;
+    $[4] = t1;
+  } else {
+    t1 = $[4];
+  }
+  return t1;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>0</div>

compiler/packages/babel-plugin-react-forget/src/__tests__/fixtures/compiler/hooks-with-prefix.js

@@ -0,0 +1,20 @@
+// @hookPattern:"React\$(\w+)"
+
+import * as React from "react";
+import { makeArray } from "shared-runtime";
+
+const React$useState = React.useState;
+const React$useMemo = React.useMemo;
+
+function Component() {
+  const [state, setState] = React$useState(0);
+  const doubledArray = React$useMemo(() => {
+    return makeArray(state);
+  }, [state]);
+  return <div>{doubledArray.join("")}</div>;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{}],
+};

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

@@ -37,6 +37,7 @@ export function transformFixtureInput(
   let compilationMode: CompilationMode = "all";
   let enableUseMemoCachePolyfill = false;
   let panicThreshold: PanicThresholdOptions = "ALL_ERRORS";
+  let hookPattern: string | null = null;
 
   if (firstLine.indexOf("@compilationMode(annotation)") !== -1) {
     assert(
@@ -85,9 +86,24 @@ export function transformFixtureInput(
   }
 
   let eslintSuppressionRules: Array<string> | null = null;
-  const match = /@eslintSuppressionRules\(([^)]+)\)/.exec(firstLine);
-  if (match != null) {
-    eslintSuppressionRules = match[1].split("|");
+  const eslintSuppressionMatch = /@eslintSuppressionRules\(([^)]+)\)/.exec(
+    firstLine
+  );
+  if (eslintSuppressionMatch != null) {
+    eslintSuppressionRules = eslintSuppressionMatch[1].split("|");
+  }
+
+  const hookPatternMatch = /@hookPattern:"([^"]+)"/.exec(firstLine);
+  if (
+    hookPatternMatch &&
+    hookPatternMatch.length > 1 &&
+    hookPatternMatch[1].trim().length > 0
+  ) {
+    hookPattern = hookPatternMatch[1].trim();
+  } else if (firstLine.includes("@hookPattern")) {
+    throw new Error(
+      'Invalid @hookPattern:"..." pragma, must contain the prefix between balanced double quotes eg @hookPattern:"pattern"'
+    );
   }
 
   const config = parseConfigPragmaFn(firstLine);
@@ -131,6 +147,7 @@ export function transformFixtureInput(
         enableEmitInstrumentForget,
         enableEmitHookGuards,
         assertValidMutableRanges: true,
+        hookPattern,
       },
       compilationMode,
       logger: null,