From 6be688d88d66e16701d5fab08c364d40e4dcbdbb Mon Sep 17 00:00:00 2001 From: Mike Vitousek Date: Mon, 12 Aug 2024 22:00:15 -0400 Subject: [PATCH] [compiler] Custom type definitions in config Summary: This PR allows custom type definitions for globals to be loaded in the Forget config. We use a slightly different representation than we do internally, in order to ensure that things are well formed before we pass them to the shape registry. The most fundamental issue is that we introduce a notion of IDs (expected to be unique) for types: we can introduce a type and then refer to that type later by a unique ID in a function return type if we have, for example, a global value that returns itself when called. [ghstack-poisoned] --- .../src/HIR/Environment.ts | 108 ++++++++++++++++-- .../src/HIR/Globals.ts | 65 +++++++++++ .../fixtures/compiler/custom-type.expect.md | 53 +++++++++ .../fixtures/compiler/custom-type.js | 16 +++ .../packages/snap/src/SproutTodoFilter.ts | 1 + compiler/packages/snap/src/compiler.ts | 34 ++++++ 6 files changed, 269 insertions(+), 8 deletions(-) create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/custom-type.expect.md create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/custom-type.js 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(