[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: 8381ca30f7
Pull Request resolved: https://github.com/facebook/react/pull/30771

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

@@ -17,6 +17,7 @@ import {
   Global,
   GlobalRegistry,
   installReAnimatedTypes,
+  installTypeConfig,
 } from './Globals';
 import {
   BlockId,
@@ -45,6 +46,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
@@ -124,6 +126,11 @@ const HookSchema = z.object({
 
 export type Hook = z.infer<typeof HookSchema>;
 
+export const ModuleTypeResolver = z
+  .function()
+  .args(z.string())
+  .returns(z.nullable(TypeSchema));
+
 /*
  * TODO(mofeiZ): User defined global types (with corresponding shapes).
  * User defined global types should have inline ObjectShapes instead of directly
@@ -137,6 +144,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.
+   */
+  resolveModuleTypeSchema: z.nullable(ModuleTypeResolver).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 +590,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,6 +712,28 @@ export class Environment {
     return this.#outlinedFunctions;
   }
 
+  #resolveModuleType(moduleName: string): Global | null {
+    if (this.config.resolveModuleTypeSchema == null) {
+      return null;
+    }
+    let moduleType = this.#moduleTypes.get(moduleName);
+    if (moduleType === undefined) {
+      const moduleConfig = this.config.resolveModuleTypeSchema(moduleName);
+      if (moduleConfig != null) {
+        const moduleTypes = TypeSchema.parse(moduleConfig);
+        moduleType = installTypeConfig(
+          this.#globals,
+          this.#shapes,
+          moduleTypes,
+        );
+      } else {
+        moduleType = null;
+      }
+      this.#moduleTypes.set(moduleName, moduleType);
+    }
+    return moduleType;
+  }
+
   getGlobalDeclaration(binding: NonLocalBinding): Global | null {
     if (this.config.hookPattern != null) {
       const match = new RegExp(this.config.hookPattern).exec(binding.name);
@@ -736,6 +772,17 @@ export class Environment {
             (isHookName(binding.imported) ? this.#getCustomHookType() : null)
           );
         } else {
+          const moduleType = this.#resolveModuleType(binding.module);
+          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 +805,11 @@ export class Environment {
             (isHookName(binding.name) ? this.#getCustomHookType() : null)
           );
         } else {
+          const moduleType = this.#resolveModuleType(binding.module);
+          if (moduleType !== null) {
+            // TODO: distinguish default/namespace cases
+            return moduleType;
+          }
           return isHookName(binding.name) ? this.#getCustomHookType() : null;
         }
       }

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,56 @@ 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};
+        }
+        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,
+      });
+    }
+    case 'object': {
+      return addObject(
+        shapes,
+        null,
+        Object.entries(typeConfig.properties ?? {}).map(([key, value]) => [
+          key,
+          installTypeConfig(globals, shapes, value),
+        ]),
+      );
+    }
+  }
+}
+
 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';
 
 /*
  * *******************************************************************************************
@@ -1389,6 +1390,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,76 @@
+/**
+ * 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} 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 === '*' || isValidIdentifier(key),
+    );
+  }, 'Expected all "object" property names to be valid identifiers or `*` to match any property');
+
+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;
+};
+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: z.nativeEnum(ValueKind),
+});
+
+export type BuiltInTypeConfig = 'Ref' | 'Array' | 'Primitive' | 'MixedReadonly';
+export const BuiltInTypeSchema: z.ZodType<BuiltInTypeConfig> = z.union([
+  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
+  | TypeReferenceConfig;
+export const TypeSchema: z.ZodType<TypeConfig> = z.union([
+  ObjectTypeSchema,
+  FunctionTypeSchema,
+  TypeReferenceSchema,
+]);