From 529abd340a2592c82584472e1275df418eb8c04e Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Wed, 21 Aug 2024 15:40:25 -0700
Subject: [PATCH] [compiler] Type provider infra for tests

ghstack-source-id: bfdbf67e3dd0cbfd511bed0bd6ba92266cf99ab8
Pull Request resolved: https://github.com/facebook/react/pull/30776
---
 .../src/HIR/Environment.ts                    |  51 +++--
 .../src/HIR/Globals.ts                        |  23 +++
 .../src/HIR/HIR.ts                            |   9 +
 .../src/HIR/TypeSchema.ts                     |  39 +++-
 .../src/Inference/DropManualMemoization.ts    |   2 +-
 .../src/TypeInference/InferTypes.ts           |   2 +-
 ...ed-scope-declarations-and-locals.expect.md |   4 +
 ...ing-mixed-scope-declarations-and-locals.js |   2 +
 .../compiler/optional-call-logical.expect.md  |   4 +
 .../compiler/optional-call-logical.js         |   2 +
 ...ject-method-calls-mutable-lambda.expect.md |   4 +
 ...only-object-method-calls-mutable-lambda.js |   2 +
 .../readonly-object-method-calls.expect.md    |   4 +
 .../compiler/readonly-object-method-calls.js  |   2 +
 .../tagged-template-in-hook.expect.md         |   4 +
 .../compiler/tagged-template-in-hook.js       |   2 +
 ...type-provider-log-default-import.expect.md | 147 ++++++++++++++
 .../type-provider-log-default-import.tsx      |  30 +++
 .../compiler/type-provider-log.expect.md      | 145 ++++++++++++++
 .../fixtures/compiler/type-provider-log.tsx   |  29 +++
 ...r-store-capture-namespace-import.expect.md | 185 ++++++++++++++++++
 ...rovider-store-capture-namespace-import.tsx |  35 ++++
 .../type-provider-store-capture.expect.md     | 185 ++++++++++++++++++
 .../compiler/type-provider-store-capture.tsx  |  35 ++++
 compiler/packages/snap/src/compiler.ts        |  45 ++---
 compiler/packages/snap/src/constants.ts       |   1 +
 compiler/packages/snap/src/runner-worker.ts   |   6 +
 .../sprout/shared-runtime-type-provider.ts    |  69 +++++++
 .../snap/src/sprout/shared-runtime.ts         |   9 +
 29 files changed, 1020 insertions(+), 57 deletions(-)
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log-default-import.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log-default-import.tsx
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log.tsx
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture-namespace-import.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture-namespace-import.tsx
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture.expect.md
 create mode 100644 compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture.tsx
 create mode 100644 compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts

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 58e818205f439..12c741641c7e0 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
@@ -29,6 +29,7 @@ import {
   NonLocalBinding,
   PolyType,
   ScopeId,
+  SourceLocation,
   Type,
   ValidatedIdentifier,
   ValueKind,
@@ -126,11 +127,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 +144,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
@@ -712,19 +708,27 @@ export class Environment {
     return this.#outlinedFunctions;
   }
 
-  #resolveModuleType(moduleName: string): Global | null {
-    if (this.config.resolveModuleTypeSchema == null) {
+  #resolveModuleType(moduleName: string, loc: SourceLocation): Global | 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);
+      const unparsedModuleConfig = this.config.moduleTypeProvider(moduleName);
+      if (unparsedModuleConfig != null) {
+        const parsedModuleConfig = TypeSchema.safeParse(unparsedModuleConfig);
+        if (!parsedModuleConfig.success) {
+          CompilerError.throwInvalidConfig({
+            reason: `Could not parse module type, the configured \`moduleTypeProvider\` function returned an invalid module description`,
+            description: parsedModuleConfig.error.toString(),
+            loc,
+          });
+        }
+        const moduleConfig = parsedModuleConfig.data;
         moduleType = installTypeConfig(
           this.#globals,
           this.#shapes,
-          moduleTypes,
+          moduleConfig,
         );
       } else {
         moduleType = null;
@@ -734,7 +738,10 @@ export class Environment {
     return moduleType;
   }
 
-  getGlobalDeclaration(binding: NonLocalBinding): Global | null {
+  getGlobalDeclaration(
+    binding: NonLocalBinding,
+    loc: SourceLocation,
+  ): Global | null {
     if (this.config.hookPattern != null) {
       const match = new RegExp(this.config.hookPattern).exec(binding.name);
       if (
@@ -772,7 +779,7 @@ export class Environment {
             (isHookName(binding.imported) ? this.#getCustomHookType() : null)
           );
         } else {
-          const moduleType = this.#resolveModuleType(binding.module);
+          const moduleType = this.#resolveModuleType(binding.module, loc);
           if (moduleType !== null) {
             const importedType = this.getPropertyType(
               moduleType,
@@ -805,10 +812,16 @@ export class Environment {
             (isHookName(binding.name) ? this.#getCustomHookType() : null)
           );
         } else {
-          const moduleType = this.#resolveModuleType(binding.module);
+          const moduleType = this.#resolveModuleType(binding.module, loc);
           if (moduleType !== null) {
-            // TODO: distinguish default/namespace cases
-            return moduleType;
+            if (binding.kind === 'ImportDefault') {
+              const defaultType = this.getPropertyType(moduleType, 'default');
+              if (defaultType !== null) {
+                return defaultType;
+              }
+            } else {
+              return moduleType;
+            }
           }
           return isHookName(binding.name) ? this.#getCustomHookType() : null;
         }
@@ -819,9 +832,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'
     );
   }
 
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 e39b04e002208..2812394300ad5 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Globals.ts
@@ -551,6 +551,9 @@ export function installTypeConfig(
         case 'Ref': {
           return {kind: 'Object', shapeId: BuiltInUseRefId};
         }
+        case 'Any': {
+          return {kind: 'Poly'};
+        }
         default: {
           assertExhaustive(
             typeConfig.name,
@@ -566,6 +569,20 @@ export function installTypeConfig(
         calleeEffect: typeConfig.calleeEffect,
         returnType: installTypeConfig(globals, shapes, typeConfig.returnType),
         returnValueKind: typeConfig.returnValueKind,
+        noAlias: typeConfig.noAlias === true,
+        mutableOnlyIfOperandsAreMutable:
+          typeConfig.mutableOnlyIfOperandsAreMutable === true,
+      });
+    }
+    case 'hook': {
+      return addHook(shapes, {
+        hookKind: 'Custom',
+        positionalParams: typeConfig.positionalParams ?? [],
+        restParam: typeConfig.restParam ?? Effect.Freeze,
+        calleeEffect: Effect.Read,
+        returnType: installTypeConfig(globals, shapes, typeConfig.returnType),
+        returnValueKind: typeConfig.returnValueKind ?? ValueKind.Frozen,
+        noAlias: typeConfig.noAlias === true,
       });
     }
     case 'object': {
@@ -578,6 +595,12 @@ export function installTypeConfig(
         ]),
       );
     }
+    default: {
+      assertExhaustive(
+        typeConfig,
+        `Unexpected type kind '${(typeConfig as any).kind}'`,
+      );
+    }
   }
 }
 
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
index 32a67bd755994..e56c002c513bd 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/HIR.ts
+++ b/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
diff --git a/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts b/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts
index 66b949da76439..362328db72155 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/TypeSchema.ts
+++ b/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
@@ -18,9 +18,9 @@ export const ObjectPropertiesSchema: z.ZodType<ObjectPropertiesConfig> = z
   )
   .refine(record => {
     return Object.keys(record).every(
-      key => key === '*' || isValidIdentifier(key),
+      key => key === '*' || key === 'default' || isValidIdentifier(key),
     );
-  }, 'Expected all "object" property names to be valid identifiers or `*` to match any property');
+  }, 'Expected all "object" property names to be valid identifier, `*` to match any property, of `default` to define a module default export');
 
 export type ObjectTypeConfig = {
   kind: 'object';
@@ -38,6 +38,8 @@ export type FunctionTypeConfig = {
   calleeEffect: Effect;
   returnType: TypeConfig;
   returnValueKind: ValueKind;
+  noAlias?: boolean | null | undefined;
+  mutableOnlyIfOperandsAreMutable?: boolean | null | undefined;
 };
 export const FunctionTypeSchema: z.ZodType<FunctionTypeConfig> = z.object({
   kind: z.literal('function'),
@@ -45,11 +47,36 @@ export const FunctionTypeSchema: z.ZodType<FunctionTypeConfig> = z.object({
   restParam: EffectSchema.nullable(),
   calleeEffect: EffectSchema,
   returnType: z.lazy(() => TypeSchema),
-  returnValueKind: z.nativeEnum(ValueKind),
+  returnValueKind: ValueKindSchema,
+  noAlias: z.boolean().nullable().optional(),
+  mutableOnlyIfOperandsAreMutable: z.boolean().nullable().optional(),
 });
 
-export type BuiltInTypeConfig = 'Ref' | 'Array' | 'Primitive' | 'MixedReadonly';
+export type HookTypeConfig = {
+  kind: 'hook';
+  positionalParams?: Array<Effect> | null | undefined;
+  restParam?: Effect | null | undefined;
+  returnType: TypeConfig;
+  returnValueKind?: ValueKind | null | undefined;
+  noAlias?: boolean | null | undefined;
+};
+export const HookTypeSchema: z.ZodType<HookTypeConfig> = z.object({
+  kind: z.literal('hook'),
+  positionalParams: z.array(EffectSchema).nullable().optional(),
+  restParam: EffectSchema.nullable().optional(),
+  returnType: z.lazy(() => TypeSchema),
+  returnValueKind: ValueKindSchema.nullable().optional(),
+  noAlias: z.boolean().nullable().optional(),
+});
+
+export type BuiltInTypeConfig =
+  | 'Any'
+  | 'Ref'
+  | 'Array'
+  | 'Primitive'
+  | 'MixedReadonly';
 export const BuiltInTypeSchema: z.ZodType<BuiltInTypeConfig> = z.union([
+  z.literal('Any'),
   z.literal('Ref'),
   z.literal('Array'),
   z.literal('Primitive'),
@@ -68,9 +95,11 @@ export const TypeReferenceSchema: z.ZodType<TypeReferenceConfig> = z.object({
 export type TypeConfig =
   | ObjectTypeConfig
   | FunctionTypeConfig
+  | HookTypeConfig
   | TypeReferenceConfig;
 export const TypeSchema: z.ZodType<TypeConfig> = z.union([
   ObjectTypeSchema,
   FunctionTypeSchema,
+  HookTypeSchema,
   TypeReferenceSchema,
 ]);
diff --git a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts b/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
index 2d9e21af1d61b..c9d2a7e1412c3 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/Inference/DropManualMemoization.ts
@@ -127,7 +127,7 @@ function collectTemporaries(
       break;
     }
     case 'LoadGlobal': {
-      const global = env.getGlobalDeclaration(value.binding);
+      const global = env.getGlobalDeclaration(value.binding, value.loc);
       const hookKind = global !== null ? getHookKindForType(env, global) : null;
       const lvalId = instr.lvalue.identifier.id;
       if (hookKind === 'useMemo' || hookKind === 'useCallback') {
diff --git a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts b/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts
index 0b8949e1977ff..4dfeb676a3d0b 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/TypeInference/InferTypes.ts
@@ -227,7 +227,7 @@ function* generateInstructionTypes(
     }
 
     case 'LoadGlobal': {
-      const globalType = env.getGlobalDeclaration(value.binding);
+      const globalType = env.getGlobalDeclaration(value.binding, value.loc);
       if (globalType) {
         yield equation(left, globalType);
       }
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md
index ed566a605fd5a..f69149aba4175 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.expect.md
@@ -2,6 +2,8 @@
 ## Input
 
 ```javascript
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const post = useFragment(
     graphql`
@@ -36,6 +38,8 @@ function Component(props) {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
+import { useFragment } from "shared-runtime";
+
 function Component(props) {
   const $ = _c(4);
   const post = useFragment(
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.js
index f44fe5c57baf3..5d1377c458855 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/destructuring-mixed-scope-declarations-and-locals.js
@@ -1,3 +1,5 @@
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const post = useFragment(
     graphql`
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/optional-call-logical.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/optional-call-logical.expect.md
index 4dc011c85f420..fa6f8cd9bc66c 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/optional-call-logical.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/optional-call-logical.expect.md
@@ -2,6 +2,8 @@
 ## Input
 
 ```javascript
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const item = useFragment(
     graphql`
@@ -20,6 +22,8 @@ function Component(props) {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
+import { useFragment } from "shared-runtime";
+
 function Component(props) {
   const $ = _c(2);
   const item = useFragment(
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/optional-call-logical.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/optional-call-logical.js
index fac4efc9ba90b..6fa11a2dab494 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/optional-call-logical.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/optional-call-logical.js
@@ -1,3 +1,5 @@
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const item = useFragment(
     graphql`
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls-mutable-lambda.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls-mutable-lambda.expect.md
index 0cbd9caf169f0..3abd8cac949c3 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls-mutable-lambda.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls-mutable-lambda.expect.md
@@ -2,6 +2,8 @@
 ## Input
 
 ```javascript
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const x = makeObject();
   const user = useFragment(
@@ -28,6 +30,8 @@ function Component(props) {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
+import { useFragment } from "shared-runtime";
+
 function Component(props) {
   const $ = _c(3);
   const x = makeObject();
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls-mutable-lambda.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls-mutable-lambda.js
index d78be74ac9957..2658a048939d7 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls-mutable-lambda.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls-mutable-lambda.js
@@ -1,3 +1,5 @@
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const x = makeObject();
   const user = useFragment(
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls.expect.md
index 1a5f49631f661..05ab1c533b9b1 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls.expect.md
@@ -2,6 +2,8 @@
 ## Input
 
 ```javascript
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const user = useFragment(
     graphql`
@@ -26,6 +28,8 @@ function Component(props) {
 
 ```javascript
 import { c as _c } from "react/compiler-runtime";
+import { useFragment } from "shared-runtime";
+
 function Component(props) {
   const $ = _c(5);
   const user = useFragment(
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls.js
index 7cccb397cabb2..bb1e52dc7645c 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/readonly-object-method-calls.js
@@ -1,3 +1,5 @@
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const user = useFragment(
     graphql`
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/tagged-template-in-hook.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/tagged-template-in-hook.expect.md
index 3ffc93ada58ea..d52a7713bb875 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/tagged-template-in-hook.expect.md
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/tagged-template-in-hook.expect.md
@@ -2,6 +2,8 @@
 ## Input
 
 ```javascript
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const user = useFragment(
     graphql`
@@ -19,6 +21,8 @@ function Component(props) {
 ## Code
 
 ```javascript
+import { useFragment } from "shared-runtime";
+
 function Component(props) {
   const user = useFragment(
     graphql`
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/tagged-template-in-hook.js b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/tagged-template-in-hook.js
index cb31d2978a610..dbcd2b6d21803 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/tagged-template-in-hook.js
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/tagged-template-in-hook.js
@@ -1,3 +1,5 @@
+import {useFragment} from 'shared-runtime';
+
 function Component(props) {
   const user = useFragment(
     graphql`
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log-default-import.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log-default-import.expect.md
new file mode 100644
index 0000000000000..54d5be2d6bf44
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log-default-import.expect.md
@@ -0,0 +1,147 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {ValidateMemoization} from 'shared-runtime';
+import typedLog 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 { ValidateMemoization } from "shared-runtime";
+import typedLog 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 }]
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log-default-import.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log-default-import.tsx
new file mode 100644
index 0000000000000..ec5dcf41e004c
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log-default-import.tsx
@@ -0,0 +1,30 @@
+import {useMemo} from 'react';
+import {ValidateMemoization} from 'shared-runtime';
+import typedLog 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},
+  ],
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log.expect.md
new file mode 100644
index 0000000000000..072c6d03d9adb
--- /dev/null
+++ b/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 }]
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-log.tsx
new file mode 100644
index 0000000000000..5fb53d9ca85d4
--- /dev/null
+++ b/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},
+  ],
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture-namespace-import.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture-namespace-import.expect.md
new file mode 100644
index 0000000000000..caa74267f326b
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture-namespace-import.expect.md
@@ -0,0 +1,185 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import * as SharedRuntime from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => ({a}), [a]);
+  const item2 = useMemo(() => ({b}), [b]);
+  const items = useMemo(() => {
+    const items = [];
+    SharedRuntime.typedArrayPush(items, item1);
+    SharedRuntime.typedArrayPush(items, item2);
+    return items;
+  }, [item1, item2]);
+
+  return (
+    <>
+      <SharedRuntime.ValidateMemoization inputs={[a]} output={items[0]} />
+      <SharedRuntime.ValidateMemoization inputs={[b]} output={items[1]} />
+      <SharedRuntime.ValidateMemoization inputs={[a, b]} output={items} />
+    </>
+  );
+}
+
+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 * as SharedRuntime from "shared-runtime";
+
+export function Component(t0) {
+  const $ = _c(27);
+  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;
+  let t5;
+  let items;
+  if ($[4] !== item1 || $[5] !== item2) {
+    items = [];
+    SharedRuntime.typedArrayPush(items, item1);
+    SharedRuntime.typedArrayPush(items, item2);
+    $[4] = item1;
+    $[5] = item2;
+    $[6] = items;
+  } else {
+    items = $[6];
+  }
+  t5 = items;
+  const items_0 = t5;
+  let t6;
+  if ($[7] !== a) {
+    t6 = [a];
+    $[7] = a;
+    $[8] = t6;
+  } else {
+    t6 = $[8];
+  }
+  const t7 = items_0[0];
+  let t8;
+  if ($[9] !== t6 || $[10] !== t7) {
+    t8 = <SharedRuntime.ValidateMemoization inputs={t6} output={t7} />;
+    $[9] = t6;
+    $[10] = t7;
+    $[11] = t8;
+  } else {
+    t8 = $[11];
+  }
+  let t9;
+  if ($[12] !== b) {
+    t9 = [b];
+    $[12] = b;
+    $[13] = t9;
+  } else {
+    t9 = $[13];
+  }
+  const t10 = items_0[1];
+  let t11;
+  if ($[14] !== t9 || $[15] !== t10) {
+    t11 = <SharedRuntime.ValidateMemoization inputs={t9} output={t10} />;
+    $[14] = t9;
+    $[15] = t10;
+    $[16] = t11;
+  } else {
+    t11 = $[16];
+  }
+  let t12;
+  if ($[17] !== a || $[18] !== b) {
+    t12 = [a, b];
+    $[17] = a;
+    $[18] = b;
+    $[19] = t12;
+  } else {
+    t12 = $[19];
+  }
+  let t13;
+  if ($[20] !== t12 || $[21] !== items_0) {
+    t13 = <SharedRuntime.ValidateMemoization inputs={t12} output={items_0} />;
+    $[20] = t12;
+    $[21] = items_0;
+    $[22] = t13;
+  } else {
+    t13 = $[22];
+  }
+  let t14;
+  if ($[23] !== t8 || $[24] !== t11 || $[25] !== t13) {
+    t14 = (
+      <>
+        {t8}
+        {t11}
+        {t13}
+      </>
+    );
+    $[23] = t8;
+    $[24] = t11;
+    $[25] = t13;
+    $[26] = t14;
+  } else {
+    t14 = $[26];
+  }
+  return t14;
+}
+
+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":[0,0],"output":[{"a":0},{"b":0}]}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[0],"output":{"b":0}}</div><div>{"inputs":[1,0],"output":[{"a":1},{"b":0}]}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[1],"output":{"b":1}}</div><div>{"inputs":[1,1],"output":[{"a":1},{"b":1}]}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[2],"output":{"b":2}}</div><div>{"inputs":[1,2],"output":[{"a":1},{"b":2}]}</div>
+<div>{"inputs":[2],"output":{"a":2}}</div><div>{"inputs":[2],"output":{"b":2}}</div><div>{"inputs":[2,2],"output":[{"a":2},{"b":2}]}</div>
+<div>{"inputs":[3],"output":{"a":3}}</div><div>{"inputs":[2],"output":{"b":2}}</div><div>{"inputs":[3,2],"output":[{"a":3},{"b":2}]}</div>
+<div>{"inputs":[0],"output":{"a":0}}</div><div>{"inputs":[0],"output":{"b":0}}</div><div>{"inputs":[0,0],"output":[{"a":0},{"b":0}]}</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture-namespace-import.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture-namespace-import.tsx
new file mode 100644
index 0000000000000..6479df9a5a86e
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture-namespace-import.tsx
@@ -0,0 +1,35 @@
+import {useMemo} from 'react';
+import * as SharedRuntime from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => ({a}), [a]);
+  const item2 = useMemo(() => ({b}), [b]);
+  const items = useMemo(() => {
+    const items = [];
+    SharedRuntime.typedArrayPush(items, item1);
+    SharedRuntime.typedArrayPush(items, item2);
+    return items;
+  }, [item1, item2]);
+
+  return (
+    <>
+      <SharedRuntime.ValidateMemoization inputs={[a]} output={items[0]} />
+      <SharedRuntime.ValidateMemoization inputs={[b]} output={items[1]} />
+      <SharedRuntime.ValidateMemoization inputs={[a, b]} output={items} />
+    </>
+  );
+}
+
+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},
+  ],
+};
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture.expect.md b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture.expect.md
new file mode 100644
index 0000000000000..a92abd4ca597c
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture.expect.md
@@ -0,0 +1,185 @@
+
+## Input
+
+```javascript
+import {useMemo} from 'react';
+import {typedArrayPush, ValidateMemoization} from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => ({a}), [a]);
+  const item2 = useMemo(() => ({b}), [b]);
+  const items = useMemo(() => {
+    const items = [];
+    typedArrayPush(items, item1);
+    typedArrayPush(items, item2);
+    return items;
+  }, [item1, item2]);
+
+  return (
+    <>
+      <ValidateMemoization inputs={[a]} output={items[0]} />
+      <ValidateMemoization inputs={[b]} output={items[1]} />
+      <ValidateMemoization inputs={[a, b]} output={items} />
+    </>
+  );
+}
+
+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 { typedArrayPush, ValidateMemoization } from "shared-runtime";
+
+export function Component(t0) {
+  const $ = _c(27);
+  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;
+  let t5;
+  let items;
+  if ($[4] !== item1 || $[5] !== item2) {
+    items = [];
+    typedArrayPush(items, item1);
+    typedArrayPush(items, item2);
+    $[4] = item1;
+    $[5] = item2;
+    $[6] = items;
+  } else {
+    items = $[6];
+  }
+  t5 = items;
+  const items_0 = t5;
+  let t6;
+  if ($[7] !== a) {
+    t6 = [a];
+    $[7] = a;
+    $[8] = t6;
+  } else {
+    t6 = $[8];
+  }
+  const t7 = items_0[0];
+  let t8;
+  if ($[9] !== t6 || $[10] !== t7) {
+    t8 = <ValidateMemoization inputs={t6} output={t7} />;
+    $[9] = t6;
+    $[10] = t7;
+    $[11] = t8;
+  } else {
+    t8 = $[11];
+  }
+  let t9;
+  if ($[12] !== b) {
+    t9 = [b];
+    $[12] = b;
+    $[13] = t9;
+  } else {
+    t9 = $[13];
+  }
+  const t10 = items_0[1];
+  let t11;
+  if ($[14] !== t9 || $[15] !== t10) {
+    t11 = <ValidateMemoization inputs={t9} output={t10} />;
+    $[14] = t9;
+    $[15] = t10;
+    $[16] = t11;
+  } else {
+    t11 = $[16];
+  }
+  let t12;
+  if ($[17] !== a || $[18] !== b) {
+    t12 = [a, b];
+    $[17] = a;
+    $[18] = b;
+    $[19] = t12;
+  } else {
+    t12 = $[19];
+  }
+  let t13;
+  if ($[20] !== t12 || $[21] !== items_0) {
+    t13 = <ValidateMemoization inputs={t12} output={items_0} />;
+    $[20] = t12;
+    $[21] = items_0;
+    $[22] = t13;
+  } else {
+    t13 = $[22];
+  }
+  let t14;
+  if ($[23] !== t8 || $[24] !== t11 || $[25] !== t13) {
+    t14 = (
+      <>
+        {t8}
+        {t11}
+        {t13}
+      </>
+    );
+    $[23] = t8;
+    $[24] = t11;
+    $[25] = t13;
+    $[26] = t14;
+  } else {
+    t14 = $[26];
+  }
+  return t14;
+}
+
+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":[0,0],"output":[{"a":0},{"b":0}]}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[0],"output":{"b":0}}</div><div>{"inputs":[1,0],"output":[{"a":1},{"b":0}]}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[1],"output":{"b":1}}</div><div>{"inputs":[1,1],"output":[{"a":1},{"b":1}]}</div>
+<div>{"inputs":[1],"output":{"a":1}}</div><div>{"inputs":[2],"output":{"b":2}}</div><div>{"inputs":[1,2],"output":[{"a":1},{"b":2}]}</div>
+<div>{"inputs":[2],"output":{"a":2}}</div><div>{"inputs":[2],"output":{"b":2}}</div><div>{"inputs":[2,2],"output":[{"a":2},{"b":2}]}</div>
+<div>{"inputs":[3],"output":{"a":3}}</div><div>{"inputs":[2],"output":{"b":2}}</div><div>{"inputs":[3,2],"output":[{"a":3},{"b":2}]}</div>
+<div>{"inputs":[0],"output":{"a":0}}</div><div>{"inputs":[0],"output":{"b":0}}</div><div>{"inputs":[0,0],"output":[{"a":0},{"b":0}]}</div>
\ No newline at end of file
diff --git a/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture.tsx b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture.tsx
new file mode 100644
index 0000000000000..3afef5439bec3
--- /dev/null
+++ b/compiler/packages/babel-plugin-react-compiler/src/__tests__/fixtures/compiler/type-provider-store-capture.tsx
@@ -0,0 +1,35 @@
+import {useMemo} from 'react';
+import {typedArrayPush, ValidateMemoization} from 'shared-runtime';
+
+export function Component({a, b}) {
+  const item1 = useMemo(() => ({a}), [a]);
+  const item2 = useMemo(() => ({b}), [b]);
+  const items = useMemo(() => {
+    const items = [];
+    typedArrayPush(items, item1);
+    typedArrayPush(items, item2);
+    return items;
+  }, [item1, item2]);
+
+  return (
+    <>
+      <ValidateMemoization inputs={[a]} output={items[0]} />
+      <ValidateMemoization inputs={[b]} output={items[1]} />
+      <ValidateMemoization inputs={[a, b]} output={items} />
+    </>
+  );
+}
+
+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},
+  ],
+};
diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts
index de0184f0bddb1..417a657d28012 100644
--- a/compiler/packages/snap/src/compiler.ts
+++ b/compiler/packages/snap/src/compiler.ts
@@ -31,6 +31,7 @@ import path from 'path';
 import prettier from 'prettier';
 import SproutTodoFilter from './SproutTodoFilter';
 import {isExpectError} from './fixture-utils';
+import {makeSharedRuntimeTypeProvider} from './sprout/shared-runtime-type-provider';
 export function parseLanguage(source: string): 'flow' | 'typescript' {
   return source.indexOf('@flow') !== -1 ? 'flow' : 'typescript';
 }
@@ -38,6 +39,8 @@ export function parseLanguage(source: string): 'flow' | 'typescript' {
 function makePluginOptions(
   firstLine: string,
   parseConfigPragmaFn: typeof ParseConfigPragma,
+  EffectEnum: typeof Effect,
+  ValueKindEnum: typeof ValueKind,
 ): [PluginOptions, Array<{filename: string | null; event: LoggerEvent}>] {
   let gating = null;
   let enableEmitInstrumentForget = null;
@@ -212,35 +215,10 @@ function makePluginOptions(
   const options = {
     environment: {
       ...config,
-      customHooks: new Map([
-        [
-          'useFreeze',
-          {
-            valueKind: 'frozen' as ValueKind,
-            effectKind: 'freeze' as Effect,
-            transitiveMixedData: false,
-            noAlias: false,
-          },
-        ],
-        [
-          'useFragment',
-          {
-            valueKind: 'frozen' as ValueKind,
-            effectKind: 'freeze' as Effect,
-            transitiveMixedData: true,
-            noAlias: true,
-          },
-        ],
-        [
-          'useNoAlias',
-          {
-            valueKind: 'mutable' as ValueKind,
-            effectKind: 'read' as Effect,
-            transitiveMixedData: false,
-            noAlias: true,
-          },
-        ],
-      ]),
+      moduleTypeProvider: makeSharedRuntimeTypeProvider({
+        EffectEnum,
+        ValueKindEnum,
+      }),
       customMacros,
       enableEmitFreeze,
       enableEmitInstrumentForget,
@@ -383,6 +361,8 @@ export async function transformFixtureInput(
   parseConfigPragmaFn: typeof ParseConfigPragma,
   plugin: BabelCore.PluginObj,
   includeEvaluator: boolean,
+  EffectEnum: typeof Effect,
+  ValueKindEnum: typeof ValueKind,
 ): Promise<{kind: 'ok'; value: TransformResult} | {kind: 'err'; msg: string}> {
   // Extract the first line to quickly check for custom test directives
   const firstLine = input.substring(0, input.indexOf('\n'));
@@ -405,7 +385,12 @@ export async function transformFixtureInput(
   /**
    * Get Forget compiled code
    */
-  const [options, logs] = makePluginOptions(firstLine, parseConfigPragmaFn);
+  const [options, logs] = makePluginOptions(
+    firstLine,
+    parseConfigPragmaFn,
+    EffectEnum,
+    ValueKindEnum,
+  );
   const forgetResult = transformFromAstSync(inputAst, input, {
     filename: virtualFilepath,
     highlightCode: false,
diff --git a/compiler/packages/snap/src/constants.ts b/compiler/packages/snap/src/constants.ts
index bcc0b0ff1ca03..abee06c55be8a 100644
--- a/compiler/packages/snap/src/constants.ts
+++ b/compiler/packages/snap/src/constants.ts
@@ -17,6 +17,7 @@ export const COMPILER_PATH = path.join(
   'Babel',
   'BabelPlugin.js',
 );
+export const COMPILER_INDEX_PATH = path.join(process.cwd(), 'dist', 'index');
 export const LOGGER_PATH = path.join(
   process.cwd(),
   'dist',
diff --git a/compiler/packages/snap/src/runner-worker.ts b/compiler/packages/snap/src/runner-worker.ts
index 55c85b9466925..9447b2cddc52c 100644
--- a/compiler/packages/snap/src/runner-worker.ts
+++ b/compiler/packages/snap/src/runner-worker.ts
@@ -11,6 +11,7 @@ import type {parseConfigPragma as ParseConfigPragma} from 'babel-plugin-react-co
 import {TransformResult, transformFixtureInput} from './compiler';
 import {
   COMPILER_PATH,
+  COMPILER_INDEX_PATH,
   LOGGER_PATH,
   PARSE_CONFIG_PRAGMA_PATH,
 } from './constants';
@@ -60,6 +61,9 @@ async function compile(
     const {default: BabelPluginReactCompiler} = require(COMPILER_PATH) as {
       default: PluginObj;
     };
+    const {Effect: EffectEnum, ValueKind: ValueKindEnum} = require(
+      COMPILER_INDEX_PATH,
+    );
     const {toggleLogging} = require(LOGGER_PATH);
     const {parseConfigPragma} = require(PARSE_CONFIG_PRAGMA_PATH) as {
       parseConfigPragma: typeof ParseConfigPragma;
@@ -74,6 +78,8 @@ async function compile(
       parseConfigPragma,
       BabelPluginReactCompiler,
       includeEvaluator,
+      EffectEnum,
+      ValueKindEnum,
     );
 
     if (result.kind === 'err') {
diff --git a/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts b/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts
new file mode 100644
index 0000000000000..fb0877d11474f
--- /dev/null
+++ b/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts
@@ -0,0 +1,69 @@
+/**
+ * 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 type {Effect, ValueKind} from 'babel-plugin-react-compiler/src';
+import type {TypeConfig} from 'babel-plugin-react-compiler/src/HIR/TypeSchema';
+
+export function makeSharedRuntimeTypeProvider({
+  EffectEnum,
+  ValueKindEnum,
+}: {
+  EffectEnum: typeof Effect;
+  ValueKindEnum: typeof ValueKind;
+}) {
+  return function sharedRuntimeTypeProvider(
+    moduleName: string,
+  ): TypeConfig | null {
+    if (moduleName !== 'shared-runtime') {
+      return null;
+    }
+    return {
+      kind: 'object',
+      properties: {
+        default: {
+          kind: 'function',
+          calleeEffect: EffectEnum.Read,
+          positionalParams: [],
+          restParam: EffectEnum.Read,
+          returnType: {kind: 'type', name: 'Primitive'},
+          returnValueKind: ValueKindEnum.Primitive,
+        },
+        typedArrayPush: {
+          kind: 'function',
+          calleeEffect: EffectEnum.Read,
+          positionalParams: [EffectEnum.Store, EffectEnum.Capture],
+          restParam: EffectEnum.Capture,
+          returnType: {kind: 'type', name: 'Primitive'},
+          returnValueKind: ValueKindEnum.Primitive,
+        },
+        typedLog: {
+          kind: 'function',
+          calleeEffect: EffectEnum.Read,
+          positionalParams: [],
+          restParam: EffectEnum.Read,
+          returnType: {kind: 'type', name: 'Primitive'},
+          returnValueKind: ValueKindEnum.Primitive,
+        },
+        useFreeze: {
+          kind: 'hook',
+          returnType: {kind: 'type', name: 'Any'},
+        },
+        useFragment: {
+          kind: 'hook',
+          returnType: {kind: 'type', name: 'MixedReadonly'},
+          noAlias: true,
+        },
+        useNoAlias: {
+          kind: 'hook',
+          returnType: {kind: 'type', name: 'Any'},
+          returnValueKind: ValueKindEnum.Mutable,
+          noAlias: true,
+        },
+      },
+    };
+  };
+}
diff --git a/compiler/packages/snap/src/sprout/shared-runtime.ts b/compiler/packages/snap/src/sprout/shared-runtime.ts
index f15aaaaa4a2e5..bb1c65a6574ac 100644
--- a/compiler/packages/snap/src/sprout/shared-runtime.ts
+++ b/compiler/packages/snap/src/sprout/shared-runtime.ts
@@ -347,3 +347,12 @@ export function useFragment(..._args: Array<any>): object {
     b: {c: {d: 4}},
   };
 }
+
+export function typedArrayPush<T>(array: Array<T>, item: T): void {
+  array.push(item);
+}
+
+export function typedLog(...values: Array<any>): void {
+  console.log(...values);
+}
+export default typedLog;