Fix for nested boundaries and suspense in root container

This is a follow up to facebook#16673 which didn't have a test because it wasn't
observable yet. This shows that it had a bug.

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

@@ -1944,6 +1944,164 @@ describe('ReactDOMServerPartialHydration', () => {
     document.body.removeChild(container);
   });
 
+  it('invokes discrete events on nested suspense boundaries in a root (legacy system)', async () => {
+    let suspend = false;
+    let resolve;
+    let promise = new Promise(resolvePromise => (resolve = resolvePromise));
+
+    let clicks = 0;
+
+    function Button() {
+      return (
+        <a
+          onClick={() => {
+            clicks++;
+          }}>
+          Click me
+        </a>
+      );
+    }
+
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return (
+          <Suspense fallback="Loading...">
+            <Button />
+          </Suspense>
+        );
+      }
+    }
+
+    function App() {
+      return (
+        <Suspense fallback="Loading...">
+          <Child />
+        </Suspense>
+      );
+    }
+
+    suspend = false;
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    // We need this to be in the document since we'll dispatch events on it.
+    document.body.appendChild(container);
+
+    let a = container.getElementsByTagName('a')[0];
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    suspend = true;
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+
+    // We'll do one click before hydrating.
+    a.click();
+    // This should be delayed.
+    expect(clicks).toBe(0);
+
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    // We're now partially hydrated.
+    a.click();
+    expect(clicks).toBe(0);
+
+    // Resolving the promise so that rendering can complete.
+    suspend = false;
+    resolve();
+    await promise;
+
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    expect(clicks).toBe(2);
+
+    document.body.removeChild(container);
+  });
+
+  it('invokes discrete events on nested suspense boundaries in a root (responder system)', async () => {
+    let suspend = false;
+    let resolve;
+    let promise = new Promise(resolvePromise => (resolve = resolvePromise));
+
+    const onEvent = jest.fn();
+    const TestResponder = React.unstable_createResponder('TestEventResponder', {
+      targetEventTypes: ['click'],
+      onEvent,
+    });
+
+    function Button() {
+      let listener = React.unstable_useResponder(TestResponder, {});
+      return <a listeners={listener}>Click me</a>;
+    }
+
+    function Child() {
+      if (suspend) {
+        throw promise;
+      } else {
+        return (
+          <Suspense fallback="Loading...">
+            <Button />
+          </Suspense>
+        );
+      }
+    }
+
+    function App() {
+      return (
+        <Suspense fallback="Loading...">
+          <Child />
+        </Suspense>
+      );
+    }
+
+    suspend = false;
+    let finalHTML = ReactDOMServer.renderToString(<App />);
+    let container = document.createElement('div');
+    container.innerHTML = finalHTML;
+
+    // We need this to be in the document since we'll dispatch events on it.
+    document.body.appendChild(container);
+
+    let a = container.getElementsByTagName('a')[0];
+
+    // On the client we don't have all data yet but we want to start
+    // hydrating anyway.
+    suspend = true;
+    let root = ReactDOM.unstable_createRoot(container, {hydrate: true});
+    root.render(<App />);
+
+    // We'll do one click before hydrating.
+    a.click();
+    // This should be delayed.
+    expect(onEvent).toHaveBeenCalledTimes(0);
+
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    // We're now partially hydrated.
+    a.click();
+    // We should not have invoked the event yet because we're not
+    // yet hydrated.
+    expect(onEvent).toHaveBeenCalledTimes(0);
+
+    // Resolving the promise so that rendering can complete.
+    suspend = false;
+    resolve();
+    await promise;
+
+    Scheduler.unstable_flushAll();
+    jest.runAllTimers();
+
+    expect(onEvent).toHaveBeenCalledTimes(2);
+
+    document.body.removeChild(container);
+  });
+
   it('does not invoke the parent of dehydrated boundary event', async () => {
     let suspend = false;
     let resolve;

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

@@ -62,14 +62,21 @@ export function getClosestInstanceFromNode(targetNode) {
       // a nested suspense boundary within it. So we can use this as a fast
       // bailout. Most of the time, when people add non-React children to
       // the tree, it is using a ref to a child-less DOM node.
-      // We only need to check one of the fibers because if it has ever
-      // gone from having children to deleting them or vice versa it would
-      // have deleted the dehydrated boundary nested inside already.
-      if (targetInst.child !== null) {
+      // Normally we'd only need to check one of the fibers because if it
+      // has ever gone from having children to deleting them or vice versa
+      // it would have deleted the dehydrated boundary nested inside already.
+      // However, since the HostRoot starts out with an alternate it might
+      // have one on the alternate so we need to check in case this was a
+      // root.
+      const alternate = targetInst.alternate;
+      if (
+        targetInst.child !== null ||
+        (alternate !== null && alternate.child !== null)
+      ) {
         // Next we need to figure out if the node that skipped past is
         // nested within a dehydrated boundary and if so, which one.
         let suspenseInstance = getParentSuspenseInstance(targetNode);
-        if (suspenseInstance !== null) {
+        while (suspenseInstance !== null) {
           // We found a suspense instance. That means that we haven't
           // hydrated it yet. Even though we leave the comments in the
           // DOM after hydrating, and there are boundaries in the DOM
@@ -83,10 +90,13 @@ export function getClosestInstanceFromNode(targetNode) {
             return targetSuspenseInst;
           }
           // If we don't find a Fiber on the comment, it might be because
-          // we haven't gotten to hydrate it yet. That should mean that
-          // the parent component also hasn't hydrated yet but we can
-          // just return that since it will bail out on the isMounted
-          // check.
+          // we haven't gotten to hydrate it yet. There might still be a
+          // parent boundary that hasn't above this one so we need to find
+          // the outer most that is known.
+          suspenseInstance = getParentSuspenseInstance(suspenseInstance);
+          // If we don't find one, then that should mean that the parent
+          // host component also hasn't hydrated yet. We can return it
+          // below since it will bail out on the isMounted check later.
         }
       }
       return targetInst;