[compiler] Override type provider for hook-like names

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

@@ -33,6 +33,7 @@ import {
   Type,
   ValidatedIdentifier,
   ValueKind,
+  getHookKindForType,
   makeBlockId,
   makeIdentifierId,
   makeIdentifierName,
@@ -788,14 +789,9 @@ export class Environment {
           );
         } else {
           const moduleType = this.#resolveModuleType(binding.module, loc);
+          let propertyType: Type | null = null;
           if (moduleType !== null) {
-            const importedType = this.getPropertyType(
-              moduleType,
-              binding.imported,
-            );
-            if (importedType != null) {
-              return importedType;
-            }
+            propertyType = this.getPropertyType(moduleType, binding.imported);
           }
 
           /**
@@ -806,9 +802,18 @@ export class Environment {
            * `import {useHook as foo} ...`
            * `import {foo as useHook} ...`
            */
-          return isHookName(binding.imported) || isHookName(binding.name)
-            ? this.#getCustomHookType()
-            : null;
+          const expectHook =
+            isHookName(binding.imported) || isHookName(binding.name);
+          if (expectHook) {
+            if (
+              propertyType &&
+              getHookKindForType(this, propertyType) !== null
+            ) {
+              return propertyType;
+            }
+            return this.#getCustomHookType();
+          }
+          return propertyType;
         }
       }
       case 'ImportDefault':
@@ -821,17 +826,27 @@ export class Environment {
           );
         } else {
           const moduleType = this.#resolveModuleType(binding.module, loc);
+          let importedType: Type | null = null;
           if (moduleType !== null) {
             if (binding.kind === 'ImportDefault') {
               const defaultType = this.getPropertyType(moduleType, 'default');
               if (defaultType !== null) {
-                return defaultType;
+                importedType = defaultType;
               }
             } else {
-              return moduleType;
+              importedType = moduleType;
+            }
+          }
+          if (isHookName(binding.name)) {
+            if (
+              importedType !== null &&
+              getHookKindForType(this, importedType) !== null
+            ) {
+              return importedType;
             }
+            return this.#getCustomHookType();
           }
-          return isHookName(binding.name) ? this.#getCustomHookType() : null;
+          return importedType;
         }
       }
     }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-hooklike-name-not-typed-as-hook-import.expect.md

@@ -0,0 +1,123 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {
+  useArrayConcatNotTypedAsHook,
+  ValidateMemoization,
+} from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => [a], [a]);
+  const item2 = useMemo(() => [b], [b]);
+  const item3 = useArrayConcatNotTypedAsHook(item1, item2);
+
+  return (
+    <>
+      <ValidateMemoization inputs={[a, b]} output={item3} />
+    </>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 1, b: 2},
+    {a: 2, b: 2},
+    {a: 3, b: 2},
+    {a: 0, b: 0},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+import {
+  useArrayConcatNotTypedAsHook,
+  ValidateMemoization,
+} from "shared-runtime";
+
+export function Component(t0) {
+  const $ = _c(10);
+  const { a, b } = t0;
+  let t1;
+  let t2;
+  if ($[0] !== a) {
+    t2 = [a];
+    $[0] = a;
+    $[1] = t2;
+  } else {
+    t2 = $[1];
+  }
+  t1 = t2;
+  const item1 = t1;
+  let t3;
+  let t4;
+  if ($[2] !== b) {
+    t4 = [b];
+    $[2] = b;
+    $[3] = t4;
+  } else {
+    t4 = $[3];
+  }
+  t3 = t4;
+  const item2 = t3;
+  const item3 = useArrayConcatNotTypedAsHook(item1, item2);
+  let t5;
+  if ($[4] !== a || $[5] !== b) {
+    t5 = [a, b];
+    $[4] = a;
+    $[5] = b;
+    $[6] = t5;
+  } else {
+    t5 = $[6];
+  }
+  let t6;
+  if ($[7] !== t5 || $[8] !== item3) {
+    t6 = (
+      <>
+        <ValidateMemoization inputs={t5} output={item3} />
+      </>
+    );
+    $[7] = t5;
+    $[8] = item3;
+    $[9] = t6;
+  } else {
+    t6 = $[9];
+  }
+  return t6;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 0, b: 0 }],
+  sequentialRenders: [
+    { a: 0, b: 0 },
+    { a: 1, b: 0 },
+    { a: 1, b: 1 },
+    { a: 1, b: 2 },
+    { a: 2, b: 2 },
+    { a: 3, b: 2 },
+    { a: 0, b: 0 },
+  ],
+};
+
+```
+
+### Eval output
+(kind: ok) <div>{"inputs":[0,0],"output":[0,0]}</div>
+<div>{"inputs":[1,0],"output":[1,0]}</div>
+<div>{"inputs":[1,1],"output":[1,1]}</div>
+<div>{"inputs":[1,2],"output":[1,2]}</div>
+<div>{"inputs":[2,2],"output":[2,2]}</div>
+<div>{"inputs":[3,2],"output":[3,2]}</div>
+<div>{"inputs":[0,0],"output":[0,0]}</div>

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-hooklike-name-not-typed-as-hook-import.tsx

@@ -0,0 +1,31 @@
+import {useMemo} from 'react';
+import {
+  useArrayConcatNotTypedAsHook,
+  ValidateMemoization,
+} from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => [a], [a]);
+  const item2 = useMemo(() => [b], [b]);
+  const item3 = useArrayConcatNotTypedAsHook(item1, item2);
+
+  return (
+    <>
+      <ValidateMemoization inputs={[a, b]} output={item3} />
+    </>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 1, b: 2},
+    {a: 2, b: 2},
+    {a: 3, b: 2},
+    {a: 0, b: 0},
+  ],
+};

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-import-object-typed-module-as-hook-name.expect.md

@@ -0,0 +1,145 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import useHook from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => ({a}), [a]);
+  const item2 = useMemo(() => ({b}), [b]);
+  useHook(item1, item2);
+
+  return (
+    <>
+      <ValidateMemoization inputs={[a]} output={item1} />
+      <ValidateMemoization inputs={[b]} output={item2} />
+    </>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 1, b: 2},
+    {a: 2, b: 2},
+    {a: 3, b: 2},
+    {a: 0, b: 0},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+import useHook from "shared-runtime";
+
+export function Component(t0) {
+  const $ = _c(17);
+  const { a, b } = t0;
+  let t1;
+  let t2;
+  if ($[0] !== a) {
+    t2 = { a };
+    $[0] = a;
+    $[1] = t2;
+  } else {
+    t2 = $[1];
+  }
+  t1 = t2;
+  const item1 = t1;
+  let t3;
+  let t4;
+  if ($[2] !== b) {
+    t4 = { b };
+    $[2] = b;
+    $[3] = t4;
+  } else {
+    t4 = $[3];
+  }
+  t3 = t4;
+  const item2 = t3;
+  useHook(item1, item2);
+  let t5;
+  if ($[4] !== a) {
+    t5 = [a];
+    $[4] = a;
+    $[5] = t5;
+  } else {
+    t5 = $[5];
+  }
+  let t6;
+  if ($[6] !== t5 || $[7] !== item1) {
+    t6 = <ValidateMemoization inputs={t5} output={item1} />;
+    $[6] = t5;
+    $[7] = item1;
+    $[8] = t6;
+  } else {
+    t6 = $[8];
+  }
+  let t7;
+  if ($[9] !== b) {
+    t7 = [b];
+    $[9] = b;
+    $[10] = t7;
+  } else {
+    t7 = $[10];
+  }
+  let t8;
+  if ($[11] !== t7 || $[12] !== item2) {
+    t8 = <ValidateMemoization inputs={t7} output={item2} />;
+    $[11] = t7;
+    $[12] = item2;
+    $[13] = t8;
+  } else {
+    t8 = $[13];
+  }
+  let t9;
+  if ($[14] !== t6 || $[15] !== t8) {
+    t9 = (
+      <>
+        {t6}
+        {t8}
+      </>
+    );
+    $[14] = t6;
+    $[15] = t8;
+    $[16] = t9;
+  } else {
+    t9 = $[16];
+  }
+  return t9;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 0, b: 0 }],
+  sequentialRenders: [
+    { a: 0, b: 0 },
+    { a: 1, b: 0 },
+    { a: 1, b: 1 },
+    { a: 1, b: 2 },
+    { a: 2, b: 2 },
+    { a: 3, b: 2 },
+    { a: 0, b: 0 },
+  ],
+};
+
+```
+
+### Eval output
+(kind: ok) [[ (exception in render) ReferenceError: ValidateMemoization is not defined ]]
+[[ (exception in render) ReferenceError: ValidateMemoization is not defined ]]
+[[ (exception in render) ReferenceError: ValidateMemoization is not defined ]]
+[[ (exception in render) ReferenceError: ValidateMemoization is not defined ]]
+[[ (exception in render) ReferenceError: ValidateMemoization is not defined ]]
+[[ (exception in render) ReferenceError: ValidateMemoization is not defined ]]
+[[ (exception in render) ReferenceError: ValidateMemoization is not defined ]]
+logs: [{ a: 0 },{ b: 0 },{ a: 0 },{ b: 0 },{ a: 1 },{ b: 1 },{ a: 1 },{ b: 2 },{ a: 2 },{ b: 2 },{ a: 3 },{ b: 2 },{ a: 0 },{ b: 0 }]