From c3ff38349af8dad5dd1ac04527bf45f2373bf1aa Mon Sep 17 00:00:00 2001
From: Joe Savona <joesavona@fb.com>
Date: Wed, 21 Aug 2024 09:30:37 -0700
Subject: [PATCH] [compiler] Type provider infra for tests

[ghstack-poisoned]
---
 .../src/HIR/Environment.ts                    | 12 +++---
 compiler/packages/snap/src/compiler.ts        |  2 +
 .../sprout/shared-runtime-type-provider.ts    | 37 +++++++++++++++++++
 .../snap/src/sprout/shared-runtime.ts         |  8 ++++
 4 files changed, 53 insertions(+), 6 deletions(-)
 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 58e818205f..4062817d7f 100644
--- a/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
+++ b/compiler/packages/babel-plugin-react-compiler/src/HIR/Environment.ts
@@ -148,7 +148,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(ModuleTypeResolver).default(null),
 
   /**
    * A list of functions which the application compiles as macros, where
@@ -713,18 +713,18 @@ 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);
+      const unparsedModuleConfig = this.config.moduleTypeProvider(moduleName);
+      if (unparsedModuleConfig != null) {
+        const moduleConfig = TypeSchema.parse(unparsedModuleConfig);
         moduleType = installTypeConfig(
           this.#globals,
           this.#shapes,
-          moduleTypes,
+          moduleConfig,
         );
       } else {
         moduleType = null;
diff --git a/compiler/packages/snap/src/compiler.ts b/compiler/packages/snap/src/compiler.ts
index de0184f0bd..fb87768b69 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 {sharedRuntimeTypeProvider} from './sprout/shared-runtime-type-provider';
 export function parseLanguage(source: string): 'flow' | 'typescript' {
   return source.indexOf('@flow') !== -1 ? 'flow' : 'typescript';
 }
@@ -241,6 +242,7 @@ function makePluginOptions(
           },
         ],
       ]),
+      moduleTypeProvider: sharedRuntimeTypeProvider,
       customMacros,
       enableEmitFreeze,
       enableEmitInstrumentForget,
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 0000000000..2f7da68cba
--- /dev/null
+++ b/compiler/packages/snap/src/sprout/shared-runtime-type-provider.ts
@@ -0,0 +1,37 @@
+/**
+ * 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 TypeConfig} from 'babel-plugin-react-compiler/src/HIR/TypeSchema';
+
+export function sharedRuntimeTypeProvider(
+  moduleName: string,
+): TypeConfig | null {
+  if (moduleName !== './shared-runtime') {
+    return null;
+  }
+  return {
+    kind: 'object',
+    properties: {
+      typedArrayPush: {
+        kind: 'function',
+        calleeEffect: 'read',
+        positionalParams: ['store', 'capture'],
+        restParam: 'capture',
+        returnType: {kind: 'type', name: 'Primitive'},
+        returnValueKind: 'primitive',
+      },
+      typedLog: {
+        kind: 'function',
+        calleeEffect: 'read',
+        positionalParams: [],
+        restParam: 'read',
+        returnType: {kind: 'type', name: 'Primitive'},
+        returnValueKind: 'primitive',
+      },
+    },
+  };
+}
diff --git a/compiler/packages/snap/src/sprout/shared-runtime.ts b/compiler/packages/snap/src/sprout/shared-runtime.ts
index f15aaaaa4a..69af42e52a 100644
--- a/compiler/packages/snap/src/sprout/shared-runtime.ts
+++ b/compiler/packages/snap/src/sprout/shared-runtime.ts
@@ -347,3 +347,11 @@ 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<string>): void {
+  console.log(...values);
+}