Hide/unhide the content of dehydrated suspense boundaries if they resuspend

Author: sebmarkbage

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

@@ -1127,6 +1127,61 @@ export function clearSuspenseBoundaryFromContainer(
   retryIfBlockedOn(container);
 }
 
+function hideOrUnhideSuspenseBoundary(
+  suspenseInstance: SuspenseInstance,
+  isHidden: boolean,
+) {
+  let node: Node = suspenseInstance;
+  // Unhide all nodes within this suspense boundary.
+  let depth = 0;
+  do {
+    const nextNode = node.nextSibling;
+    if (node.nodeType === ELEMENT_NODE) {
+      const instance = ((node: any): HTMLElement & {_stashedDisplay?: string});
+      if (isHidden) {
+        instance._stashedDisplay = instance.style.display;
+        instance.style.display = 'none';
+      } else {
+        instance.style.display = instance._stashedDisplay || '';
+        if (instance.getAttribute('style') === '') {
+          instance.removeAttribute('style');
+        }
+      }
+    } else if (node.nodeType === TEXT_NODE) {
+      const textNode = ((node: any): Text & {_stashedText?: string});
+      if (isHidden) {
+        textNode._stashedText = textNode.nodeValue;
+        textNode.nodeValue = '';
+      } else {
+        textNode.nodeValue = textNode._stashedText || '';
+      }
+    }
+    if (nextNode && nextNode.nodeType === COMMENT_NODE) {
+      const data = ((nextNode: any).data: string);
+      if (data === SUSPENSE_END_DATA) {
+        if (depth === 0) {
+          return;
+        } else {
+          depth--;
+        }
+      } else if (
+        data === SUSPENSE_START_DATA ||
+        data === SUSPENSE_PENDING_START_DATA ||
+        data === SUSPENSE_FALLBACK_START_DATA
+      ) {
+        depth++;
+      }
+      // TODO: Should we hide preamble contribution in this case?
+    }
+    // $FlowFixMe[incompatible-type] we bail out when we get a null
+    node = nextNode;
+  } while (node);
+}
+
+export function hideSuspenseBoundary(suspenseInstance: SuspenseInstance): void {
+  hideOrUnhideSuspenseBoundary(suspenseInstance, true);
+}
+
 export function hideInstance(instance: Instance): void {
   // TODO: Does this work for all element types? What about MathML? Should we
   // pass host context to this method?
@@ -1144,6 +1199,12 @@ export function hideTextInstance(textInstance: TextInstance): void {
   textInstance.nodeValue = '';
 }
 
+export function unhideSuspenseBoundary(
+  suspenseInstance: SuspenseInstance,
+): void {
+  hideOrUnhideSuspenseBoundary(suspenseInstance, false);
+}
+
 export function unhideInstance(instance: Instance, props: Props): void {
   instance = ((instance: any): HTMLElement);
   const styleProp = props[STYLE];

packages/react-dom/src/__tests__/ReactDOMServerPartialHydration-test.internal.js

@@ -3945,4 +3945,94 @@ describe('ReactDOMServerPartialHydration', () => {
       "onRecoverableError: Hydration failed because the server rendered text didn't match the client.",
     ]);
   });
+
+  it('hides a dehydrated suspense boundary if the parent resuspends', async () => {
+    let suspend = false;
+    let resolve;
+    const promise = new Promise(resolvePromise => (resolve = resolvePromise));
+    const ref = React.createRef();
+
+    function Child({text}) {
+      if (suspend) {
+        throw promise;
+      } else {
+        return text;
+      }
+    }
+
+    function Sibling({resuspend}) {
+      if (suspend && resuspend) {
+        throw promise;
+      } else {
+        return null;
+      }
+    }
+
+    function Component({text}) {
+      return (
+        <Suspense>
+          <Child text={text} />
+          <span ref={ref}>World</span>
+        </Suspense>
+      );
+    }
+
+    function App({text, resuspend}) {
+      const memoized = React.useMemo(() => <Component text={text} />, [text]);
+      return (
+        <div>
+          <Suspense fallback="Loading...">
+            {memoized}
+            <Sibling resuspend={resuspend} />
+          </Suspense>
+        </div>
+      );
+    }
+
+    suspend = false;
+    const finalHTML = ReactDOMServer.renderToString(<App text="Hello" />);
+    const container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    suspend = true;
+    const root = ReactDOMClient.hydrateRoot(container, <App text="Hello" />, {
+      onRecoverableError(error) {
+        Scheduler.log('onRecoverableError: ' + normalizeError(error.message));
+        if (error.cause) {
+          Scheduler.log('Cause: ' + normalizeError(error.cause.message));
+        }
+      },
+    });
+    await waitForAll([]);
+
+    expect(ref.current).toBe(null); // Still dehydrated
+    const span = container.getElementsByTagName('span')[0];
+    const textNode = span.previousSibling;
+    expect(textNode.nodeValue).toBe('Hello');
+    expect(span.textContent).toBe('World');
+
+    // Render an update, that resuspends the parent boundary.
+    // Flushing now now hide the text content.
+    await act(() => {
+      root.render(<App text="Hello" resuspend={true} />);
+    });
+
+    expect(ref.current).toBe(null);
+    expect(span.style.display).toBe('none');
+    expect(textNode.nodeValue).toBe('');
+
+    // Unsuspending shows the content.
+    await act(async () => {
+      suspend = false;
+      resolve();
+      await promise;
+    });
+
+    expect(textNode.nodeValue).toBe('Hello');
+    expect(span.textContent).toBe('World');
+    expect(span.style.display).toBe('');
+    expect(ref.current).toBe(span);
+  });
 });

packages/react-reconciler/src/ReactFiberCommitHostEffects.js

@@ -41,8 +41,10 @@ import {
   insertBefore,
   insertInContainerBefore,
   replaceContainerChildren,
+  hideSuspenseBoundary,
   hideInstance,
   hideTextInstance,
+  unhideSuspenseBoundary,
   unhideInstance,
   unhideTextInstance,
   commitHydratedContainer,
@@ -152,6 +154,27 @@ export function commitHostResetTextContent(finishedWork: Fiber) {
   }
 }
 
+export function commitShowHideSuspenseBoundary(node: Fiber, isHidden: boolean) {
+  try {
+    const instance = node.stateNode;
+    if (isHidden) {
+      if (__DEV__) {
+        runWithFiberInDEV(node, hideSuspenseBoundary, instance);
+      } else {
+        hideSuspenseBoundary(instance);
+      }
+    } else {
+      if (__DEV__) {
+        runWithFiberInDEV(node, unhideSuspenseBoundary, node.stateNode);
+      } else {
+        unhideSuspenseBoundary(node.stateNode);
+      }
+    }
+  } catch (error) {
+    captureCommitPhaseError(node, node.return, error);
+  }
+}
+
 export function commitShowHideHostInstance(node: Fiber, isHidden: boolean) {
   try {
     const instance = node.stateNode;

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -227,6 +227,7 @@ import {
   commitHostUpdate,
   commitHostTextUpdate,
   commitHostResetTextContent,
+  commitShowHideSuspenseBoundary,
   commitShowHideHostInstance,
   commitShowHideHostTextInstance,
   commitHostPlacement,
@@ -1158,6 +1159,10 @@ function hideOrUnhideAllChildren(finishedWork: Fiber, isHidden: boolean) {
         if (hostSubtreeRoot === null) {
           commitShowHideHostTextInstance(node, isHidden);
         }
+      } else if (node.tag === DehydratedFragment) {
+        if (hostSubtreeRoot === null) {
+          commitShowHideSuspenseBoundary(node, isHidden);
+        }
       } else if (
         (node.tag === OffscreenComponent ||
           node.tag === LegacyHiddenComponent) &&

packages/react-reconciler/src/ReactFiberConfigWithNoHydration.js

@@ -44,6 +44,7 @@ export const commitHydratedContainer = shim;
 export const commitHydratedSuspenseInstance = shim;
 export const clearSuspenseBoundary = shim;
 export const clearSuspenseBoundaryFromContainer = shim;
+export const hideSuspenseBoundary = shim;
 export const shouldDeleteUnhydratedTailInstances = shim;
 export const diffHydratedPropsForDevWarnings = shim;
 export const diffHydratedTextForDevWarnings = shim;

packages/react-reconciler/src/forks/ReactFiberConfig.custom.js

@@ -220,6 +220,7 @@ export const commitHydratedSuspenseInstance =
 export const clearSuspenseBoundary = $$$config.clearSuspenseBoundary;
 export const clearSuspenseBoundaryFromContainer =
   $$$config.clearSuspenseBoundaryFromContainer;
+export const hideSuspenseBoundary = $$$config.hideSuspenseBoundary;
 export const shouldDeleteUnhydratedTailInstances =
   $$$config.shouldDeleteUnhydratedTailInstances;
 export const diffHydratedPropsForDevWarnings =