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 7a83f7e3a0..c0a38f8241 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts @@ -16,6 +16,7 @@ import { DEFAULT_SHAPES, Global, GlobalRegistry, + installCustomGlobals, installReAnimatedTypes, } from './Globals'; import { @@ -125,18 +126,72 @@ const HookSchema = z.object({ export type Hook = z.infer; /* - * TODO(mofeiZ): User defined global types (with corresponding shapes). - * User defined global types should have inline ObjectShapes instead of directly - * using ObjectShapes.ShapeRegistry, as a user-provided ShapeRegistry may be - * accidentally be not well formed. - * i.e. - * missing required shapes (BuiltInArray for [] and BuiltInObject for {}) - * missing some recursive Object / Function shapeIds + * Extremely hacky Zod typing in order to get recursive types working for + * external type schema definitions. */ +const EffectSchema = z.enum([ + Effect.Read, + Effect.Mutate, + Effect.ConditionallyMutate, + Effect.Capture, + Effect.Store, + Effect.Freeze, +]); + +const TypeIDSchema = z.number(); + +const TypeReferenceSchema = z.union([TypeIDSchema, z.literal('array')]); + +const FunctionTypeSchema = z.object({ + positionalParams: z.array(EffectSchema), + restParam: EffectSchema.nullable(), + calleeEffect: EffectSchema, + returnType: TypeReferenceSchema.nullable(), + returnValueKind: z.nativeEnum(ValueKind), +}); + +const baseTypeSchema = z.object({ + id: TypeIDSchema.nullable().default(null), + fn: FunctionTypeSchema.nullable(), +}); + +export type ZodCompositeType< + T extends z.ZodTypeAny, + U extends {[K in keyof U]: z.ZodType}, +> = z.ZodType< + z.output & {[K in keyof U]: z.output}, + z.ZodTypeDef, + z.input & {[K in keyof U]: z.input} +>; + +export const TypeSchema: ZodCompositeType< + typeof baseTypeSchema, + { + properties: typeof propertiesField; + } +> = baseTypeSchema.extend({ + properties: z.lazy(() => propertiesField), +}); + +const propertiesField = z.array(z.tuple([z.string(), TypeSchema])); + +export type GlobalEffect = z.infer; + +export type GlobalType = z.infer & { + properties: Array<[string, GlobalType]>; +}; + +export type GlobalFunctionType = z.infer; + const EnvironmentConfigSchema = z.object({ customHooks: z.map(z.string(), HookSchema).optional().default(new Map()), + typedGlobals: z + .array(z.tuple([z.string(), TypeSchema])) + .optional() + .default([]), + /** * 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 @@ -514,12 +569,46 @@ export function parseConfigPragma(pragma: string): EnvironmentConfig { props.push({type: 'name', name: elt}); } } - console.log([valSplit[0], props.map(x => x.name ?? '*').join('.')]); maybeConfig[key] = [[valSplit[0], props]]; } continue; } + if (key === 'customType') { + maybeConfig['typedGlobals'] = [ + [ + 'custom', + { + id: null, + fn: { + positionalParams: [], + restParam: 'read' as Effect.Read, + calleeEffect: 'read' as Effect.Read, + returnType: null, + returnValueKind: 'primitive' as ValueKind.Primitive, + }, + properties: [ + [ + 'prop', + { + id: null, + fn: { + positionalParams: [], + restParam: 'read' as Effect.Read, + calleeEffect: 'read' as Effect.Read, + returnType: null, + returnValueKind: 'primitive' as ValueKind.Primitive, + }, + properties: [], + }, + ], + ], + }, + ], + ]; + continue; + } + if (typeof defaultConfig[key as keyof EnvironmentConfig] !== 'boolean') { // skip parsing non-boolean properties continue; @@ -644,6 +733,8 @@ export class Environment { installReAnimatedTypes(this.#globals, this.#shapes); } + installCustomGlobals(this.#globals, this.#shapes, this.config.typedGlobals); + this.#contextIdentifiers = contextIdentifiers; this.#hoistedIdentifiers = new Set(); } @@ -756,6 +847,7 @@ export class Environment { return ( moduleName.toLowerCase() === 'react' || moduleName.toLowerCase() === 'react-dom' || + [...this.#globals.keys()].includes(moduleName) || (this.config.enableSharedRuntime__testonly && moduleName === 'shared-runtime') ); 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 884372b986..3c577c2e8e 100644 --- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts +++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts @@ -5,6 +5,8 @@ * LICENSE file in the root directory of this source tree. */ +import {CompilerError} from '..'; +import {GlobalType} from './Environment'; import {Effect, ValueKind, ValueReason} from './HIR'; import { BUILTIN_SHAPES, @@ -18,6 +20,7 @@ import { BuiltInUseReducerId, BuiltInUseRefId, BuiltInUseStateId, + FunctionSignature, ShapeRegistry, addFunction, addHook, @@ -516,6 +519,68 @@ DEFAULT_GLOBALS.set( addObject(DEFAULT_SHAPES, 'global', TYPED_GLOBALS), ); +export function installCustomGlobals( + globals: GlobalRegistry, + registry: ShapeRegistry, + custom: Array<[string, GlobalType]>, +): void { + const refMap: Map = new Map(); + const toSet: Map void> = new Map(); + function installShape(name: string, type: GlobalType): BuiltInType { + const props: Array<[string, BuiltInType]> = type.properties.map( + ([name, type]) => [name, installShape(name, type)], + ); + + let ty; + if (type.fn == null) { + ty = addObject(registry, name, props); + } else { + const func: Omit = { + positionalParams: type.fn.positionalParams, + restParam: type.fn.restParam, + calleeEffect: type.fn.calleeEffect, + returnValueKind: type.fn.returnValueKind, + returnType: {kind: 'Primitive'}, // Placeholder, mutated below + }; + + if (type.fn.returnType == null) { + func.returnType = {kind: 'Poly'}; + } else if (type.fn.returnType === 'array') { + func.returnType = {kind: 'Object', shapeId: BuiltInArrayId}; + } else if (refMap.has(type.fn.returnType)) { + func.returnType = refMap.get(type.fn.returnType)!; + } else { + const cur = toSet.get(type.fn.returnType); + toSet.set(type.fn.returnType, (ty: BuiltInType) => { + cur?.(ty); + func.returnType = ty; + }); + } + ty = addFunction(registry, props, func, name); + } + + if (type.id != null) { + CompilerError.invariant(refMap.get(type.id) == null, { + reason: 'Duplicate type id', + loc: null, + }); + refMap.set(type.id, ty); + toSet.get(type.id)?.(ty); + toSet.delete(type.id); + } + return ty; + } + + for (const [name, type] of custom) { + globals.set(name, installShape(name, type)); + } + + CompilerError.invariant(toSet.size === 0, { + reason: 'Unresolved type ids', + loc: null, + }); +} + export function installReAnimatedTypes( globals: GlobalRegistry, registry: ShapeRegistry, diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/custom-type.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/custom-type.expect.md new file mode 100644 index 0000000000..85d70008f9 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/custom-type.expect.md @@ -0,0 +1,53 @@ + +## Input + +```javascript +// @customType + +import custom from 'custom'; + +function Component(props) { + const x = [props.x]; + const y = [props.y]; + + useHook(); + + custom(x); + custom.prop(x); + custom.notPresent(y); + + return ; +} + +``` + +## Code + +```javascript +import { c as _c } from "react/compiler-runtime"; // @customType + +import custom from "custom"; + +function Component(props) { + const $ = _c(2); + let t0; + if ($[0] !== props.x) { + t0 = [props.x]; + $[0] = props.x; + $[1] = t0; + } else { + t0 = $[1]; + } + const x = t0; + const y = [props.y]; + + useHook(); + + custom(x); + custom.prop(x); + custom.notPresent(y); + return ; +} + +``` + \ No newline at end of file diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/custom-type.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/custom-type.js new file mode 100644 index 0000000000..0773160b39 --- /dev/null +++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/custom-type.js @@ -0,0 +1,16 @@ +// @customType + +import custom from 'custom'; + +function Component(props) { + const x = [props.x]; + const y = [props.y]; + + useHook(); + + custom(x); + custom.prop(x); + custom.notPresent(y); + + return ; +} diff --git a/compiler/packages/snap/src/SproutTodoFilter.ts b/compiler/packages/snap/src/SproutTodoFilter.ts index 2d87543026..80c20fc7c0 100644 --- a/compiler/packages/snap/src/SproutTodoFilter.ts +++ b/compiler/packages/snap/src/SproutTodoFilter.ts @@ -502,6 +502,7 @@ const skipFilter = new Set([ // Depends on external functions 'idx-method-no-outlining-wildcard', 'idx-method-no-outlining', + 'custom-type', // needs to be executed as a module 'meta-property', diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts index de0184f0bd..76656464ed 100644 --- a/compiler/packages/snap/src/compiler.ts +++ b/compiler/packages/snap/src/compiler.ts @@ -21,6 +21,7 @@ import type { } from 'babel-plugin-react-compiler/src/Entrypoint'; import type {Effect, ValueKind} from 'babel-plugin-react-compiler/src/HIR'; import type { + GlobalType, Macro, MacroMethod, parseConfigPragma as ParseConfigPragma, @@ -52,6 +53,39 @@ function makePluginOptions( let enableChangeDetectionForDebugging = null; let customMacros: null | Array = null; let validateBlocklistedImports = null; + let typedGlobals: Array<[string, GlobalType]> = []; + + if (firstLine.indexOf('@customType') !== -1) { + typedGlobals.push([ + 'custom', + { + id: null, + fn: { + positionalParams: [], + restParam: 'read' as Effect.Read, + calleeEffect: 'read' as Effect.Read, + returnType: null, + returnValueKind: 'primitive' as ValueKind.Primitive, + }, + properties: [ + [ + 'prop', + { + id: null, + fn: { + positionalParams: [], + restParam: 'read' as Effect.Read, + calleeEffect: 'read' as Effect.Read, + returnType: null, + returnValueKind: 'primitive' as ValueKind.Primitive, + }, + properties: [], + }, + ], + ], + }, + ]); + } if (firstLine.indexOf('@compilationMode(annotation)') !== -1) { assert(