[compiler] Validate type configs for hooks/non-hooks

Alternative to #30868. The goal is to ensure that the types coming out of moduleTypeProvider are valid wrt to hook typing. If something is named like a hook, then it must be typed as a hook (or don't type it).

ghstack-source-id: 787575b
Pull Request resolved: #30888

Author: josephsavona

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

@@ -734,9 +734,11 @@ export class Environment {
         }
         const moduleConfig = parsedModuleConfig.data;
         moduleType = installTypeConfig(
+          moduleName,
           this.#globals,
           this.#shapes,
           moduleConfig,
+          moduleName,
         );
       } else {
         moduleType = null;

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

@@ -5,10 +5,11 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-import {Effect, ValueKind, ValueReason} from './HIR';
+import {Effect, GeneratedSource, ValueKind, ValueReason} from './HIR';
 import {
   BUILTIN_SHAPES,
   BuiltInArrayId,
+  BuiltInJsxId,
   BuiltInMixedReadonlyId,
   BuiltInUseActionStateId,
   BuiltInUseContextHookId,
@@ -28,6 +29,8 @@ import {
 import {BuiltInType, PolyType} from './Types';
 import {TypeConfig} from './TypeSchema';
 import {assertExhaustive} from '../Utils/utils';
+import {isHookName} from './Environment';
+import {CompilerError} from '..';
 
 /*
  * This file exports types and defaults for JavaScript global objects.
@@ -532,6 +535,50 @@ DEFAULT_GLOBALS.set(
 );
 
 export function installTypeConfig(
+  moduleName: string,
+  globals: GlobalRegistry,
+  shapes: ShapeRegistry,
+  typeConfig: TypeConfig,
+  moduleOrPropertyName: string | null,
+): Global {
+  const type = convertTypeConfig(moduleName, globals, shapes, typeConfig);
+  if (moduleOrPropertyName !== null) {
+    if (isHookName(moduleOrPropertyName)) {
+      // Named like a hook => must be typed as a hook
+      if (type.kind !== 'Function' || type.shapeId == null) {
+        CompilerError.throwInvalidConfig({
+          reason: `Invalid moduleTypeProvider result for module '${moduleName}', expected type for '${moduleOrPropertyName}' to be a hook based on its name, but the type was not a hook`,
+          loc: GeneratedSource,
+        });
+      }
+      const functionType = shapes.get(type.shapeId);
+      if (functionType == null || functionType.functionType?.hookKind == null) {
+        CompilerError.throwInvalidConfig({
+          reason: `Invalid moduleTypeProvider result for module '${moduleName}', expected type for '${moduleOrPropertyName}' to be a hook based on its name, but the type was not a hook`,
+          loc: GeneratedSource,
+        });
+      }
+    } else {
+      // Not named like a hook => must not be a hook
+      if (type.kind === 'Function' && type.shapeId != null) {
+        const functionType = shapes.get(type.shapeId);
+        if (
+          functionType != null &&
+          functionType.functionType?.hookKind != null
+        ) {
+          CompilerError.throwInvalidConfig({
+            reason: `Invalid moduleTypeProvider result for module '${moduleName}', expected type for '${moduleOrPropertyName}' not to be a hook, but it was typed as a hook`,
+            loc: GeneratedSource,
+          });
+        }
+      }
+    }
+  }
+  return type;
+}
+
+function convertTypeConfig(
+  moduleName: string,
   globals: GlobalRegistry,
   shapes: ShapeRegistry,
   typeConfig: TypeConfig,
@@ -554,6 +601,9 @@ export function installTypeConfig(
         case 'Any': {
           return {kind: 'Poly'};
         }
+        case 'JSX': {
+          return {kind: 'Object', shapeId: BuiltInJsxId};
+        }
         default: {
           assertExhaustive(
             typeConfig.name,
@@ -567,7 +617,12 @@ export function installTypeConfig(
         positionalParams: typeConfig.positionalParams,
         restParam: typeConfig.restParam,
         calleeEffect: typeConfig.calleeEffect,
-        returnType: installTypeConfig(globals, shapes, typeConfig.returnType),
+        returnType: convertTypeConfig(
+          moduleName,
+          globals,
+          shapes,
+          typeConfig.returnType,
+        ),
         returnValueKind: typeConfig.returnValueKind,
         noAlias: typeConfig.noAlias === true,
         mutableOnlyIfOperandsAreMutable:
@@ -580,7 +635,12 @@ export function installTypeConfig(
         positionalParams: typeConfig.positionalParams ?? [],
         restParam: typeConfig.restParam ?? Effect.Freeze,
         calleeEffect: Effect.Read,
-        returnType: installTypeConfig(globals, shapes, typeConfig.returnType),
+        returnType: convertTypeConfig(
+          moduleName,
+          globals,
+          shapes,
+          typeConfig.returnType,
+        ),
         returnValueKind: typeConfig.returnValueKind ?? ValueKind.Frozen,
         noAlias: typeConfig.noAlias === true,
       });
@@ -591,7 +651,7 @@ export function installTypeConfig(
         null,
         Object.entries(typeConfig.properties ?? {}).map(([key, value]) => [
           key,
-          installTypeConfig(globals, shapes, value),
+          installTypeConfig(moduleName, globals, shapes, value, key),
         ]),
       );
     }

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

@@ -74,13 +74,15 @@ export type BuiltInTypeConfig =
   | 'Ref'
   | 'Array'
   | 'Primitive'
-  | 'MixedReadonly';
+  | 'MixedReadonly'
+  | 'JSX';
 export const BuiltInTypeSchema: z.ZodType<BuiltInTypeConfig> = z.union([
   z.literal('Any'),
   z.literal('Ref'),
   z.literal('Array'),
   z.literal('Primitive'),
   z.literal('MixedReadonly'),
+  z.literal('JSX'),
 ]);
 
 export type TypeReferenceConfig = {