compiler: only resolve globals and react imports

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

@@ -11,7 +11,6 @@ import { fromZodError } from "zod-validation-error";
 import { CompilerError } from "../CompilerError";
 import { Logger } from "../Entrypoint";
 import { Err, Ok, Result } from "../Utils/Result";
-import { log } from "../Utils/logger";
 import {
   DEFAULT_GLOBALS,
   DEFAULT_SHAPES,
@@ -320,6 +319,8 @@ const EnvironmentConfigSchema = z.object({
    */
   throwUnknownException__testonly: z.boolean().default(false),
 
+  enableSharedRuntime__testonly: z.boolean().default(false),
+
   /**
    * Enables deps of a function epxression to be treated as conditional. This
    * makes sure we don't load a dep when it's a property (to check if it has
@@ -514,7 +515,6 @@ export class Environment {
 
   getGlobalDeclaration(binding: NonLocalBinding): Global | null {
     const name = binding.name;
-    let resolvedName = name;
 
     if (this.config.hookPattern != null) {
       const match = new RegExp(this.config.hookPattern).exec(name);
@@ -523,21 +523,45 @@ export class Environment {
         typeof match[1] === "string" &&
         isHookName(match[1])
       ) {
-        resolvedName = match[1];
+        const resolvedName = match[1];
+        return this.#globals.get(resolvedName) ?? this.#getCustomHookType();
       }
     }
 
-    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(resolvedName)) {
-        return this.#getCustomHookType();
-      } else {
-        log(() => `Undefined global \`${name}\``);
+    let global: Global | null = null;
+    switch (binding.kind) {
+      case "ModuleLocal": {
+        // don't resolve module locals
+        break;
+      }
+      case "Global": {
+        global = this.#globals.get(name) ?? null;
+        break;
+      }
+      case "ImportDefault":
+      case "ImportNamespace":
+      case "ImportSpecifier": {
+        if (
+          binding.module.toLowerCase() === "react" ||
+          binding.module.toLowerCase() === "react-dom" ||
+          (this.config.enableSharedRuntime__testonly &&
+            binding.module === "shared-runtime")
+        ) {
+          // only resolve imports to modules we know about
+          global = this.#globals.get(name) ?? null;
+        }
+        break;
       }
     }
 
-    return resolvedGlobal;
+    if (global === null && isHookName(name)) {
+      /**
+       * Type inference relies on all hooks being resolved as such, so if we don't have
+       * a global declaration and its a hook name, return the default custom hook type.
+       */
+      return this.#getCustomHookType();
+    }
+    return global;
   }
 
   getPropertyType(

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

@@ -588,7 +588,38 @@ export function printInstructionValue(instrValue: ReactiveValue): string {
       break;
     }
     case "LoadGlobal": {
-      value = `LoadGlobal ${instrValue.binding.name}`;
+      switch (instrValue.binding.kind) {
+        case "Global": {
+          value = `LoadGlobal(global) ${instrValue.binding.name}`;
+          break;
+        }
+        case "ModuleLocal": {
+          value = `LoadGlobal(module) ${instrValue.binding.name}`;
+          break;
+        }
+        case "ImportDefault": {
+          value = `LoadGlobal import ${instrValue.binding.name} from '${instrValue.binding.module}'`;
+          break;
+        }
+        case "ImportNamespace": {
+          value = `LoadGlobal import * as ${instrValue.binding.name} from '${instrValue.binding.module}'`;
+          break;
+        }
+        case "ImportSpecifier": {
+          if (instrValue.binding.imported !== instrValue.binding.name) {
+            value = `LoadGlobal import { ${instrValue.binding.imported} as ${instrValue.binding.name} } from '${instrValue.binding.module}'`;
+          } else {
+            value = `LoadGlobal import { ${instrValue.binding.name} } from '${instrValue.binding.module}'`;
+          }
+          break;
+        }
+        default: {
+          assertExhaustive(
+            instrValue.binding,
+            `Unexpected binding kind \`${(instrValue.binding as any).kind}\``
+          );
+        }
+      }
       break;
     }
     case "StoreGlobal": {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/globals-dont-resolve-local-useState.expect.md

@@ -0,0 +1,64 @@
+
+## Input
+
+```javascript
+import { useState as _useState } from "react";
+
+function useState(value) {
+  return _useState(value);
+}
+
+function Component() {
+  const [state, setState] = useState("hello");
+
+  return <div onClick={() => setState("bye")}>{state}</div>;
+}
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useState as _useState } from "react";
+
+function useState(value) {
+  const $ = _c(2);
+  let t0;
+  if ($[0] !== value) {
+    t0 = _useState(value);
+    $[0] = value;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  return t0;
+}
+
+function Component() {
+  const $ = _c(5);
+  const [state, setState] = useState("hello");
+  let t0;
+  if ($[0] !== setState) {
+    t0 = () => setState("bye");
+    $[0] = setState;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
+  let t1;
+  if ($[2] !== t0 || $[3] !== state) {
+    t1 = <div onClick={t0}>{state}</div>;
+    $[2] = t0;
+    $[3] = state;
+    $[4] = t1;
+  } else {
+    t1 = $[4];
+  }
+  return t1;
+}
+
+```
+
+### Eval output
+(kind: exception) Fixture not implemented

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/globals-dont-resolve-local-useState.js

@@ -0,0 +1,11 @@
+import { useState as _useState } from "react";
+
+function useState(value) {
+  return _useState(value);
+}
+
+function Component() {
+  const [state, setState] = useState("hello");
+
+  return <div onClick={() => setState("bye")}>{state}</div>;
+}

compiler/packages/snap/src/compiler.ts

@@ -170,6 +170,7 @@ function makePluginOptions(
       enableEmitInstrumentForget,
       enableEmitHookGuards,
       assertValidMutableRanges: true,
+      enableSharedRuntime__testonly: true,
       hookPattern,
       validatePreserveExistingMemoizationGuarantees,
     },