From 7c99e6f7332bba8ef510249315e366d221ee4c70 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sebastian@calyptus.eu>
Date: Wed, 24 Aug 2022 21:22:48 -0400
Subject: [PATCH] Allow functions to be used as module references

---
 .../react-server/src/ReactFlightServer.js     | 93 +++++++++++--------
 1 file changed, 54 insertions(+), 39 deletions(-)

diff --git a/packages/react-server/src/ReactFlightServer.js b/packages/react-server/src/ReactFlightServer.js
index 814f8f1a4f9d2..1577687c0040e 100644
--- a/packages/react-server/src/ReactFlightServer.js
+++ b/packages/react-server/src/ReactFlightServer.js
@@ -195,6 +195,10 @@ function attemptResolveElement(
     );
   }
   if (typeof type === 'function') {
+    if (isModuleReference(type)) {
+      // This is a reference to a client component.
+      return [REACT_ELEMENT_TYPE, type, key, props];
+    }
     // This is a server-side component.
     return type(props);
   } else if (typeof type === 'string') {
@@ -295,6 +299,52 @@ function serializeByRefID(id: number): string {
   return '@' + id.toString(16);
 }
 
+function serializeModuleReference(
+  request: Request,
+  parent: {+[key: string | number]: ReactModel} | $ReadOnlyArray<ReactModel>,
+  key: string,
+  moduleReference: ModuleReference<any>,
+): string {
+  const moduleKey: ModuleKey = getModuleKey(moduleReference);
+  const writtenModules = request.writtenModules;
+  const existingId = writtenModules.get(moduleKey);
+  if (existingId !== undefined) {
+    if (parent[0] === REACT_ELEMENT_TYPE && key === '1') {
+      // If we're encoding the "type" of an element, we can refer
+      // to that by a lazy reference instead of directly since React
+      // knows how to deal with lazy values. This lets us suspend
+      // on this component rather than its parent until the code has
+      // loaded.
+      return serializeByRefID(existingId);
+    }
+    return serializeByValueID(existingId);
+  }
+  try {
+    const moduleMetaData: ModuleMetaData = resolveModuleMetaData(
+      request.bundlerConfig,
+      moduleReference,
+    );
+    request.pendingChunks++;
+    const moduleId = request.nextChunkId++;
+    emitModuleChunk(request, moduleId, moduleMetaData);
+    writtenModules.set(moduleKey, moduleId);
+    if (parent[0] === REACT_ELEMENT_TYPE && key === '1') {
+      // If we're encoding the "type" of an element, we can refer
+      // to that by a lazy reference instead of directly since React
+      // knows how to deal with lazy values. This lets us suspend
+      // on this component rather than its parent until the code has
+      // loaded.
+      return serializeByRefID(moduleId);
+    }
+    return serializeByValueID(moduleId);
+  } catch (x) {
+    request.pendingChunks++;
+    const errorId = request.nextChunkId++;
+    emitErrorChunk(request, errorId, x);
+    return serializeByValueID(errorId);
+  }
+}
+
 function escapeStringValue(value: string): string {
   if (value[0] === '$' || value[0] === '@') {
     // We need to escape $ or @ prefixed strings since we use those to encode
@@ -561,45 +611,7 @@ export function resolveModelToJSON(
 
   if (typeof value === 'object') {
     if (isModuleReference(value)) {
-      const moduleReference: ModuleReference<any> = (value: any);
-      const moduleKey: ModuleKey = getModuleKey(moduleReference);
-      const writtenModules = request.writtenModules;
-      const existingId = writtenModules.get(moduleKey);
-      if (existingId !== undefined) {
-        if (parent[0] === REACT_ELEMENT_TYPE && key === '1') {
-          // If we're encoding the "type" of an element, we can refer
-          // to that by a lazy reference instead of directly since React
-          // knows how to deal with lazy values. This lets us suspend
-          // on this component rather than its parent until the code has
-          // loaded.
-          return serializeByRefID(existingId);
-        }
-        return serializeByValueID(existingId);
-      }
-      try {
-        const moduleMetaData: ModuleMetaData = resolveModuleMetaData(
-          request.bundlerConfig,
-          moduleReference,
-        );
-        request.pendingChunks++;
-        const moduleId = request.nextChunkId++;
-        emitModuleChunk(request, moduleId, moduleMetaData);
-        writtenModules.set(moduleKey, moduleId);
-        if (parent[0] === REACT_ELEMENT_TYPE && key === '1') {
-          // If we're encoding the "type" of an element, we can refer
-          // to that by a lazy reference instead of directly since React
-          // knows how to deal with lazy values. This lets us suspend
-          // on this component rather than its parent until the code has
-          // loaded.
-          return serializeByRefID(moduleId);
-        }
-        return serializeByValueID(moduleId);
-      } catch (x) {
-        request.pendingChunks++;
-        const errorId = request.nextChunkId++;
-        emitErrorChunk(request, errorId, x);
-        return serializeByValueID(errorId);
-      }
+      return serializeModuleReference(request, parent, key, (value: any));
     } else if ((value: any).$$typeof === REACT_PROVIDER_TYPE) {
       const providerKey = ((value: any): ReactProviderType<any>)._context
         ._globalName;
@@ -673,6 +685,9 @@ export function resolveModelToJSON(
   }
 
   if (typeof value === 'function') {
+    if (isModuleReference(value)) {
+      return serializeModuleReference(request, parent, key, (value: any));
+    }
     if (/^on[A-Z]/.test(key)) {
       throw new Error(
         'Event handlers cannot be passed to client component props. ' +