[compiler][wip] Environment option for resolving imported module types

Adds a new Environment config option which allows specifying a function that is called to resolve types of imported modules. The function is passed the name of the imported module (the RHS of the import stmt) and can return a TypeConfig, which is a recursive type of the following form:

* Object of valid identifier keys (or "*" for wildcard) and values that are TypeConfigs
* Function with various properties, whose return type is a TypeConfig
* or a reference to a builtin type using one of a small list (currently Ref, Array, MixedReadonly, Primitive)

Rather than have to eagerly supply all known types (most of which may not be used) when creating the config, this function can do so lazily. During InferTypes we call `getGlobalDeclaration()` to resolve global types. Originally this was just for known react modules, but if the new config option is passed we also call it to see if it can resolve a type. For `import {name} from 'module'` syntax, we first resolve the module type and then call `getPropertyType(moduleType, 'name')` to attempt to retrieve the property of the module (the module would obviously have to be typed as an object type for this to have a chance of yielding a result). If the module type is returned as null, or the property doesn't exist, we fall through to the original checking of whether the name was hook-like.

TODO:
* testing
* cache the results of modules so we don't have to re-parse/install their types on each LoadGlobal of the same module
* decide what to do if the module types are invalid. probably better to fatal rather than bail out, since this would indicate an invalid configuration.

ghstack-source-id: bfdbf67e3dd0cbfd511bed0bd6ba92266cf99ab8
Pull Request resolved: #30771

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

@@ -17,6 +17,7 @@ import {
   Global,
   GlobalRegistry,
   installReAnimatedTypes,
+  installTypeConfig,
 } from './Globals';
 import {
   BlockId,
@@ -28,6 +29,7 @@ import {
   NonLocalBinding,
   PolyType,
   ScopeId,
+  SourceLocation,
   Type,
   ValidatedIdentifier,
   ValueKind,
@@ -45,6 +47,7 @@ import {
   addHook,
 } from './ObjectShape';
 import {Scope as BabelScope} from '@babel/traverse';
+import {TypeSchema} from './TypeSchema';
 
 export const ExternalFunctionSchema = z.object({
   // Source for the imported module that exports the `importSpecifierName` functions
@@ -137,6 +140,12 @@ export type Hook = z.infer<typeof HookSchema>;
 const EnvironmentConfigSchema = z.object({
   customHooks: z.map(z.string(), HookSchema).optional().default(new Map()),
 
+  /**
+   * A function that, given the name of a module, can optionally return a description
+   * of that module's type signature.
+   */
+  moduleTypeProvider: z.nullable(z.function().args(z.string())).default(null),
+
   /**
    * A list of functions which the application compiles as macros, where
    * the compiler must ensure they are not compiled to rename the macro or separate the
@@ -577,6 +586,7 @@ export function printFunctionType(type: ReactFunctionType): string {
 export class Environment {
   #globals: GlobalRegistry;
   #shapes: ShapeRegistry;
+  #moduleTypes: Map<string, Global | null> = new Map();
   #nextIdentifer: number = 0;
   #nextBlock: number = 0;
   #nextScope: number = 0;
@@ -698,7 +708,40 @@ export class Environment {
     return this.#outlinedFunctions;
   }
 
-  getGlobalDeclaration(binding: NonLocalBinding): Global | null {
+  #resolveModuleType(moduleName: string, loc: SourceLocation): Global | null {
+    if (this.config.moduleTypeProvider == null) {
+      return null;
+    }
+    let moduleType = this.#moduleTypes.get(moduleName);
+    if (moduleType === undefined) {
+      const unparsedModuleConfig = this.config.moduleTypeProvider(moduleName);
+      if (unparsedModuleConfig != null) {
+        const parsedModuleConfig = TypeSchema.safeParse(unparsedModuleConfig);
+        if (!parsedModuleConfig.success) {
+          CompilerError.throwInvalidConfig({
+            reason: `Could not parse module type, the configured \`moduleTypeProvider\` function returned an invalid module description`,
+            description: parsedModuleConfig.error.toString(),
+            loc,
+          });
+        }
+        const moduleConfig = parsedModuleConfig.data;
+        moduleType = installTypeConfig(
+          this.#globals,
+          this.#shapes,
+          moduleConfig,
+        );
+      } else {
+        moduleType = null;
+      }
+      this.#moduleTypes.set(moduleName, moduleType);
+    }
+    return moduleType;
+  }
+
+  getGlobalDeclaration(
+    binding: NonLocalBinding,
+    loc: SourceLocation,
+  ): Global | null {
     if (this.config.hookPattern != null) {
       const match = new RegExp(this.config.hookPattern).exec(binding.name);
       if (
@@ -736,6 +779,17 @@ export class Environment {
             (isHookName(binding.imported) ? this.#getCustomHookType() : null)
           );
         } else {
+          const moduleType = this.#resolveModuleType(binding.module, loc);
+          if (moduleType !== null) {
+            const importedType = this.getPropertyType(
+              moduleType,
+              binding.imported,
+            );
+            if (importedType != null) {
+              return importedType;
+            }
+          }
+
           /**
            * 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
@@ -758,6 +812,17 @@ export class Environment {
             (isHookName(binding.name) ? this.#getCustomHookType() : null)
           );
         } else {
+          const moduleType = this.#resolveModuleType(binding.module, loc);
+          if (moduleType !== null) {
+            if (binding.kind === 'ImportDefault') {
+              const defaultType = this.getPropertyType(moduleType, 'default');
+              if (defaultType !== null) {
+                return defaultType;
+              }
+            } else {
+              return moduleType;
+            }
+          }
           return isHookName(binding.name) ? this.#getCustomHookType() : null;
         }
       }
@@ -767,9 +832,7 @@ export class Environment {
   #isKnownReactModule(moduleName: string): boolean {
     return (
       moduleName.toLowerCase() === 'react' ||
-      moduleName.toLowerCase() === 'react-dom' ||
-      (this.config.enableSharedRuntime__testonly &&
-        moduleName === 'shared-runtime')
+      moduleName.toLowerCase() === 'react-dom'
     );
   }
 

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

@@ -9,6 +9,7 @@ import {Effect, ValueKind, ValueReason} from './HIR';
 import {
   BUILTIN_SHAPES,
   BuiltInArrayId,
+  BuiltInMixedReadonlyId,
   BuiltInUseActionStateId,
   BuiltInUseContextHookId,
   BuiltInUseEffectHookId,
@@ -25,6 +26,8 @@ import {
   addObject,
 } from './ObjectShape';
 import {BuiltInType, PolyType} from './Types';
+import {TypeConfig} from './TypeSchema';
+import {assertExhaustive} from '../Utils/utils';
 
 /*
  * This file exports types and defaults for JavaScript global objects.
@@ -528,6 +531,79 @@ DEFAULT_GLOBALS.set(
   addObject(DEFAULT_SHAPES, 'global', TYPED_GLOBALS),
 );
 
+export function installTypeConfig(
+  globals: GlobalRegistry,
+  shapes: ShapeRegistry,
+  typeConfig: TypeConfig,
+): Global {
+  switch (typeConfig.kind) {
+    case 'type': {
+      switch (typeConfig.name) {
+        case 'Array': {
+          return {kind: 'Object', shapeId: BuiltInArrayId};
+        }
+        case 'MixedReadonly': {
+          return {kind: 'Object', shapeId: BuiltInMixedReadonlyId};
+        }
+        case 'Primitive': {
+          return {kind: 'Primitive'};
+        }
+        case 'Ref': {
+          return {kind: 'Object', shapeId: BuiltInUseRefId};
+        }
+        case 'Any': {
+          return {kind: 'Poly'};
+        }
+        default: {
+          assertExhaustive(
+            typeConfig.name,
+            `Unexpected type '${(typeConfig as any).name}'`,
+          );
+        }
+      }
+    }
+    case 'function': {
+      return addFunction(shapes, [], {
+        positionalParams: typeConfig.positionalParams,
+        restParam: typeConfig.restParam,
+        calleeEffect: typeConfig.calleeEffect,
+        returnType: installTypeConfig(globals, shapes, typeConfig.returnType),
+        returnValueKind: typeConfig.returnValueKind,
+        noAlias: typeConfig.noAlias === true,
+        mutableOnlyIfOperandsAreMutable:
+          typeConfig.mutableOnlyIfOperandsAreMutable === true,
+      });
+    }
+    case 'hook': {
+      return addHook(shapes, {
+        hookKind: 'Custom',
+        positionalParams: typeConfig.positionalParams ?? [],
+        restParam: typeConfig.restParam ?? Effect.Freeze,
+        calleeEffect: Effect.Read,
+        returnType: installTypeConfig(globals, shapes, typeConfig.returnType),
+        returnValueKind: typeConfig.returnValueKind ?? ValueKind.Frozen,
+        noAlias: typeConfig.noAlias === true,
+      });
+    }
+    case 'object': {
+      return addObject(
+        shapes,
+        null,
+        Object.entries(typeConfig.properties ?? {}).map(([key, value]) => [
+          key,
+          installTypeConfig(globals, shapes, value),
+        ]),
+      );
+    }
+    default: {
+      assertExhaustive(
+        typeConfig,
+        `Unexpected type kind '${(typeConfig as any).kind}'`,
+      );
+    }
+  }
+}
+
 export function installReAnimatedTypes(
   globals: GlobalRegistry,
   registry: ShapeRegistry,

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

@@ -12,6 +12,7 @@ import {assertExhaustive} from '../Utils/utils';
 import {Environment, ReactFunctionType} from './Environment';
 import {HookKind} from './ObjectShape';
 import {Type, makeType} from './Types';
+import {z} from 'zod';
 
 /*
  * *******************************************************************************************
@@ -1360,6 +1361,15 @@ export enum ValueKind {
   Context = 'context',
 }
 
+export const ValueKindSchema = z.enum([
+  ValueKind.MaybeFrozen,
+  ValueKind.Frozen,
+  ValueKind.Primitive,
+  ValueKind.Global,
+  ValueKind.Mutable,
+  ValueKind.Context,
+]);
+
 // The effect with which a value is modified.
 export enum Effect {
   // Default value: not allowed after lifetime inference
@@ -1389,6 +1399,15 @@ export enum Effect {
   Store = 'store',
 }
 
+export const EffectSchema = z.enum([
+  Effect.Read,
+  Effect.Mutate,
+  Effect.ConditionallyMutate,
+  Effect.Capture,
+  Effect.Store,
+  Effect.Freeze,
+]);
+
 export function isMutableEffect(
   effect: Effect,
   location: SourceLocation,

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

@@ -0,0 +1,105 @@
+/**
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+import {isValidIdentifier} from '@babel/types';
+import {z} from 'zod';
+import {Effect, ValueKind} from '..';
+import {EffectSchema, ValueKindSchema} from './HIR';
+
+export type ObjectPropertiesConfig = {[key: string]: TypeConfig};
+export const ObjectPropertiesSchema: z.ZodType<ObjectPropertiesConfig> = z
+  .record(
+    z.string(),
+    z.lazy(() => TypeSchema),
+  )
+  .refine(record => {
+    return Object.keys(record).every(
+      key => key === '*' || key === 'default' || isValidIdentifier(key),
+    );
+  }, 'Expected all "object" property names to be valid identifier, `*` to match any property, of `default` to define a module default export');
+
+export type ObjectTypeConfig = {
+  kind: 'object';
+  properties: ObjectPropertiesConfig | null;
+};
+export const ObjectTypeSchema: z.ZodType<ObjectTypeConfig> = z.object({
+  kind: z.literal('object'),
+  properties: ObjectPropertiesSchema.nullable(),
+});
+
+export type FunctionTypeConfig = {
+  kind: 'function';
+  positionalParams: Array<Effect>;
+  restParam: Effect | null;
+  calleeEffect: Effect;
+  returnType: TypeConfig;
+  returnValueKind: ValueKind;
+  noAlias?: boolean | null | undefined;
+  mutableOnlyIfOperandsAreMutable?: boolean | null | undefined;
+};
+export const FunctionTypeSchema: z.ZodType<FunctionTypeConfig> = z.object({
+  kind: z.literal('function'),
+  positionalParams: z.array(EffectSchema),
+  restParam: EffectSchema.nullable(),
+  calleeEffect: EffectSchema,
+  returnType: z.lazy(() => TypeSchema),
+  returnValueKind: ValueKindSchema,
+  noAlias: z.boolean().nullable().optional(),
+  mutableOnlyIfOperandsAreMutable: z.boolean().nullable().optional(),
+});
+
+export type HookTypeConfig = {
+  kind: 'hook';
+  positionalParams?: Array<Effect> | null | undefined;
+  restParam?: Effect | null | undefined;
+  returnType: TypeConfig;
+  returnValueKind?: ValueKind | null | undefined;
+  noAlias?: boolean | null | undefined;
+};
+export const HookTypeSchema: z.ZodType<HookTypeConfig> = z.object({
+  kind: z.literal('hook'),
+  positionalParams: z.array(EffectSchema).nullable().optional(),
+  restParam: EffectSchema.nullable().optional(),
+  returnType: z.lazy(() => TypeSchema),
+  returnValueKind: ValueKindSchema.nullable().optional(),
+  noAlias: z.boolean().nullable().optional(),
+});
+
+export type BuiltInTypeConfig =
+  | 'Any'
+  | 'Ref'
+  | 'Array'
+  | 'Primitive'
+  | 'MixedReadonly';
+export const BuiltInTypeSchema: z.ZodType<BuiltInTypeConfig> = z.union([
+  z.literal('Any'),
+  z.literal('Ref'),
+  z.literal('Array'),
+  z.literal('Primitive'),
+  z.literal('MixedReadonly'),
+]);
+
+export type TypeReferenceConfig = {
+  kind: 'type';
+  name: BuiltInTypeConfig;
+};
+export const TypeReferenceSchema: z.ZodType<TypeReferenceConfig> = z.object({
+  kind: z.literal('type'),
+  name: BuiltInTypeSchema,
+});
+
+export type TypeConfig =
+  | ObjectTypeConfig
+  | FunctionTypeConfig
+  | HookTypeConfig
+  | TypeReferenceConfig;
+export const TypeSchema: z.ZodType<TypeConfig> = z.union([
+  ObjectTypeSchema,
+  FunctionTypeSchema,
+  HookTypeSchema,
+  TypeReferenceSchema,
+]);

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

@@ -127,7 +127,7 @@ function collectTemporaries(
       break;
     }
     case 'LoadGlobal': {
-      const global = env.getGlobalDeclaration(value.binding);
+      const global = env.getGlobalDeclaration(value.binding, value.loc);
       const hookKind = global !== null ? getHookKindForType(env, global) : null;
       const lvalId = instr.lvalue.identifier.id;
       if (hookKind === 'useMemo' || hookKind === 'useCallback') {

compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts

@@ -227,7 +227,7 @@ function* generateInstructionTypes(
     }
 
     case 'LoadGlobal': {
-      const globalType = env.getGlobalDeclaration(value.binding);
+      const globalType = env.getGlobalDeclaration(value.binding, value.loc);
       if (globalType) {
         yield equation(left, globalType);
       }