[Flight] Ensure blocked debug info is handled properly

packages/react-client/src/ReactFlightClient.js

@@ -512,6 +512,27 @@ function wakeChunk<T>(
       fulfillReference(listener, value, chunk);
     }
   }
+
+  if (__DEV__ && chunk.status === INITIALIZED) {
+    const resolvedValue = resolveLazy(value);
+    if (isReactElementOrArrayLike(resolvedValue) || isLazy(resolvedValue)) {
+      const debugInfo = chunk._debugInfo.splice(0);
+      if (resolvedValue._debugInfo) {
+        // $FlowFixMe[method-unbinding]
+        resolvedValue._debugInfo.push.apply(
+          resolvedValue._debugInfo,
+          debugInfo,
+        );
+      } else {
+        Object.defineProperty(resolvedValue, '_debugInfo', {
+          configurable: false,
+          enumerable: false,
+          writable: true,
+          value: debugInfo,
+        });
+      }
+    }
+  }
 }
 
 function rejectChunk(
@@ -649,7 +670,6 @@ function triggerErrorOnChunk<T>(
       }
       try {
         initializeDebugChunk(response, chunk);
-        chunk._debugChunk = null;
         if (initializingHandler !== null) {
           if (initializingHandler.errored) {
             // Ignore error parsing debug info, we'll report the original error instead.
@@ -932,9 +952,9 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
   }
 
   if (__DEV__) {
-    // Lazily initialize any debug info and block the initializing chunk on any unresolved entries.
+    // Initialize any debug info and block the initializing chunk on any
+    // unresolved entries.
     initializeDebugChunk(response, chunk);
-    chunk._debugChunk = null;
   }
 
   try {
@@ -960,6 +980,22 @@ function initializeModelChunk<T>(chunk: ResolvedModelChunk<T>): void {
         return;
       }
     }
+
+    if (__DEV__ && isReactElementOrArrayLike(value)) {
+      const debugInfo = chunk._debugInfo.splice(0);
+      if (value._debugInfo) {
+        // $FlowFixMe[method-unbinding]
+        value._debugInfo.push.apply(value._debugInfo, debugInfo);
+      } else {
+        Object.defineProperty(value, '_debugInfo', {
+          configurable: false,
+          enumerable: false,
+          writable: true,
+          value: debugInfo,
+        });
+      }
+    }
+
     const initializedChunk: InitializedChunk<T> = (chunk: any);
     initializedChunk.status = INITIALIZED;
     initializedChunk.value = value;
@@ -1053,11 +1089,7 @@ function getTaskName(type: mixed): string {
     // the client. There should only be one for any given owner chain.
     return '"use client"';
   }
-  if (
-    typeof type === 'object' &&
-    type !== null &&
-    type.$$typeof === REACT_LAZY_TYPE
-  ) {
+  if (isLazy(type)) {
     if (type._init === readChunk) {
       // This is a lazy node created by Flight. It is probably a client reference.
       // We use the "use client" string to indicate that this is the boundary into
@@ -1079,7 +1111,7 @@ function getTaskName(type: mixed): string {
 function initializeElement(
   response: Response,
   element: any,
-  lazyType: null | LazyComponent<
+  lazyNode: null | LazyComponent<
     React$Element<any>,
     SomeChunk<React$Element<any>>,
   >,
@@ -1151,15 +1183,31 @@ function initializeElement(
     initializeFakeStack(response, owner);
   }
 
-  // In case the JSX runtime has validated the lazy type as a static child, we
-  // need to transfer this information to the element.
-  if (
-    lazyType &&
-    lazyType._store &&
-    lazyType._store.validated &&
-    !element._store.validated
-  ) {
-    element._store.validated = lazyType._store.validated;
+  if (lazyNode !== null) {
+    // In case the JSX runtime has validated the lazy type as a static child, we
+    // need to transfer this information to the element.
+    if (
+      lazyNode._store &&
+      lazyNode._store.validated &&
+      !element._store.validated
+    ) {
+      element._store.validated = lazyNode._store.validated;
+    }
+
+    if (lazyNode._payload.status === INITIALIZED && lazyNode._debugInfo) {
+      const debugInfo = lazyNode._debugInfo.splice(0);
+      if (element._debugInfo) {
+        // $FlowFixMe[method-unbinding]
+        element._debugInfo.push.apply(element._debugInfo, debugInfo);
+      } else {
+        Object.defineProperty(element, '_debugInfo', {
+          configurable: false,
+          enumerable: false,
+          writable: true,
+          value: debugInfo,
+        });
+      }
+    }
   }
 
   // TODO: We should be freezing the element but currently, we might write into
@@ -1279,13 +1327,13 @@ function createElement(
         createBlockedChunk(response);
       handler.value = element;
       handler.chunk = blockedChunk;
-      const lazyType = createLazyChunkWrapper(blockedChunk, validated);
+      const lazyNode = createLazyChunkWrapper(blockedChunk, validated);
       if (__DEV__) {
         // After we have initialized any blocked references, initialize stack etc.
-        const init = initializeElement.bind(null, response, element, lazyType);
+        const init = initializeElement.bind(null, response, element, lazyNode);
         blockedChunk.then(init, init);
       }
-      return lazyType;
+      return lazyNode;
     }
   }
   if (__DEV__) {
@@ -1337,11 +1385,7 @@ function fulfillReference(
   const {response, handler, parentObject, key, map, path} = reference;
 
   for (let i = 1; i < path.length; i++) {
-    while (
-      typeof value === 'object' &&
-      value !== null &&
-      value.$$typeof === REACT_LAZY_TYPE
-    ) {
+    while (isLazy(value)) {
       // We never expect to see a Lazy node on this path because we encode those as
       // separate models. This must mean that we have inserted an extra lazy node
       // e.g. to replace a blocked element. We must instead look for it inside.
@@ -1413,11 +1457,7 @@ function fulfillReference(
     value = value[path[i]];
   }
 
-  while (
-    typeof value === 'object' &&
-    value !== null &&
-    value.$$typeof === REACT_LAZY_TYPE
-  ) {
+  while (isLazy(value)) {
     // If what we're referencing is a Lazy it must be because we inserted one as a virtual node
     // while it was blocked by other data. If it's no longer blocked, we can unwrap it.
     const referencedChunk: SomeChunk<any> = value._payload;
@@ -1466,7 +1506,7 @@ function fulfillReference(
     const element: any = handler.value;
     switch (key) {
       case '3':
-        transferReferencedDebugInfo(handler.chunk, fulfilledChunk, mappedValue);
+        transferReferencedDebugInfo(handler.chunk, fulfilledChunk);
         element.props = mappedValue;
         break;
       case '4':
@@ -1482,11 +1522,11 @@ function fulfillReference(
         }
         break;
       default:
-        transferReferencedDebugInfo(handler.chunk, fulfilledChunk, mappedValue);
+        transferReferencedDebugInfo(handler.chunk, fulfilledChunk);
         break;
     }
   } else if (__DEV__ && !reference.isDebug) {
-    transferReferencedDebugInfo(handler.chunk, fulfilledChunk, mappedValue);
+    transferReferencedDebugInfo(handler.chunk, fulfilledChunk);
   }
 
   handler.deps--;
@@ -1808,47 +1848,57 @@ function loadServerReference<A: Iterable<any>, T>(
   return (null: any);
 }
 
+function isReactElementOrArrayLike(
+  value: any,
+  // eslint-disable-next-line no-undef
+): value is {_debugInfo: null | ReactDebugInfo, ...} {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    (isArray(value) ||
+      typeof value[ASYNC_ITERATOR] === 'function' ||
+      value.$$typeof === REACT_ELEMENT_TYPE)
+  );
+}
+
+function isLazy(
+  value: any,
+  // eslint-disable-next-line no-undef
+): implies value is LazyComponent<
+  React$Element<any>,
+  SomeChunk<React$Element<any>>,
+> {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    value.$$typeof === REACT_LAZY_TYPE
+  );
+}
+
+function resolveLazy(value: mixed): mixed {
+  while (isLazy(value)) {
+    const payload: SomeChunk<any> = value._payload;
+    if (payload.status === INITIALIZED) {
+      value = payload.value;
+      continue;
+    }
+    break;
+  }
+
+  return value;
+}
+
 function transferReferencedDebugInfo(
   parentChunk: null | SomeChunk<any>,
   referencedChunk: SomeChunk<any>,
-  referencedValue: mixed,
 ): void {
   if (__DEV__) {
-    const referencedDebugInfo = referencedChunk._debugInfo;
-    // If we have a direct reference to an object that was rendered by a synchronous
-    // server component, it might have some debug info about how it was rendered.
-    // We forward this to the underlying object. This might be a React Element or
-    // an Array fragment.
-    // If this was a string / number return value we lose the debug info. We choose
-    // that tradeoff to allow sync server components to return plain values and not
-    // use them as React Nodes necessarily. We could otherwise wrap them in a Lazy.
-    if (
-      typeof referencedValue === 'object' &&
-      referencedValue !== null &&
-      (isArray(referencedValue) ||
-        typeof referencedValue[ASYNC_ITERATOR] === 'function' ||
-        referencedValue.$$typeof === REACT_ELEMENT_TYPE)
-    ) {
-      // We should maybe use a unique symbol for arrays but this is a React owned array.
-      // $FlowFixMe[prop-missing]: This should be added to elements.
-      const existingDebugInfo: ?ReactDebugInfo =
-        (referencedValue._debugInfo: any);
-      if (existingDebugInfo == null) {
-        Object.defineProperty((referencedValue: any), '_debugInfo', {
-          configurable: false,
-          enumerable: false,
-          writable: true,
-          value: referencedDebugInfo.slice(0), // Clone so that pushing later isn't going into the original
-        });
-      } else {
-        // $FlowFixMe[method-unbinding]
-        existingDebugInfo.push.apply(existingDebugInfo, referencedDebugInfo);
-      }
-    }
-    // We also add the debug info to the initializing chunk since the resolution of that promise is
-    // also blocked by the referenced debug info. By adding it to both we can track it even if the array/element
-    // is extracted, or if the root is rendered as is.
+    // We add the debug info to the initializing chunk since the resolution of
+    // that promise is also blocked by the referenced debug info. By adding it
+    // to both we can track it even if the array/element/lazy is extracted, or
+    // if the root is rendered as is.
     if (parentChunk !== null) {
+      const referencedDebugInfo = referencedChunk._debugInfo;
       const parentDebugInfo = parentChunk._debugInfo;
       for (let i = 0; i < referencedDebugInfo.length; ++i) {
         const debugInfoEntry = referencedDebugInfo[i];
@@ -1892,11 +1942,7 @@ function getOutlinedModel<T>(
     case INITIALIZED:
       let value = chunk.value;
       for (let i = 1; i < path.length; i++) {
-        while (
-          typeof value === 'object' &&
-          value !== null &&
-          value.$$typeof === REACT_LAZY_TYPE
-        ) {
+        while (isLazy(value)) {
           const referencedChunk: SomeChunk<any> = value._payload;
           switch (referencedChunk.status) {
             case RESOLVED_MODEL:
@@ -1966,11 +2012,7 @@ function getOutlinedModel<T>(
         value = value[path[i]];
       }
 
-      while (
-        typeof value === 'object' &&
-        value !== null &&
-        value.$$typeof === REACT_LAZY_TYPE
-      ) {
+      while (isLazy(value)) {
         // If what we're referencing is a Lazy it must be because we inserted one as a virtual node
         // while it was blocked by other data. If it's no longer blocked, we can unwrap it.
         const referencedChunk: SomeChunk<any> = value._payload;
@@ -1999,7 +2041,7 @@ function getOutlinedModel<T>(
         // If we're resolving the "owner" or "stack" slot of an Element array, we don't call
         // transferReferencedDebugInfo because this reference is to a debug chunk.
       } else {
-        transferReferencedDebugInfo(initializingChunk, chunk, chunkValue);
+        transferReferencedDebugInfo(initializingChunk, chunk);
       }
       return chunkValue;
     case PENDING:
@@ -2714,9 +2756,35 @@ function resolveChunkDebugInfo(
   chunk: SomeChunk<any>,
 ): void {
   if (__DEV__ && enableAsyncDebugInfo) {
-    // Push the currently resolving chunk's debug info representing the stream on the Promise
-    // that was waiting on the stream.
-    chunk._debugInfo.push({awaited: streamState._debugInfo});
+    // Add the currently resolving chunk's debug info representing the stream
+    // to the Promise that was waiting on the stream, or its underlying value.
+    const debugInfoEntry: ReactAsyncInfo = {awaited: streamState._debugInfo};
+
+    const addDebugInfo = () => {
+      const value = resolveLazy(chunk.value);
+      if (isReactElementOrArrayLike(value)) {
+        const debugInfo: ReactDebugInfo = [debugInfoEntry];
+        if (value._debugInfo) {
+          // $FlowFixMe[method-unbinding]
+          value._debugInfo.push.apply(value._debugInfo, debugInfo);
+        } else {
+          Object.defineProperty(value, '_debugInfo', {
+            configurable: false,
+            enumerable: false,
+            writable: true,
+            value: debugInfo,
+          });
+        }
+      } else {
+        chunk._debugInfo.push(debugInfoEntry);
+      }
+    };
+
+    if (chunk.status === PENDING || chunk.status === BLOCKED) {
+      chunk.then(addDebugInfo, addDebugInfo);
+    } else {
+      addDebugInfo();
+    }
   }
 }
 
@@ -2909,7 +2977,8 @@ function resolveStream<T: ReadableStream | $AsyncIterable<any, any, void>>(
   const resolveListeners = chunk.value;
 
   if (__DEV__) {
-    // Lazily initialize any debug info and block the initializing chunk on any unresolved entries.
+    // Initialize any debug info and block the initializing chunk on any
+    // unresolved entries.
     if (chunk._debugChunk != null) {
       const prevHandler = initializingHandler;
       const prevChunk = initializingChunk;
@@ -2923,7 +2992,6 @@ function resolveStream<T: ReadableStream | $AsyncIterable<any, any, void>>(
       }
       try {
         initializeDebugChunk(response, chunk);
-        chunk._debugChunk = null;
         if (initializingHandler !== null) {
           if (initializingHandler.errored) {
             // Ignore error parsing debug info, we'll report the original error instead.