Fix fragmentInstance#compareDocumentPosition nesting and portal cases

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

@@ -38,9 +38,16 @@ import hasOwnProperty from 'shared/hasOwnProperty';
 import {checkAttributeStringCoercion} from 'shared/CheckStringCoercion';
 import {REACT_CONTEXT_TYPE} from 'shared/ReactSymbols';
 import {
-  isFiberContainedBy,
+  isFiberContainedByFragment,
   isFiberFollowing,
   isFiberPreceding,
+  isFragmentContainedByFiber,
+  traverseFragmentInstance,
+  getFragmentParentHostFiber,
+  getNextSiblingHostFiber,
+  getInstanceFromHostFiber,
+  traverseFragmentInstanceDeeply,
+  fiberIsPortaledIntoHost,
 } from 'react-reconciler/src/ReactFiberTreeReflection';
 
 export {
@@ -63,13 +70,6 @@ import {
   markNodeAsHoistable,
   isOwnedInstance,
 } from './ReactDOMComponentTree';
-import {
-  traverseFragmentInstance,
-  getFragmentParentHostFiber,
-  getNextSiblingHostFiber,
-  getInstanceFromHostFiber,
-  traverseFragmentInstanceDeeply,
-} from 'react-reconciler/src/ReactFiberTreeReflection';
 
 export {detachDeletedInstance};
 import {hasRole} from './DOMAccessibilityRoles';
@@ -3052,13 +3052,13 @@ FragmentInstance.prototype.compareDocumentPosition = function (
   }
   const children: Array<Fiber> = [];
   traverseFragmentInstance(this._fragmentFiber, collectChildren, children);
+  const parentHostInstance =
+    getInstanceFromHostFiber<Instance>(parentHostFiber);
 
   let result = Node.DOCUMENT_POSITION_DISCONNECTED;
   if (children.length === 0) {
     // If the fragment has no children, we can use the parent and
     // siblings to determine a position.
-    const parentHostInstance =
-      getInstanceFromHostFiber<Instance>(parentHostFiber);
     const parentResult = parentHostInstance.compareDocumentPosition(otherNode);
     result = parentResult;
     if (parentHostInstance === otherNode) {
@@ -3095,15 +3095,53 @@ FragmentInstance.prototype.compareDocumentPosition = function (
   const lastElement = getInstanceFromHostFiber<Instance>(
     children[children.length - 1],
   );
+
+  // If the fragment has been portaled into another host instance, we need to
+  // our best guess is to use the parent of the child instance, rather than
+  // the fiber tree host parent.
+  const parentHostInstanceFromDOM = fiberIsPortaledIntoHost(this._fragmentFiber)
+    ? (getInstanceFromHostFiber<Instance>(children[0]).parentElement: ?Instance)
+    : parentHostInstance;
+
+  if (parentHostInstanceFromDOM == null) {
+    return Node.DOCUMENT_POSITION_DISCONNECTED;
+  }
+
+  // Check if first and last element are actually in the expected document position
+  // before relying on them as source of truth for other contained elements
+  const firstElementIsContained =
+    parentHostInstanceFromDOM.compareDocumentPosition(firstElement) &
+    Node.DOCUMENT_POSITION_CONTAINED_BY;
+  const lastElementIsContained =
+    parentHostInstanceFromDOM.compareDocumentPosition(lastElement) &
+    Node.DOCUMENT_POSITION_CONTAINED_BY;
   const firstResult = firstElement.compareDocumentPosition(otherNode);
   const lastResult = lastElement.compareDocumentPosition(otherNode);
+
+  const otherNodeIsFirstOrLastChild =
+    (firstElementIsContained && firstElement === otherNode) ||
+    (lastElementIsContained && lastElement === otherNode);
+  const otherNodeIsFirstOrLastChildDisconnected =
+    (!firstElementIsContained && firstElement === otherNode) ||
+    (!lastElementIsContained && lastElement === otherNode);
+  const otherNodeIsWithinFirstOrLastChild =
+    firstResult & Node.DOCUMENT_POSITION_CONTAINED_BY ||
+    lastResult & Node.DOCUMENT_POSITION_CONTAINED_BY;
+  const otherNodeIsBetweenFirstAndLastChildren =
+    firstElementIsContained &&
+    lastElementIsContained &&
+    firstResult & Node.DOCUMENT_POSITION_FOLLOWING &&
+    lastResult & Node.DOCUMENT_POSITION_PRECEDING;
+
   if (
-    (firstResult & Node.DOCUMENT_POSITION_FOLLOWING &&
-      lastResult & Node.DOCUMENT_POSITION_PRECEDING) ||
-    otherNode === firstElement ||
-    otherNode === lastElement
+    otherNodeIsFirstOrLastChild ||
+    otherNodeIsWithinFirstOrLastChild ||
+    otherNodeIsBetweenFirstAndLastChildren
   ) {
     result = Node.DOCUMENT_POSITION_CONTAINED_BY;
+  } else if (otherNodeIsFirstOrLastChildDisconnected) {
+    // otherNode has been portaled into another container
+    result = Node.DOCUMENT_POSITION_IMPLEMENTATION_SPECIFIC;
   } else {
     result = firstResult;
   }
@@ -3141,15 +3179,17 @@ function validateDocumentPositionWithFiberTree(
 ): boolean {
   const otherFiber = getClosestInstanceFromNode(otherNode);
   if (documentPosition & Node.DOCUMENT_POSITION_CONTAINED_BY) {
-    return !!otherFiber && isFiberContainedBy(fragmentFiber, otherFiber);
+    return (
+      !!otherFiber && isFiberContainedByFragment(otherFiber, fragmentFiber)
+    );
   }
   if (documentPosition & Node.DOCUMENT_POSITION_CONTAINS) {
     if (otherFiber === null) {
       // otherFiber could be null if its the document or body element
       const ownerDocument = otherNode.ownerDocument;
       return otherNode === ownerDocument || otherNode === ownerDocument.body;
     }
-    return isFiberContainedBy(otherFiber, fragmentFiber);
+    return isFragmentContainedByFiber(fragmentFiber, otherFiber);
   }
   if (documentPosition & Node.DOCUMENT_POSITION_PRECEDING) {
     return (

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

@@ -1197,14 +1197,14 @@ describe('FragmentRefs', () => {
 
       function Test() {
         return (
-          <div ref={containerRef}>
-            <div ref={beforeRef} />
+          <div ref={containerRef} id="container">
+            <div ref={beforeRef} id="before" />
             <React.Fragment ref={fragmentRef}>
-              <div ref={firstChildRef} />
-              <div ref={middleChildRef} />
-              <div ref={lastChildRef} />
+              <div ref={firstChildRef} id="first" />
+              <div ref={middleChildRef} id="middle" />
+              <div ref={lastChildRef} id="last" />
             </React.Fragment>
-            <div ref={afterRef} />
+            <div ref={afterRef} id="after" />
           </div>
         );
       }
@@ -1289,7 +1289,7 @@ describe('FragmentRefs', () => {
         },
       );
 
-      // containerRef preceds and contains the fragment
+      // containerRef precedes and contains the fragment
       expectPosition(
         fragmentRef.current.compareDocumentPosition(containerRef.current),
         {
@@ -1328,7 +1328,7 @@ describe('FragmentRefs', () => {
       function Test() {
         return (
           <div id="container" ref={containerRef}>
-            <div>
+            <div id="innercontainer">
               <div ref={beforeRef} id="before" />
               <React.Fragment ref={fragmentRef}>
                 <div ref={onlyChildRef} id="within" />
@@ -1491,6 +1491,77 @@ describe('FragmentRefs', () => {
       );
     });
 
+    // @gate enableFragmentRefs
+    it('handles nested children', async () => {
+      const fragmentRef = React.createRef();
+      const nestedFragmentRef = React.createRef();
+      const childARef = React.createRef();
+      const childBRef = React.createRef();
+      const childCRef = React.createRef();
+      document.body.appendChild(container);
+      const root = ReactDOMClient.createRoot(container);
+
+      function Child() {
+        return (
+          <div ref={childCRef} id="C">
+            C
+          </div>
+        );
+      }
+
+      function Test() {
+        return (
+          <React.Fragment ref={fragmentRef}>
+            <div ref={childARef} id="A">
+              A
+            </div>
+            <React.Fragment ref={nestedFragmentRef}>
+              <div ref={childBRef} id="B">
+                B
+              </div>
+            </React.Fragment>
+            <Child />
+          </React.Fragment>
+        );
+      }
+
+      await act(() => root.render(<Test />));
+
+      expectPosition(
+        fragmentRef.current.compareDocumentPosition(childARef.current),
+        {
+          preceding: false,
+          following: false,
+          contains: false,
+          containedBy: true,
+          disconnected: false,
+          implementationSpecific: false,
+        },
+      );
+      expectPosition(
+        fragmentRef.current.compareDocumentPosition(childBRef.current),
+        {
+          preceding: false,
+          following: false,
+          contains: false,
+          containedBy: true,
+          disconnected: false,
+          implementationSpecific: false,
+        },
+      );
+      expectPosition(
+        fragmentRef.current.compareDocumentPosition(childCRef.current),
+        {
+          preceding: false,
+          following: false,
+          contains: false,
+          containedBy: true,
+          disconnected: false,
+          implementationSpecific: false,
+        },
+      );
+    });
+
     // @gate enableFragmentRefs
     it('returns disconnected for comparison with an unmounted fragment instance', async () => {
       const fragmentRef = React.createRef();
@@ -1551,11 +1622,11 @@ describe('FragmentRefs', () => {
 
         function Test() {
           return (
-            <div>
-              {createPortal(<div ref={portaledSiblingRef} />, document.body)}
+            <div id="wrapper">
+              {createPortal(<div ref={portaledSiblingRef} id="A" />, container)}
               <Fragment ref={fragmentRef}>
-                {createPortal(<div ref={portaledChildRef} />, document.body)}
-                <div />
+                {createPortal(<div ref={portaledChildRef} id="B" />, container)}
+                <div id="C" />
               </Fragment>
             </div>
           );
@@ -1600,6 +1671,8 @@ describe('FragmentRefs', () => {
         const childARef = React.createRef();
         const childBRef = React.createRef();
         const childCRef = React.createRef();
+        const childDRef = React.createRef();
+        const childERef = React.createRef();
 
         function Test() {
           const [c, setC] = React.useState(false);
@@ -1612,23 +1685,30 @@ describe('FragmentRefs', () => {
               {createPortal(
                 <Fragment ref={fragmentRef}>
                   <div id="A" ref={childARef} />
-                  {c ? <div id="C" ref={childCRef} /> : null}
+                  {c ? (
+                    <div id="C" ref={childCRef}>
+                      <div id="D" ref={childDRef} />
+                    </div>
+                  ) : null}
                 </Fragment>,
                 document.body,
               )}
               {createPortal(<p id="B" ref={childBRef} />, document.body)}
+              <div id="E" ref={childERef} />
             </>
           );
         }
 
         await act(() => root.render(<Test />));
 
-        // Due to effect, order is A->B->C
-        expect(document.body.innerHTML).toBe(
-          '<div></div>' +
+        // Due to effect, order is E / A->B->C->D
+        expect(document.body.outerHTML).toBe(
+          '<body>' +
+            '<div><div id="E"></div></div>' +
             '<div id="A"></div>' +
             '<p id="B"></p>' +
-            '<div id="C"></div>',
+            '<div id="C"><div id="D"></div></div>' +
+            '</body>',
         );
 
         expectPosition(
@@ -1642,7 +1722,6 @@ describe('FragmentRefs', () => {
             implementationSpecific: false,
           },
         );
-
         expectPosition(
           fragmentRef.current.compareDocumentPosition(childARef.current),
           {
@@ -1654,6 +1733,7 @@ describe('FragmentRefs', () => {
             implementationSpecific: false,
           },
         );
+        // Contained by in DOM, but following in React tree
         expectPosition(
           fragmentRef.current.compareDocumentPosition(childBRef.current),
           {
@@ -1676,6 +1756,29 @@ describe('FragmentRefs', () => {
             implementationSpecific: false,
           },
         );
+        expectPosition(
+          fragmentRef.current.compareDocumentPosition(childDRef.current),
+          {
+            preceding: false,
+            following: false,
+            contains: false,
+            containedBy: true,
+            disconnected: false,
+            implementationSpecific: false,
+          },
+        );
+        // Preceding DOM but following in React tree
+        expectPosition(
+          fragmentRef.current.compareDocumentPosition(childERef.current),
+          {
+            preceding: false,
+            following: false,
+            contains: false,
+            containedBy: false,
+            disconnected: false,
+            implementationSpecific: true,
+          },
+        );
       });
 
       // @gate enableFragmentRefs