Delete all children within a dehydrated suspense boundary when it's d…

…eleted

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

@@ -137,4 +137,57 @@ describe('ReactDOMServerPartialHydration', () => {
     expect(container.firstChild.firstChild.tagName).toBe('DIV');
     expect(container.firstChild.firstChild.textContent).toBe('First');
   });
+
+  it('can delete the dehydrated boundary before it is hydrated', () => {
+    let suspend = false;
+    let promise = new Promise(() => {});
+    let hideMiddle;
+
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return (
+          <React.Fragment>
+            <div>Middle</div>
+            Some text
+          </React.Fragment>
+        );
+      }
+    }
+
+    function App() {
+      let [visible, setVisibilty] = React.useState(true);
+      hideMiddle = () => setVisibilty(false);
+
+      return (
+        <div>
+          <div>Before</div>
+          {visible ? (
+            <Suspense fallback="Loading...">
+              <Child />
+            </Suspense>
+          ) : null}
+          <div>After</div>
+        </div>
+      );
+    }
+
+    suspend = false;
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+    let 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;
+    ReactDOM.hydrate(<App />, container);
+
+    expect(container.firstChild.children[1].textContent).toBe('Middle');
+
+    // In this state, we can still delete the boundary.
+    hideMiddle();
+
+    expect(container.firstChild.children[1].textContent).toBe('After');
+  });
 });

packages/react-dom/src/client/ReactDOMHostConfig.js

@@ -436,6 +436,49 @@ export function removeChildFromContainer(
   }
 }
 
+export function clearSuspenseBoundary(
+  parentInstance: Instance,
+  suspenseInstance: SuspenseInstance,
+): void {
+  let node = suspenseInstance;
+  // 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.
+  let depth = 0;
+  do {
+    let nextNode = node.nextSibling;
+    parentInstance.removeChild(node);
+    if (nextNode && nextNode.nodeType === COMMENT_NODE) {
+      let data = ((nextNode: any).data: string);
+      if (data === SUSPENSE_END_DATA) {
+        if (depth === 0) {
+          parentInstance.removeChild(nextNode);
+          return;
+        } else {
+          depth--;
+        }
+      } else if (nextNode === SUSPENSE_START_DATA) {
+        depth++;
+      }
+    }
+    node = nextNode;
+  } while (node);
+  // TODO: Warn, we didn't find the end comment boundary.
+}
+
+export function clearSuspenseBoundaryFromContainer(
+  container: Container,
+  suspenseInstance: SuspenseInstance,
+): void {
+  if (container.nodeType === COMMENT_NODE) {
+    clearSuspenseBoundary((container.parentNode: any), suspenseInstance);
+  } else if (container.nodeType === ELEMENT_NODE) {
+    clearSuspenseBoundary((container: any), suspenseInstance);
+  } else {
+    // Document nodes should never contain suspense boundaries.
+  }
+}
+
 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?

packages/react-reconciler/src/ReactFiberCommitWork.js

@@ -10,6 +10,7 @@
 import type {
   Instance,
   TextInstance,
+  SuspenseInstance,
   Container,
   ChildSet,
   UpdatePayload,
@@ -80,6 +81,8 @@ import {
   insertInContainerBefore,
   removeChild,
   removeChildFromContainer,
+  clearSuspenseBoundary,
+  clearSuspenseBoundaryFromContainer,
   replaceContainerChildren,
   createContainerChildSet,
   hideInstance,
@@ -1037,11 +1040,30 @@ function unmountHostComponents(current): void {
       // After all the children have unmounted, it is now safe to remove the
       // node from the tree.
       if (currentParentIsContainer) {
-        removeChildFromContainer((currentParent: any), node.stateNode);
+        removeChildFromContainer(
+          ((currentParent: any): Container),
+          (node.stateNode: Instance | TextInstance),
+        );
       } else {
-        removeChild((currentParent: any), node.stateNode);
+        removeChild(
+          ((currentParent: any): Instance),
+          (node.stateNode: Instance | TextInstance),
+        );
       }
       // Don't visit children because we already visited them.
+    } else if (node.tag === DehydratedSuspenseComponent) {
+      // Delete the dehydrated suspense boundary and all of its content.
+      if (currentParentIsContainer) {
+        clearSuspenseBoundaryFromContainer(
+          ((currentParent: any): Container),
+          (node.stateNode: SuspenseInstance),
+        );
+      } else {
+        clearSuspenseBoundary(
+          ((currentParent: any): Instance),
+          (node.stateNode: SuspenseInstance),
+        );
+      }
     } else if (node.tag === HostPortal) {
       // When we go into a portal, it becomes the parent to remove from.
       // We will reassign it back when we pop the portal on the way up.

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

@@ -29,6 +29,7 @@ export opaque type Props = mixed; // eslint-disable-line no-undef
 export opaque type Container = mixed; // eslint-disable-line no-undef
 export opaque type Instance = mixed; // eslint-disable-line no-undef
 export opaque type TextInstance = mixed; // eslint-disable-line no-undef
+export opaque type SuspenseInstance = mixed; // eslint-disable-line no-undef
 export opaque type HydratableInstance = mixed; // eslint-disable-line no-undef
 export opaque type PublicInstance = mixed; // eslint-disable-line no-undef
 export opaque type HostContext = mixed; // eslint-disable-line no-undef
@@ -112,6 +113,9 @@ export const hydrateInstance = $$$hostConfig.hydrateInstance;
 export const hydrateTextInstance = $$$hostConfig.hydrateTextInstance;
 export const getNextHydratableInstanceAfterSuspenseInstance =
   $$$hostConfig.getNextHydratableInstanceAfterSuspenseInstance;
+export const clearSuspenseBoundary = $$$hostConfig.clearSuspenseBoundary;
+export const clearSuspenseBoundaryFromContainer =
+  $$$hostConfig.clearSuspenseBoundaryFromContainer;
 export const didNotMatchHydratedContainerTextInstance =
   $$$hostConfig.didNotMatchHydratedContainerTextInstance;
 export const didNotMatchHydratedTextInstance =

packages/shared/HostConfigWithNoHydration.js

@@ -22,6 +22,7 @@ function shim(...args: any) {
 }
 
 // Hydration (when unsupported)
+export type SuspenseInstance = mixed;
 export const supportsHydration = false;
 export const canHydrateInstance = shim;
 export const canHydrateTextInstance = shim;
@@ -31,6 +32,8 @@ export const getFirstHydratableChild = shim;
 export const hydrateInstance = shim;
 export const hydrateTextInstance = shim;
 export const getNextHydratableInstanceAfterSuspenseInstance = shim;
+export const clearSuspenseBoundary = shim;
+export const clearSuspenseBoundaryFromContainer = shim;
 export const didNotMatchHydratedContainerTextInstance = shim;
 export const didNotMatchHydratedTextInstance = shim;
 export const didNotHydrateContainerInstance = shim;