[Fiber] support hydration when rendering Suspense anywhere (#32224)

follow up to #32163

This continues the work of making Suspense workable anywhere in a
react-dom tree. See the prior PRs for how we handle server rendering and
client rendering. In this change we update the hydration implementation
to be able to locate expected nodes. In particular this means hydration
understands now that the default hydration context is the document body
when the container is above the body.

One case that is unique to hydration is clearing Suspense boundaries.
When hydration fails or when the server instructs the client to recover
an errored boundary it's possible that the html, head, and body tags in
the initial document were written from a fallback or a different primary
content on the server and need to be replaced by the client render.
However these tags (and in the case of head, their content) won't be
inside the comment nodes that identify the bounds of the Suspense
boundary. And when client rendering you may not even render the same
singletons that were server rendered. So when server rendering a
boudnary which contributes to the preamble (the html, head, and body tag
openings plus the head contents) we emit a special marker comment just
before closing the boundary out. This marker encodes which parts of the
preamble this boundary owned. If we need to clear the suspense boundary
on the client we read this marker and use it to reset the appropriate
singleton state.

Author: gnoff

packages/react-dom-bindings/src/client/ReactFiberConfigDOM.js

@@ -207,6 +207,9 @@ const SUSPENSE_START_DATA = '$';
 const SUSPENSE_END_DATA = '/$';
 const SUSPENSE_PENDING_START_DATA = '$?';
 const SUSPENSE_FALLBACK_START_DATA = '$!';
+const PREAMBLE_CONTRIBUTION_HTML = 0b001;
+const PREAMBLE_CONTRIBUTION_BODY = 0b010;
+const PREAMBLE_CONTRIBUTION_HEAD = 0b100;
 const FORM_STATE_IS_MATCHING = 'F!';
 const FORM_STATE_IS_NOT_MATCHING = 'F';
 
@@ -963,6 +966,7 @@ export function clearSuspenseBoundary(
   suspenseInstance: SuspenseInstance,
 ): void {
   let node: Node = suspenseInstance;
+  let possiblePreambleContribution: number = 0;
   // Delete all nodes within this suspense boundary.
   // There might be nested nodes so we need to keep track of how
   // deep we are and only break out when we're back on top.
@@ -973,6 +977,36 @@ export function clearSuspenseBoundary(
     if (nextNode && nextNode.nodeType === COMMENT_NODE) {
       const data = ((nextNode: any).data: string);
       if (data === SUSPENSE_END_DATA) {
+        if (
+          // represents 3 bits where at least one bit is set (1-7)
+          possiblePreambleContribution > 0 &&
+          possiblePreambleContribution < 8
+        ) {
+          const code = possiblePreambleContribution;
+          // It's not normally possible to insert a comment immediately preceding Suspense boundary
+          // closing comment marker so we can infer that if the comment preceding starts with "1" through "7"
+          // then it is in fact a preamble contribution marker comment. We do this value test to avoid the case
+          // where the Suspense boundary is empty and the preceding comment marker is the Suspense boundary
+          // opening marker or the closing marker of an inner boundary. In those cases the first character won't
+          // have the requisite value to be interpreted as a Preamble contribution
+          const ownerDocument = parentInstance.ownerDocument;
+          if (code & PREAMBLE_CONTRIBUTION_HTML) {
+            const documentElement: Element =
+              (ownerDocument.documentElement: any);
+            releaseSingletonInstance(documentElement);
+          }
+          if (code & PREAMBLE_CONTRIBUTION_BODY) {
+            const body: Element = (ownerDocument.body: any);
+            releaseSingletonInstance(body);
+          }
+          if (code & PREAMBLE_CONTRIBUTION_HEAD) {
+            const head: Element = (ownerDocument.head: any);
+            releaseSingletonInstance(head);
+            // We need to clear the head because this is the only singleton that can have children that
+            // were part of this boundary but are not inside this boundary.
+            clearHead(head);
+          }
+        }
         if (depth === 0) {
           parentInstance.removeChild(nextNode);
           // Retry if any event replaying was blocked on this.
@@ -987,7 +1021,11 @@ export function clearSuspenseBoundary(
         data === SUSPENSE_FALLBACK_START_DATA
       ) {
         depth++;
+      } else {
+        possiblePreambleContribution = data.charCodeAt(0) - 48;
       }
+    } else {
+      possiblePreambleContribution = 0;
     }
     // $FlowFixMe[incompatible-type] we bail out when we get a null
     node = nextNode;
@@ -1501,7 +1539,7 @@ function clearContainerSparingly(container: Node) {
       case 'STYLE': {
         continue;
       }
-      // Stylesheet tags are retained because tehy may likely come from 3rd party scripts and extensions
+      // Stylesheet tags are retained because they may likely come from 3rd party scripts and extensions
       case 'LINK': {
         if (((node: any): HTMLLinkElement).rel.toLowerCase() === 'stylesheet') {
           continue;
@@ -1513,6 +1551,27 @@ function clearContainerSparingly(container: Node) {
   return;
 }
 
+function clearHead(head: Element): void {
+  let node = head.firstChild;
+  while (node) {
+    const nextNode = node.nextSibling;
+    const nodeName = node.nodeName;
+    if (
+      isMarkedHoistable(node) ||
+      nodeName === 'SCRIPT' ||
+      nodeName === 'STYLE' ||
+      (nodeName === 'LINK' &&
+        ((node: any): HTMLLinkElement).rel.toLowerCase() === 'stylesheet')
+    ) {
+      // retain these nodes
+    } else {
+      head.removeChild(node);
+    }
+    node = nextNode;
+  }
+  return;
+}
+
 // Making this so we can eventually move all of the instance caching to the commit phase.
 // Currently this is only used to associate fiber and props to instances for hydrating
 // HostSingletons. The reason we need it here is we only want to make this binding on commit
@@ -1874,7 +1933,20 @@ export function getFirstHydratableChild(
 export function getFirstHydratableChildWithinContainer(
   parentContainer: Container,
 ): null | HydratableInstance {
-  return getNextHydratable(parentContainer.firstChild);
+  let parentElement: Element;
+  switch (parentContainer.nodeType) {
+    case DOCUMENT_NODE:
+      parentElement = (parentContainer: any).body;
+      break;
+    default: {
+      if (parentContainer.nodeName === 'HTML') {
+        parentElement = (parentContainer: any).ownerDocument.body;
+      } else {
+        parentElement = (parentContainer: any);
+      }
+    }
+  }
+  return getNextHydratable(parentElement.firstChild);
 }
 
 export function getFirstHydratableChildWithinSuspenseInstance(
@@ -1883,6 +1955,40 @@ export function getFirstHydratableChildWithinSuspenseInstance(
   return getNextHydratable(parentInstance.nextSibling);
 }
 
+// If it were possible to have more than one scope singleton in a DOM tree
+// we would need to model this as a stack but since you can only have one <head>
+// and head is the only singleton that is a scope in DOM we can get away with
+// tracking this as a single value.
+let previousHydratableOnEnteringScopedSingleton: null | HydratableInstance =
+  null;
+
+export function getFirstHydratableChildWithinSingleton(
+  type: string,
+  singletonInstance: Instance,
+  currentHydratableInstance: null | HydratableInstance,
+): null | HydratableInstance {
+  if (isSingletonScope(type)) {
+    previousHydratableOnEnteringScopedSingleton = currentHydratableInstance;
+    return getNextHydratable(singletonInstance.firstChild);
+  } else {
+    return currentHydratableInstance;
+  }
+}
+
+export function getNextHydratableSiblingAfterSingleton(
+  type: string,
+  currentHydratableInstance: null | HydratableInstance,
+): null | HydratableInstance {
+  if (isSingletonScope(type)) {
+    const previousHydratableInstance =
+      previousHydratableOnEnteringScopedSingleton;
+    previousHydratableOnEnteringScopedSingleton = null;
+    return previousHydratableInstance;
+  } else {
+    return currentHydratableInstance;
+  }
+}
+
 export function describeHydratableInstanceForDevWarnings(
   instance: HydratableInstance,
 ): string | {type: string, props: $ReadOnly<Props>} {

packages/react-dom-bindings/src/server/ReactFizzConfigDOM.js

@@ -684,16 +684,23 @@ export function completeResumableState(resumableState: ResumableState): void {
   resumableState.bootstrapModules = undefined;
 }
 
+const NoContribution /*     */ = 0b000;
+const HTMLContribution /*   */ = 0b001;
+const BodyContribution /*   */ = 0b010;
+const HeadContribution /*   */ = 0b100;
+
 export type PreambleState = {
   htmlChunks: null | Array<Chunk | PrecomputedChunk>,
   headChunks: null | Array<Chunk | PrecomputedChunk>,
   bodyChunks: null | Array<Chunk | PrecomputedChunk>,
+  contribution: number,
 };
 export function createPreambleState(): PreambleState {
   return {
     htmlChunks: null,
     headChunks: null,
     bodyChunks: null,
+    contribution: NoContribution,
   };
 }
 
@@ -3227,7 +3234,7 @@ function pushStartHead(
       throw new Error(`The ${'`<head>`'} tag may only be rendered once.`);
     }
     preamble.headChunks = [];
-    return pushStartGenericElement(preamble.headChunks, props, 'head');
+    return pushStartSingletonElement(preamble.headChunks, props, 'head');
   } else {
     // This <head> is deep and is likely just an error. we emit it inline though.
     // Validation should warn that this tag is the the wrong spot.
@@ -3251,7 +3258,7 @@ function pushStartBody(
     }
 
     preamble.bodyChunks = [];
-    return pushStartGenericElement(preamble.bodyChunks, props, 'body');
+    return pushStartSingletonElement(preamble.bodyChunks, props, 'body');
   } else {
     // This <head> is deep and is likely just an error. we emit it inline though.
     // Validation should warn that this tag is the the wrong spot.
@@ -3275,7 +3282,7 @@ function pushStartHtml(
     }
 
     preamble.htmlChunks = [DOCTYPE];
-    return pushStartGenericElement(preamble.htmlChunks, props, 'html');
+    return pushStartSingletonElement(preamble.htmlChunks, props, 'html');
   } else {
     // This <html> is deep and is likely just an error. we emit it inline though.
     // Validation should warn that this tag is the the wrong spot.
@@ -3416,6 +3423,43 @@ function pushScriptImpl(
   return null;
 }
 
+// This is a fork of pushStartGenericElement because we don't ever want to do
+// the children as strign optimization on that path when rendering singletons.
+// When we eliminate that special path we can delete this fork and unify it again
+function pushStartSingletonElement(
+  target: Array<Chunk | PrecomputedChunk>,
+  props: Object,
+  tag: string,
+): ReactNodeList {
+  target.push(startChunkForTag(tag));
+
+  let children = null;
+  let innerHTML = null;
+  for (const propKey in props) {
+    if (hasOwnProperty.call(props, propKey)) {
+      const propValue = props[propKey];
+      if (propValue == null) {
+        continue;
+      }
+      switch (propKey) {
+        case 'children':
+          children = propValue;
+          break;
+        case 'dangerouslySetInnerHTML':
+          innerHTML = propValue;
+          break;
+        default:
+          pushAttribute(target, propKey, propValue);
+          break;
+      }
+    }
+  }
+
+  target.push(endOfStartTag);
+  pushInnerHTML(target, innerHTML, children);
+  return children;
+}
+
 function pushStartGenericElement(
   target: Array<Chunk | PrecomputedChunk>,
   props: Object,
@@ -3907,14 +3951,17 @@ export function hoistPreambleState(
   preambleState: PreambleState,
 ) {
   const rootPreamble = renderState.preamble;
-  if (rootPreamble.htmlChunks === null) {
+  if (rootPreamble.htmlChunks === null && preambleState.htmlChunks) {
     rootPreamble.htmlChunks = preambleState.htmlChunks;
+    preambleState.contribution |= HTMLContribution;
   }
-  if (rootPreamble.headChunks === null) {
+  if (rootPreamble.headChunks === null && preambleState.headChunks) {
     rootPreamble.headChunks = preambleState.headChunks;
+    preambleState.contribution |= HeadContribution;
   }
-  if (rootPreamble.bodyChunks === null) {
+  if (rootPreamble.bodyChunks === null && preambleState.bodyChunks) {
     rootPreamble.bodyChunks = preambleState.bodyChunks;
+    preambleState.contribution |= BodyContribution;
   }
 }
 
@@ -4091,7 +4138,11 @@ export function writeStartClientRenderedSuspenseBoundary(
 export function writeEndCompletedSuspenseBoundary(
   destination: Destination,
   renderState: RenderState,
+  preambleState: null | PreambleState,
 ): boolean {
+  if (preambleState) {
+    writePreambleContribution(destination, preambleState);
+  }
   return writeChunkAndReturn(destination, endSuspenseBoundary);
 }
 export function writeEndPendingSuspenseBoundary(
@@ -4103,10 +4154,31 @@ export function writeEndPendingSuspenseBoundary(
 export function writeEndClientRenderedSuspenseBoundary(
   destination: Destination,
   renderState: RenderState,
+  preambleState: null | PreambleState,
 ): boolean {
+  if (preambleState) {
+    writePreambleContribution(destination, preambleState);
+  }
   return writeChunkAndReturn(destination, endSuspenseBoundary);
 }
 
+const boundaryPreambleContributionChunkStart = stringToPrecomputedChunk('<!--');
+const boundaryPreambleContributionChunkEnd = stringToPrecomputedChunk('-->');
+
+function writePreambleContribution(
+  destination: Destination,
+  preambleState: PreambleState,
+) {
+  const contribution = preambleState.contribution;
+  if (contribution !== NoContribution) {
+    writeChunk(destination, boundaryPreambleContributionChunkStart);
+    // This is a number type so we can do the fast path without coercion checking
+    // eslint-disable-next-line react-internal/safe-string-coercion
+    writeChunk(destination, stringToChunk('' + contribution));
+    writeChunk(destination, boundaryPreambleContributionChunkEnd);
+  }
+}
+
 const startSegmentHTML = stringToPrecomputedChunk('<div hidden id="');
 const startSegmentHTML2 = stringToPrecomputedChunk('">');
 const endSegmentHTML = stringToPrecomputedChunk('</div>');

packages/react-dom-bindings/src/server/ReactFizzConfigDOMLegacy.js

@@ -244,20 +244,30 @@ export function writeStartClientRenderedSuspenseBoundary(
 export function writeEndCompletedSuspenseBoundary(
   destination: Destination,
   renderState: RenderState,
+  preambleState: null | PreambleState,
 ): boolean {
   if (renderState.generateStaticMarkup) {
     return true;
   }
-  return writeEndCompletedSuspenseBoundaryImpl(destination, renderState);
+  return writeEndCompletedSuspenseBoundaryImpl(
+    destination,
+    renderState,
+    preambleState,
+  );
 }
 export function writeEndClientRenderedSuspenseBoundary(
   destination: Destination,
   renderState: RenderState,
+  preambleState: null | PreambleState,
 ): boolean {
   if (renderState.generateStaticMarkup) {
     return true;
   }
-  return writeEndClientRenderedSuspenseBoundaryImpl(destination, renderState);
+  return writeEndClientRenderedSuspenseBoundaryImpl(
+    destination,
+    renderState,
+    preambleState,
+  );
 }
 
 export type TransitionStatus = FormStatus;