diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts index a5614ac244..58e818205f 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/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; +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; 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 = 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; } } diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts index e9066f85b8..e39b04e002 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts +++ b/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, diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts index 0810130102..32a67bd755 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts +++ b/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, diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts new file mode 100644 index 0000000000..66b949da76 --- /dev/null +++ b/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 = 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 = z.object({ + kind: z.literal('object'), + properties: ObjectPropertiesSchema.nullable(), +}); + +export type FunctionTypeConfig = { + kind: 'function'; + positionalParams: Array; + restParam: Effect | null; + calleeEffect: Effect; + returnType: TypeConfig; + returnValueKind: ValueKind; +}; +export const FunctionTypeSchema: z.ZodType = 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 = 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 = z.object({ + kind: z.literal('type'), + name: BuiltInTypeSchema, +}); + +export type TypeConfig = + | ObjectTypeConfig + | FunctionTypeConfig + | TypeReferenceConfig; +export const TypeSchema: z.ZodType = z.union([ + ObjectTypeSchema, + FunctionTypeSchema, + TypeReferenceSchema, +]);