[Fiber] support hydration when rendering Suspense anywhere

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;