Update on "[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]

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,76 @@ const HookSchema = z.object({
 export type Hook = z.infer<typeof HookSchema>;
 
 /*
- * 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'),
+  z.literal('ref'),
+]);
+
+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<any, z.ZodTypeDef, any>},
+> = z.ZodType<
+  z.output<T> & {[K in keyof U]: z.output<U[K]>},
+  z.ZodTypeDef,
+  z.input<T> & {[K in keyof U]: z.input<U[K]>}
+>;
+
+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<typeof EffectSchema>;
+
+export type GlobalType = z.infer<typeof baseTypeSchema> & {
+  properties: Array<[string, GlobalType]>;
+};
+
+export type GlobalFunctionType = z.infer<typeof FunctionTypeSchema>;
+
 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
@@ -520,12 +579,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;
@@ -650,6 +743,8 @@ export class Environment {
       installReAnimatedTypes(this.#globals, this.#shapes);
     }
 
+    installCustomGlobals(this.#globals, this.#shapes, this.config.typedGlobals);
+
     this.#contextIdentifiers = contextIdentifiers;
     this.#hoistedIdentifiers = new Set();
   }
@@ -717,7 +812,7 @@ export class Environment {
         );
       }
       case 'ImportSpecifier': {
-        if (this.#isKnownReactModule(binding.module)) {
+        if (this.#isKnownTypedModule(binding.module)) {
           /**
            * For `import {imported as name} from "..."` form, we use the `imported`
            * name rather than the local alias. Because we don't have definitions for
@@ -745,7 +840,7 @@ export class Environment {
       }
       case 'ImportDefault':
       case 'ImportNamespace': {
-        if (this.#isKnownReactModule(binding.module)) {
+        if (this.#isKnownTypedModule(binding.module)) {
           // only resolve imports to modules we know about
           return (
             this.#globals.get(binding.name) ??
@@ -758,10 +853,11 @@ export class Environment {
     }
   }
 
-  #isKnownReactModule(moduleName: string): boolean {
+  #isKnownTypedModule(moduleName: string): boolean {
     return (
       moduleName.toLowerCase() === 'react' ||
       moduleName.toLowerCase() === 'react-dom' ||
+      this.#globals.has(moduleName) ||
       (this.config.enableSharedRuntime__testonly &&
         moduleName === 'shared-runtime')
     );

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,
@@ -19,6 +21,7 @@ import {
   BuiltInUseRefId,
   BuiltInUseStateId,
   BuiltInUseTransitionId,
+  FunctionSignature,
   ShapeRegistry,
   addFunction,
   addHook,
@@ -528,6 +531,70 @@ DEFAULT_GLOBALS.set(
   addObject(DEFAULT_SHAPES, 'global', TYPED_GLOBALS),
 );
 
+export function installCustomGlobals(
+  globals: GlobalRegistry,
+  registry: ShapeRegistry,
+  custom: Array<[string, GlobalType]>,
+): void {
+  const refMap: Map<number, BuiltInType> = new Map();
+  const toSet: Map<number, (ty: BuiltInType) => 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<FunctionSignature, 'hookKind'> = {
+        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 (type.fn.returnType === 'ref') {
+        func.returnType = {kind: 'Object', shapeId: BuiltInUseRefId};
+      } 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,

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 <Foo x={x} y={y} />;
+}
+
+```
+
+## 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 <Foo x={x} y={y} />;
+}
+
+```
+      

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 <Foo x={x} y={y} />;
+}

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',

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<Macro> = 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(