[compiler] Type provider infra for tests

ghstack-source-id: d206061
Pull Request resolved: #30776

Author: josephsavona

compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts

@@ -126,11 +126,6 @@ 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
@@ -148,7 +143,7 @@ const EnvironmentConfigSchema = z.object({
    * 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),
+  moduleTypeProvider: z.nullable(z.function().args(z.string())).default(null),
 
   /**
    * A list of functions which the application compiles as macros, where
@@ -713,20 +708,30 @@ export class Environment {
   }
 
   #resolveModuleType(moduleName: string): Global | null {
-    if (this.config.resolveModuleTypeSchema == null) {
+    if (this.config.moduleTypeProvider == 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 {
+      try {
+        const unparsedModuleConfig = this.config.moduleTypeProvider(moduleName);
+        if (unparsedModuleConfig != null) {
+          const parsedModuleConfig = TypeSchema.safeParse(unparsedModuleConfig);
+          if (!parsedModuleConfig.success) {
+            console.error(parsedModuleConfig.error.toString());
+            process.exit(1);
+          }
+          const moduleConfig = parsedModuleConfig.data;
+          moduleType = installTypeConfig(
+            this.#globals,
+            this.#shapes,
+            moduleConfig,
+          );
+        } else {
+          moduleType = null;
+        }
+      } catch (e) {
+        console.error((e as Error).stack);
         moduleType = null;
       }
       this.#moduleTypes.set(moduleName, moduleType);
@@ -819,9 +824,7 @@ export class Environment {
   #isKnownReactModule(moduleName: string): boolean {
     return (
       moduleName.toLowerCase() === 'react' ||
-      moduleName.toLowerCase() === 'react-dom' ||
-      (this.config.enableSharedRuntime__testonly &&
-        moduleName === 'shared-runtime')
+      moduleName.toLowerCase() === 'react-dom'
     );
   }
 

compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts

@@ -1361,6 +1361,15 @@ export enum ValueKind {
   Context = 'context',
 }
 
+export const ValueKindSchema = z.enum([
+  ValueKind.MaybeFrozen,
+  ValueKind.Frozen,
+  ValueKind.Primitive,
+  ValueKind.Global,
+  ValueKind.Mutable,
+  ValueKind.Context,
+]);
+
 // The effect with which a value is modified.
 export enum Effect {
   // Default value: not allowed after lifetime inference

compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts

@@ -8,7 +8,7 @@
 import {isValidIdentifier} from '@babel/types';
 import {z} from 'zod';
 import {Effect, ValueKind} from '..';
-import {EffectSchema} from './HIR';
+import {EffectSchema, ValueKindSchema} from './HIR';
 
 export type ObjectPropertiesConfig = {[key: string]: TypeConfig};
 export const ObjectPropertiesSchema: z.ZodType<ObjectPropertiesConfig> = z
@@ -45,7 +45,7 @@ export const FunctionTypeSchema: z.ZodType<FunctionTypeConfig> = z.object({
   restParam: EffectSchema.nullable(),
   calleeEffect: EffectSchema,
   returnType: z.lazy(() => TypeSchema),
-  returnValueKind: z.nativeEnum(ValueKind),
+  returnValueKind: ValueKindSchema,
 });
 
 export type BuiltInTypeConfig = 'Ref' | 'Array' | 'Primitive' | 'MixedReadonly';

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/hook-noAlias.expect.md

@@ -27,7 +27,7 @@ import { c as _c } from "react/compiler-runtime";
 import { useNoAlias } from "shared-runtime";
 
 function Component(props) {
-  const $ = _c(5);
+  const $ = _c(9);
   let t0;
   if ($[0] !== props.a) {
     t0 = { a: props.a };
@@ -37,19 +37,35 @@ function Component(props) {
     t0 = $[1];
   }
   const item = t0;
-  const x = useNoAlias(item, () => {
-    console.log(props);
-  }, [props.a]);
   let t1;
-  if ($[2] !== x || $[3] !== item) {
-    t1 = [x, item];
-    $[2] = x;
-    $[3] = item;
-    $[4] = t1;
+  if ($[2] !== props) {
+    t1 = () => {
+      console.log(props);
+    };
+    $[2] = props;
+    $[3] = t1;
+  } else {
+    t1 = $[3];
+  }
+  let t2;
+  if ($[4] !== props.a) {
+    t2 = [props.a];
+    $[4] = props.a;
+    $[5] = t2;
+  } else {
+    t2 = $[5];
+  }
+  const x = useNoAlias(item, t1, t2);
+  let t3;
+  if ($[6] !== x || $[7] !== item) {
+    t3 = [x, item];
+    $[6] = x;
+    $[7] = item;
+    $[8] = t3;
   } else {
-    t1 = $[4];
+    t3 = $[8];
   }
-  return t1;
+  return t3;
 }
 
 export const FIXTURE_ENTRYPOINT = {

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/original-reactive-scopes-fork/allocating-logical-expression-instruction-scope.expect.md

@@ -36,16 +36,23 @@ import { c as _c } from "react/compiler-runtime"; // @enableReactiveScopesInHIR:
 import { useFragment } from "shared-runtime";
 
 function Foo() {
-  const $ = _c(2);
+  const $ = _c(4);
   const data = useFragment();
-  const t0 = data?.toString() || "";
+  let t0;
+  if ($[0] !== data) {
+    t0 = data?.toString() || "";
+    $[0] = data;
+    $[1] = t0;
+  } else {
+    t0 = $[1];
+  }
   let t1;
-  if ($[0] !== t0) {
+  if ($[2] !== t0) {
     t1 = [t0];
-    $[0] = t0;
-    $[1] = t1;
+    $[2] = t0;
+    $[3] = t1;
   } else {
-    t1 = $[1];
+    t1 = $[3];
   }
   return t1;
 }

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log.expect.md

@@ -0,0 +1,145 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {typedLog, ValidateMemoization} from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => ({a}), [a]);
+  const item2 = useMemo(() => ({b}), [b]);
+  typedLog(item1, item2);
+
+  return (
+    <>
+      <ValidateMemoization inputs={[a]} output={item1} />
+      <ValidateMemoization inputs={[b]} output={item2} />
+    </>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 1, b: 2},
+    {a: 2, b: 2},
+    {a: 3, b: 2},
+    {a: 0, b: 0},
+  ],
+};
+
+```
+
+## Code
+
+```javascript
+import { c as _c } from "react/compiler-runtime";
+import { useMemo } from "react";
+import { typedLog, ValidateMemoization } from "shared-runtime";
+
+export function Component(t0) {
+  const $ = _c(17);
+  const { a, b } = t0;
+  let t1;
+  let t2;
+  if ($[0] !== a) {
+    t2 = { a };
+    $[0] = a;
+    $[1] = t2;
+  } else {
+    t2 = $[1];
+  }
+  t1 = t2;
+  const item1 = t1;
+  let t3;
+  let t4;
+  if ($[2] !== b) {
+    t4 = { b };
+    $[2] = b;
+    $[3] = t4;
+  } else {
+    t4 = $[3];
+  }
+  t3 = t4;
+  const item2 = t3;
+  typedLog(item1, item2);
+  let t5;
+  if ($[4] !== a) {
+    t5 = [a];
+    $[4] = a;
+    $[5] = t5;
+  } else {
+    t5 = $[5];
+  }
+  let t6;
+  if ($[6] !== t5 || $[7] !== item1) {
+    t6 = <ValidateMemoization inputs={t5} output={item1} />;
+    $[6] = t5;
+    $[7] = item1;
+    $[8] = t6;
+  } else {
+    t6 = $[8];
+  }
+  let t7;
+  if ($[9] !== b) {
+    t7 = [b];
+    $[9] = b;
+    $[10] = t7;
+  } else {
+    t7 = $[10];
+  }
+  let t8;
+  if ($[11] !== t7 || $[12] !== item2) {
+    t8 = <ValidateMemoization inputs={t7} output={item2} />;
+    $[11] = t7;
+    $[12] = item2;
+    $[13] = t8;
+  } else {
+    t8 = $[13];
+  }
+  let t9;
+  if ($[14] !== t6 || $[15] !== t8) {
+    t9 = (
+      <>
+        {t6}
+        {t8}
+      </>
+    );
+    $[14] = t6;
+    $[15] = t8;
+    $[16] = t9;
+  } else {
+    t9 = $[16];
+  }
+  return t9;
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{ a: 0, b: 0 }],
+  sequentialRenders: [
+    { a: 0, b: 0 },
+    { a: 1, b: 0 },
+    { a: 1, b: 1 },
+    { a: 1, b: 2 },
+    { a: 2, b: 2 },
+    { a: 3, b: 2 },
+    { a: 0, b: 0 },
+  ],
+};
+
+```
+      
+### Eval output
+(kind: ok) <div>{"inputs":[0],"output":{"a":0}}</div><div>{"inputs":[0],"output":{"b":0}}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[0],"output":{"b":0}}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[1],"output":{"b":1}}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[2],"output":{"b":2}}</div>
+<div>{"inputs":[2],"output":{"a":2}}</div><div>{"inputs":[2],"output":{"b":2}}</div>
+<div>{"inputs":[3],"output":{"a":3}}</div><div>{"inputs":[2],"output":{"b":2}}</div>
+<div>{"inputs":[0],"output":{"a":0}}</div><div>{"inputs":[0],"output":{"b":0}}</div>
+logs: [{ a: 0 },{ b: 0 },{ a: 1 },{ b: 0 },{ a: 1 },{ b: 1 },{ a: 1 },{ b: 2 },{ a: 2 },{ b: 2 },{ a: 3 },{ b: 2 },{ a: 0 },{ b: 0 }]

compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log.tsx

@@ -0,0 +1,29 @@
+import {useMemo} from 'react';
+import {typedLog, ValidateMemoization} from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => ({a}), [a]);
+  const item2 = useMemo(() => ({b}), [b]);
+  typedLog(item1, item2);
+
+  return (
+    <>
+      <ValidateMemoization inputs={[a]} output={item1} />
+      <ValidateMemoization inputs={[b]} output={item2} />
+    </>
+  );
+}
+
+export const FIXTURE_ENTRYPOINT = {
+  fn: Component,
+  params: [{a: 0, b: 0}],
+  sequentialRenders: [
+    {a: 0, b: 0},
+    {a: 1, b: 0},
+    {a: 1, b: 1},
+    {a: 1, b: 2},
+    {a: 2, b: 2},
+    {a: 3, b: 2},
+    {a: 0, b: 0},
+  ],
+};