compiler: only resolve globals and react imports

Updates Environment#getGlobalDeclaration() to only resolve "globals" if they are a true global or an import from react/react-dom. We still keep the logic to resolve hook-like names as custom hooks. Notably, this means that a local `Array` reference won't get confused with our Array global declaration, a local `useState` (or import from something other than React) won't get confused as `React.useState()`, etc.

I tried to write a proper fixture test to test that we react to changes to a custom setState setter function, but I think there may be an issue with snap and how it handles re-renders from effects. I think the tests are good here but open to feedback if we want to go down the rabbit hole of figuring out a proper snap test for this.

ghstack-source-id: 5e9a8f6e0d23659c72a9d041e8d394b83d6e526d
Pull Request resolved: #29190

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
@@ -513,31 +514,78 @@ 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);
+      const match = new RegExp(this.config.hookPattern).exec(binding.name);
       if (
         match != null &&
         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}\``);
+    switch (binding.kind) {
+      case "ModuleLocal": {
+        // don't resolve module locals
+        return isHookName(binding.name) ? this.#getCustomHookType() : null;
+      }
+      case "Global": {
+        return (
+          this.#globals.get(binding.name) ??
+          (isHookName(binding.name) ? this.#getCustomHookType() : null)
+        );
+      }
+      case "ImportSpecifier": {
+        if (this.#isKnownReactModule(binding.module)) {
+          /**
+           * For `import {imported as name} from "..."` form, we use the `imported`
+           * name rather than the local alias. Because we don't have definitions for
+           * every React builtin hook yet, we also check to see if the imported name
+           * is hook-like (whereas the fall-through below is checking if the aliased
+           * name is hook-like)
+           */
+          return (
+            this.#globals.get(binding.imported) ??
+            (isHookName(binding.imported) ? this.#getCustomHookType() : null)
+          );
+        } else {
+          /**
+           * For modules we don't own, we look at whether the original name or import alias
+           * are hook-like. Both of the following are likely hooks so we would return a hook
+           * type for both:
+           *
+           * `import {useHook as foo} ...`
+           * `import {foo as useHook} ...`
+           */
+          return isHookName(binding.imported) || isHookName(binding.name)
+            ? this.#getCustomHookType()
+            : null;
+        }
+      }
+      case "ImportDefault":
+      case "ImportNamespace": {
+        if (this.#isKnownReactModule(binding.module)) {
+          // only resolve imports to modules we know about
+          return (
+            this.#globals.get(binding.name) ??
+            (isHookName(binding.name) ? this.#getCustomHookType() : null)
+          );
+        } else {
+          return isHookName(binding.name) ? this.#getCustomHookType() : null;
+        }
       }
     }
+  }
 
-    return resolvedGlobal;
+  #isKnownReactModule(moduleName: string): boolean {
+    return (
+      moduleName.toLowerCase() === "react" ||
+      moduleName.toLowerCase() === "react-dom" ||
+      (this.config.enableSharedRuntime__testonly &&
+        moduleName === "shared-runtime")
+    );
   }
 
   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/destructuring-mixed-scope-and-local-variables-with-default.expect.md

@@ -56,71 +56,78 @@ function useFragment(_arg1, _arg2) {
 }
 
 function Component(props) {
-  const $ = _c(14);
-  const post = useFragment(graphql`...`, props.post);
+  const $ = _c(15);
+  let t0;
+  if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
+    t0 = graphql`...`;
+    $[0] = t0;
+  } else {
+    t0 = $[0];
+  }
+  const post = useFragment(t0, props.post);
   let media;
   let allUrls;
   let onClick;
-  if ($[0] !== post) {
+  if ($[1] !== post) {
     allUrls = [];
 
-    const { media: t0, comments: t1, urls: t2 } = post;
-    media = t0 === undefined ? null : t0;
-    let t3;
-    if ($[4] !== t1) {
-      t3 = t1 === undefined ? [] : t1;
-      $[4] = t1;
-      $[5] = t3;
-    } else {
-      t3 = $[5];
-    }
-    const comments = t3;
+    const { media: t1, comments: t2, urls: t3 } = post;
+    media = t1 === undefined ? null : t1;
     let t4;
-    if ($[6] !== t2) {
+    if ($[5] !== t2) {
       t4 = t2 === undefined ? [] : t2;
-      $[6] = t2;
-      $[7] = t4;
+      $[5] = t2;
+      $[6] = t4;
     } else {
-      t4 = $[7];
+      t4 = $[6];
     }
-    const urls = t4;
+    const comments = t4;
     let t5;
-    if ($[8] !== comments.length) {
-      t5 = (e) => {
+    if ($[7] !== t3) {
+      t5 = t3 === undefined ? [] : t3;
+      $[7] = t3;
+      $[8] = t5;
+    } else {
+      t5 = $[8];
+    }
+    const urls = t5;
+    let t6;
+    if ($[9] !== comments.length) {
+      t6 = (e) => {
         if (!comments.length) {
           return;
         }
 
         console.log(comments.length);
       };
-      $[8] = comments.length;
-      $[9] = t5;
+      $[9] = comments.length;
+      $[10] = t6;
     } else {
-      t5 = $[9];
+      t6 = $[10];
     }
-    onClick = t5;
+    onClick = t6;
 
     allUrls.push(...urls);
-    $[0] = post;
-    $[1] = media;
-    $[2] = allUrls;
-    $[3] = onClick;
+    $[1] = post;
+    $[2] = media;
+    $[3] = allUrls;
+    $[4] = onClick;
   } else {
-    media = $[1];
-    allUrls = $[2];
-    onClick = $[3];
+    media = $[2];
+    allUrls = $[3];
+    onClick = $[4];
   }
-  let t0;
-  if ($[10] !== media || $[11] !== allUrls || $[12] !== onClick) {
-    t0 = <Stringify media={media} allUrls={allUrls} onClick={onClick} />;
-    $[10] = media;
-    $[11] = allUrls;
-    $[12] = onClick;
-    $[13] = t0;
+  let t1;
+  if ($[11] !== media || $[12] !== allUrls || $[13] !== onClick) {
+    t1 = <Stringify media={media} allUrls={allUrls} onClick={onClick} />;
+    $[11] = media;
+    $[12] = allUrls;
+    $[13] = onClick;
+    $[14] = t1;
   } else {
-    t0 = $[13];
+    t1 = $[14];
   }
-  return t0;
+  return t1;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-call-aliased-hook-import.expect.md

@@ -0,0 +1,30 @@
+
+## Input
+
+```javascript
+import { useFragment as readFragment } from "shared-runtime";
+
+function Component(props) {
+  let data;
+  if (props.cond) {
+    data = readFragment();
+  }
+  return data;
+}
+
+```
+
+
+## Error
+
+```
+  4 |   let data;
+  5 |   if (props.cond) {
+> 6 |     data = readFragment();
+    |            ^^^^^^^^^^^^ InvalidReact: Hooks must always be called in a consistent order, and may not be called conditionally. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning) (6:6)
+  7 |   }
+  8 |   return data;
+  9 | }
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-call-aliased-hook-import.js

@@ -0,0 +1,9 @@
+import { useFragment as readFragment } from "shared-runtime";
+
+function Component(props) {
+  let data;
+  if (props.cond) {
+    data = readFragment();
+  }
+  return data;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-call-aliased-react-hook.expect.md

@@ -0,0 +1,30 @@
+
+## Input
+
+```javascript
+import { useState as state } from "react";
+
+function Component(props) {
+  let s;
+  if (props.cond) {
+    [s] = state();
+  }
+  return s;
+}
+
+```
+
+
+## Error
+
+```
+  4 |   let s;
+  5 |   if (props.cond) {
+> 6 |     [s] = state();
+    |           ^^^^^ InvalidReact: Hooks must always be called in a consistent order, and may not be called conditionally. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning) (6:6)
+  7 |   }
+  8 |   return s;
+  9 | }
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-call-aliased-react-hook.js

@@ -0,0 +1,9 @@
+import { useState as state } from "react";
+
+function Component(props) {
+  let s;
+  if (props.cond) {
+    [s] = state();
+  }
+  return s;
+}

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-call-non-hook-imported-as-hook.expect.md

@@ -0,0 +1,30 @@
+
+## Input
+
+```javascript
+import { makeArray as useArray } from "other";
+
+function Component(props) {
+  let data;
+  if (props.cond) {
+    data = useArray();
+  }
+  return data;
+}
+
+```
+
+
+## Error
+
+```
+  4 |   let data;
+  5 |   if (props.cond) {
+> 6 |     data = useArray();
+    |            ^^^^^^^^ InvalidReact: Hooks must always be called in a consistent order, and may not be called conditionally. See the Rules of Hooks (https://react.dev/warnings/invalid-hook-call-warning) (6:6)
+  7 |   }
+  8 |   return data;
+  9 | }
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.invalid-conditional-call-non-hook-imported-as-hook.js

@@ -0,0 +1,9 @@
+import { makeArray as useArray } from "other";
+
+function Component(props) {
+  let data;
+  if (props.cond) {
+    data = useArray();
+  }
+  return data;
+}