[compiler] Custom type definitions in config

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(