Update base for Update on "[compiler] Type provider infra for tests"

Adds snap/sprout support for type providers, with declarations describing some new functions in shared-runtime. This allows us to verify that the type provider infrastructure is working. I had to make a few fixes to workaround limitations in zod, but i think the types here are mostly stable enough that we won't have too much work keeping things in sync.

[ghstack-poisoned]

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<typeof HookSchema>;
 
+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<typeof HookSchema>;
 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<string, Global | null> = 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;
         }
       }

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,

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,

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<ObjectPropertiesConfig> = 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<ObjectTypeConfig> = z.object({
+  kind: z.literal('object'),
+  properties: ObjectPropertiesSchema.nullable(),
+});
+
+export type FunctionTypeConfig = {
+  kind: 'function';
+  positionalParams: Array<Effect>;
+  restParam: Effect | null;
+  calleeEffect: Effect;
+  returnType: TypeConfig;
+  returnValueKind: ValueKind;
+};
+export const FunctionTypeSchema: z.ZodType<FunctionTypeConfig> = 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<BuiltInTypeConfig> = 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<TypeReferenceConfig> = z.object({
+  kind: z.literal('type'),
+  name: BuiltInTypeSchema,
+});
+
+export type TypeConfig =
+  | ObjectTypeConfig
+  | FunctionTypeConfig
+  | TypeReferenceConfig;
+export const TypeSchema: z.ZodType<TypeConfig> = z.union([
+  ObjectTypeSchema,
+  FunctionTypeSchema,
+  TypeReferenceSchema,
+]);

compiler/packages/babel-plugin-react-compiler/src/Inference/InferReferenceEffects.ts

@@ -453,6 +453,37 @@ class InferenceState {
     }
   }
 
+  freezeValues(values: Set<InstructionValue>, reason: Set<ValueReason>): void {
+    for (const value of values) {
+      this.#values.set(value, {
+        kind: ValueKind.Frozen,
+        reason,
+        context: new Set(),
+      });
+      if (value.kind === 'FunctionExpression') {
+        if (
+          this.#env.config.enablePreserveExistingMemoizationGuarantees ||
+          this.#env.config.enableTransitivelyFreezeFunctionExpressions
+        ) {
+          if (value.kind === 'FunctionExpression') {
+            /*
+             * We want to freeze the captured values, not mark the operands
+             * themselves as frozen. There could be mutations that occur
+             * before the freeze we are processing, and it would be invalid
+             * to overwrite those mutations as a freeze.
+             */
+            for (const operand of eachInstructionValueOperand(value)) {
+              const operandValues = this.#variables.get(operand.identifier.id);
+              if (operandValues !== undefined) {
+                this.freezeValues(operandValues, reason);
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+
   reference(
     place: Place,
     effectKind: Effect,
@@ -482,29 +513,7 @@ class InferenceState {
             reason: reasonSet,
             context: new Set(),
           };
-          values.forEach(value => {
-            this.#values.set(value, {
-              kind: ValueKind.Frozen,
-              reason: reasonSet,
-              context: new Set(),
-            });
-
-            if (
-              this.#env.config.enablePreserveExistingMemoizationGuarantees ||
-              this.#env.config.enableTransitivelyFreezeFunctionExpressions
-            ) {
-              if (value.kind === 'FunctionExpression') {
-                for (const operand of eachInstructionValueOperand(value)) {
-                  this.referenceAndRecordEffects(
-                    operand,
-                    Effect.Freeze,
-                    ValueReason.Other,
-                    [],
-                  );
-                }
-              }
-            }
-          });
+          this.freezeValues(values, reasonSet);
         } else {
           effect = Effect.Read;
         }
@@ -1241,6 +1250,7 @@ function inferBlock(
       case 'ObjectMethod':
       case 'FunctionExpression': {
         let hasMutableOperand = false;
+        const mutableOperands: Array<Place> = [];
         for (const operand of eachInstructionOperand(instr)) {
           state.referenceAndRecordEffects(
             operand,
@@ -1248,6 +1258,9 @@ function inferBlock(
             ValueReason.Other,
             [],
           );
+          if (isMutableEffect(operand.effect, operand.loc)) {
+            mutableOperands.push(operand);
+          }
           hasMutableOperand ||= isMutableEffect(operand.effect, operand.loc);
 
           /**

compiler/packages/babel-plugin-react-compiler/src/Validation/ValidatePreservedManualMemoization.ts

@@ -23,7 +23,7 @@ import {
   ScopeId,
   SourceLocation,
 } from '../HIR';
-import {printManualMemoDependency} from '../HIR/PrintHIR';
+import {printIdentifier, printManualMemoDependency} from '../HIR/PrintHIR';
 import {eachInstructionValueOperand} from '../HIR/visitors';
 import {collectMaybeMemoDependencies} from '../Inference/DropManualMemoization';
 import {
@@ -537,7 +537,9 @@ class Visitor extends ReactiveFunctionVisitor<VisitorState> {
               state.errors.push({
                 reason:
                   'React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output.',
-                description: null,
+                description: DEBUG
+                  ? `${printIdentifier(identifier)} was not memoized`
+                  : null,
                 severity: ErrorSeverity.CannotPreserveMemoization,
                 loc,
                 suggestions: null,

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/allow-mutate-global-in-effect-fixpoint.expect.md

@@ -46,54 +46,57 @@ import { useEffect, useState } from "react";
 let someGlobal = { value: null };
 
 function Component() {
-  const $ = _c(6);
+  const $ = _c(7);
   const [state, setState] = useState(someGlobal);
-
-  let x = someGlobal;
-  while (x == null) {
-    x = someGlobal;
-  }
-
-  const y = x;
   let t0;
   let t1;
+  let t2;
   if ($[0] === Symbol.for("react.memo_cache_sentinel")) {
-    t0 = () => {
+    let x = someGlobal;
+    while (x == null) {
+      x = someGlobal;
+    }
+
+    const y = x;
+    t0 = useEffect;
+    t1 = () => {
       y.value = "hello";
     };
-    t1 = [];
+    t2 = [];
     $[0] = t0;
     $[1] = t1;
+    $[2] = t2;
   } else {
     t0 = $[0];
     t1 = $[1];
+    t2 = $[2];
   }
-  useEffect(t0, t1);
-  let t2;
+  t0(t1, t2);
   let t3;
-  if ($[2] === Symbol.for("react.memo_cache_sentinel")) {
-    t2 = () => {
+  let t4;
+  if ($[3] === Symbol.for("react.memo_cache_sentinel")) {
+    t3 = () => {
       setState(someGlobal.value);
     };
-    t3 = [someGlobal];
-    $[2] = t2;
+    t4 = [someGlobal];
     $[3] = t3;
-  } else {
-    t2 = $[2];
-    t3 = $[3];
-  }
-  useEffect(t2, t3);
-
-  const t4 = String(state);
-  let t5;
-  if ($[4] !== t4) {
-    t5 = <div>{t4}</div>;
     $[4] = t4;
-    $[5] = t5;
   } else {
-    t5 = $[5];
+    t3 = $[3];
+    t4 = $[4];
   }
-  return t5;
+  useEffect(t3, t4);
+
+  const t5 = String(state);
+  let t6;
+  if ($[5] !== t5) {
+    t6 = <div>{t5}</div>;
+    $[5] = t5;
+    $[6] = t6;
+  } else {
+    t6 = $[6];
+  }
+  return t6;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-repro-missed-memoization-from-inferred-mutation-in-logger.expect.md

@@ -0,0 +1,78 @@
+
+## Input
+
+```javascript
+// @flow @validatePreserveExistingMemoizationGuarantees
+import {useFragment} from 'react-relay';
+import LogEvent from 'LogEvent';
+import {useCallback, useMemo} from 'react';
+
+component Component(id) {
+  const items = useFragment();
+
+  const [index, setIndex] = useState(0);
+
+  const logData = useMemo(() => {
+    const item = items[index];
+    return {
+      key: item.key,
+    };
+  }, [index, items]);
+
+  const setCurrentIndex = useCallback(
+    (index: number) => {
+      const object = {
+        tracking: logData.key,
+      };
+      // We infer that this may mutate `object`, which in turn aliases
+      // data from `logData`, such that `logData` may be mutated.
+      LogEvent.log(() => object);
+      setIndex(index);
+    },
+    [index, logData, items]
+  );
+
+  if (prevId !== id) {
+    setCurrentIndex(0);
+  }
+
+  return (
+    <Foo
+      index={index}
+      items={items}
+      current={mediaList[index]}
+      setCurrentIndex={setCurrentIndex}
+    />
+  );
+}
+
+```
+
+
+## Error
+
+```
+   9 |   const [index, setIndex] = useState(0);
+  10 |
+> 11 |   const logData = useMemo(() => {
+     |                   ^^^^^^^^^^^^^^^
+> 12 |     const item = items[index];
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 13 |     return {
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 14 |       key: item.key,
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 15 |     };
+     | ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+> 16 |   }, [index, items]);
+     | ^^^^^^^^^^^^^^^^^^^^^ CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (11:16)
+
+CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This dependency may be mutated later, which could cause the value to change unexpectedly (28:28)
+
+CannotPreserveMemoization: React Compiler has skipped optimizing this component because the existing manual memoization could not be preserved. This value was memoized in source but not in compilation output. (19:27)
+  17 |
+  18 |   const setCurrentIndex = useCallback(
+  19 |     (index: number) => {
+```
+          
+      

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/error.todo-repro-missed-memoization-from-inferred-mutation-in-logger.js

@@ -0,0 +1,43 @@
+// @flow @validatePreserveExistingMemoizationGuarantees
+import {useFragment} from 'react-relay';
+import LogEvent from 'LogEvent';
+import {useCallback, useMemo} from 'react';
+
+component Component(id) {
+  const items = useFragment();
+
+  const [index, setIndex] = useState(0);
+
+  const logData = useMemo(() => {
+    const item = items[index];
+    return {
+      key: item.key,
+    };
+  }, [index, items]);
+
+  const setCurrentIndex = useCallback(
+    (index: number) => {
+      const object = {
+        tracking: logData.key,
+      };
+      // We infer that this may mutate `object`, which in turn aliases
+      // data from `logData`, such that `logData` may be mutated.
+      LogEvent.log(() => object);
+      setIndex(index);
+    },
+    [index, logData, items]
+  );
+
+  if (prevId !== id) {
+    setCurrentIndex(0);
+  }
+
+  return (
+    <Foo
+      index={index}
+      items={items}
+      current={mediaList[index]}
+      setCurrentIndex={setCurrentIndex}
+    />
+  );
+}