[Fiber] Fix findDOMNode and findAllInRenderedTree

scripts/fiber/tests-failing.txt

@@ -119,6 +119,3 @@ src/renderers/shared/shared/__tests__/ReactUpdates-test.js
 
 src/renderers/shared/shared/__tests__/refs-test.js
 * Should increase refs with an increase in divs
-
-src/test/__tests__/ReactTestUtils-test.js
-* traverses children in the correct order

scripts/fiber/tests-passing.txt

@@ -507,6 +507,7 @@ src/renderers/dom/fiber/__tests__/ReactDOMFiber-test.js
 * should pass portal context when rendering subtree elsewhere
 * should update portal context if it changes due to setState
 * should update portal context if it changes due to re-render
+* findDOMNode should find dom element after expanding a fragment
 
 src/renderers/dom/shared/__tests__/CSSProperty-test.js
 * should generate browser prefixes for its `isUnitlessNumber`
@@ -713,6 +714,7 @@ src/renderers/dom/shared/__tests__/escapeTextContentForBrowser-test.js
 src/renderers/dom/shared/__tests__/findDOMNode-test.js
 * findDOMNode should return null if passed null
 * findDOMNode should find dom element
+* findDOMNode should find dom element after an update from null
 * findDOMNode should reject random objects
 * findDOMNode should reject unmounted objects with render func
 * findDOMNode should not throw an error when called within a component that is not mounted
@@ -1589,6 +1591,7 @@ src/test/__tests__/ReactTestUtils-test.js
 * can scryRenderedDOMComponentsWithClass with TextComponent
 * can scryRenderedDOMComponentsWithClass with className contains \n
 * can scryRenderedDOMComponentsWithClass with multiple classes
+* traverses children in the correct order
 * should support injected wrapper components as DOM components
 * should change the value of an input field
 * should change the value of an input field in a component

src/renderers/dom/fiber/__tests__/ReactDOMFiber-test.js

@@ -459,5 +459,29 @@ describe('ReactDOMFiber', () => {
       expect(portalContainer.innerHTML).toBe('<div>changed-changed</div>');
       expect(container.innerHTML).toBe('');
     });
+
+    it('findDOMNode should find dom element after expanding a fragment', () => {
+      class MyNode extends React.Component {
+        render() {
+          return (
+            !this.props.flag ?
+            [<div key="a" />] :
+            [<span key="b" />, <div key="a" />]
+          );
+        }
+      }
+
+      var container = document.createElement('div');
+
+      var myNodeA = ReactDOM.render(<MyNode />, container);
+      var a = ReactDOM.findDOMNode(myNodeA);
+      expect(a.tagName).toBe('DIV');
+
+      var myNodeB = ReactDOM.render(<MyNode flag={true} />, container);
+      expect(myNodeA === myNodeB).toBe(true);
+
+      var b = ReactDOM.findDOMNode(myNodeB);
+      expect(b.tagName).toBe('SPAN');
+    });
   }
 });

src/renderers/dom/shared/__tests__/findDOMNode-test.js

@@ -34,6 +34,32 @@ describe('findDOMNode', () => {
     expect(mySameDiv).toBe(myDiv);
   });
 
+  it('findDOMNode should find dom element after an update from null', () => {
+    function Bar({ flag }) {
+      if (flag) {
+        return <span>A</span>;
+      }
+      return null;
+    }
+    class MyNode extends React.Component {
+      render() {
+        return <Bar flag={this.props.flag} />;
+      }
+    }
+
+    var container = document.createElement('div');
+
+    var myNodeA = ReactDOM.render(<MyNode />, container);
+    var a = ReactDOM.findDOMNode(myNodeA);
+    expect(a).toBe(null);
+
+    var myNodeB = ReactDOM.render(<MyNode flag={true} />, container);
+    expect(myNodeA === myNodeB).toBe(true);
+
+    var b = ReactDOM.findDOMNode(myNodeB);
+    expect(b.tagName).toBe('SPAN');
+  });
+
   it('findDOMNode should reject random objects', () => {
     expect(function() {
       ReactDOM.findDOMNode({foo: 'bar'});

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -183,6 +183,9 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
               commitPlacement(effectfulFiber);
               // Clear the "placement" from effect tag so that we know that this is inserted, before
               // any life-cycles like componentDidMount gets called.
+              // TODO: findDOMNode doesn't rely on this any more but isMounted
+              // does and isMounted is deprecated anyway so we should be able
+              // to kill this.
               effectfulFiber.effectTag &= ~Placement;
               break;
             }

src/renderers/shared/fiber/ReactFiberTreeReflection.js

@@ -73,60 +73,116 @@ exports.isMounted = function(component : ReactComponent<any, any, any>) : boolea
   return isFiberMountedImpl(fiber) === MOUNTED;
 };
 
-exports.findCurrentHostFiber = function(parent : Fiber) : Fiber | null {
-  // First check if this node itself is mounted.
-  const state = isFiberMountedImpl(parent, true);
-  if (state === UNMOUNTED) {
+function assertIsMounted(fiber) {
+  invariant(
+    isFiberMountedImpl(fiber) === MOUNTED,
+    'Unable to find node on an unmounted component.'
+  );
+}
+
+function findCurrentFiberUsingSlowPath(fiber : Fiber) : Fiber | null {
+  let alternate = fiber.alternate;
+  if (!alternate) {
+    // If there is no alternate, then we only need to check if it is mounted.
+    const state = isFiberMountedImpl(fiber);
     invariant(
-      false,
+      state !== UNMOUNTED,
       'Unable to find node on an unmounted component.'
     );
-  } else if (state === MOUNTING) {
-    return null;
+    if (state === MOUNTING) {
+      return null;
+    }
+    return fiber;
   }
+  // If we have two possible branches, we'll walk backwards up to the root
+  // to see what path the root points to. On the way we may hit one of the
+  // special cases and we'll deal with them.
+  let a = fiber;
+  let b = alternate;
+  while (true) {
+    let parentA = a.return;
+    let parentB = b.return;
+    if (!parentA || !parentB) {
+      // We're at the root.
+      break;
+    }
+    if (parentA.child === parentB.child) {
+      // If both parents are the same, then that is the current parent. If
+      // they're different but point to the same child, then it doesn't matter.
+      // Regardless, whatever child they point to is the current child.
+      // So we can now determine which child is current by scanning the child
+      // list for either A or B.
+      let child = parentA.child;
+      while (child) {
+        if (child === a) {
+          // We've determined that A is the current branch.
+          assertIsMounted(parentA);
+          return fiber;
+        }
+        if (child === b) {
+          // We've determined that B is the current branch.
+          assertIsMounted(parentA);
+          return alternate;
+        }
+        child = child.sibling;
+      }
+      // We should never have an alternate for any mounting node. So the only
+      // way this could possibly happen is if this was unmounted, if at all.
+      invariant(
+        false,
+        'Unable to find node on an unmounted component.'
+      );
+    }
+    a = parentA;
+    b = parentB;
+    invariant(
+      a.alternate === b,
+      'Return fibers should always be each others\' alternates.'
+    );
+  }
+  // If the root is not a host container, we're in a disconnected tree. I.e.
+  // unmounted.
+  invariant(
+    a.tag === HostRoot,
+    'Unable to find node on an unmounted component.'
+  );
+  if (a.stateNode.current === a) {
+    // We've determined that A is the current branch.
+    return fiber;
+  }
+  // Otherwise B has to be current branch.
+  return alternate;
+}
+exports.findCurrentFiberUsingSlowPath = findCurrentFiberUsingSlowPath;
 
-  let didTryOtherTree = false;
-
-  // If the component doesn't have a child we first check the alternate to see
-  // if it has any and if so, if those were just recently inserted.
-  if (!parent.child && parent.alternate) {
-    parent = parent.alternate;
+exports.findCurrentHostFiber = function(parent : Fiber) : Fiber | null {
+  const currentParent = findCurrentFiberUsingSlowPath(parent);
+  if (!currentParent) {
+    return null;
   }
 
   // Next we'll drill down this component to find the first HostComponent/Text.
-  let node : Fiber = parent;
+  let node : Fiber = currentParent;
   while (true) {
-    if ((node.effectTag & Placement) !== NoEffect || !node.return) {
-      // If any node along the way was deleted, or is an insertion, that means
-      // that we're actually in a work in progress to update this component with
-      // a different component. We need to look in the "current" fiber instead.
-      if (!parent.alternate) {
-        return null;
-      }
-      if (didTryOtherTree) {
-        // Safety, to avoid an infinite loop if something goes wrong.
-        throw new Error('This should never hit this infinite loop.');
-      }
-      didTryOtherTree = true;
-      node = parent = parent.alternate;
-      continue;
-    }
     if (node.tag === HostComponent || node.tag === HostText) {
       return node;
     } else if (node.child) {
+      // TODO: If we hit a Portal, we're supposed to skip it.
       // TODO: Coroutines need to visit the stateNode.
+      node.child.return = node;
       node = node.child;
       continue;
     }
-    if (node === parent) {
+    if (node === currentParent) {
       return null;
     }
     while (!node.sibling) {
-      if (!node.return || node.return === parent) {
+      if (!node.return || node.return === currentParent) {
         return null;
       }
       node = node.return;
     }
+    node.sibling.return = node.return;
     node = node.sibling;
   }
   // Flow needs the return null here, but ESLint complains about it.

src/test/ReactTestUtils.js

@@ -20,9 +20,9 @@ var ReactControlledComponent = require('ReactControlledComponent');
 var ReactDOM = require('ReactDOM');
 var ReactDOMComponentTree = require('ReactDOMComponentTree');
 var ReactBrowserEventEmitter = require('ReactBrowserEventEmitter');
+var ReactFiberTreeReflection = require('ReactFiberTreeReflection');
 var ReactInstanceMap = require('ReactInstanceMap');
 var ReactTypeOfWork = require('ReactTypeOfWork');
-var ReactTypeOfSideEffect = require('ReactTypeOfSideEffect');
 var ReactGenericBatching = require('ReactGenericBatching');
 var SyntheticEvent = require('SyntheticEvent');
 var ReactShallowRenderer = require('ReactShallowRenderer');
@@ -37,10 +37,6 @@ var {
   HostComponent,
   HostText,
 } = ReactTypeOfWork;
-var {
-  NoEffect,
-  Placement,
-} = ReactTypeOfSideEffect;
 
 function Event(suffix) {}
 
@@ -84,34 +80,40 @@ function findAllInRenderedFiberTreeInternal(fiber, test) {
   if (!fiber) {
     return [];
   }
-  if (
-    fiber.tag !== ClassComponent &&
-    fiber.tag !== FunctionalComponent &&
-    fiber.tag !== HostComponent &&
-    fiber.tag !== HostText
-  ) {
+  var currentParent = ReactFiberTreeReflection.findCurrentFiberUsingSlowPath(
+    fiber
+  );
+  if (!currentParent) {
     return [];
   }
-  if ((fiber.effectTag & Placement) !== NoEffect || !fiber.return) {
-    // If any node along the way was deleted, or is an insertion, that means
-    // that we're actually in a work in progress to update this component with
-    // a different component. We need to look in the "current" fiber instead.
-    return null;
-  }
-  var publicInst = fiber.stateNode;
-  var ret = publicInst && test(publicInst) ? [publicInst] : [];
-  var child = fiber.child;
-  while (child) {
-    ret = ret.concat(
-      findAllInRenderedFiberTreeInternal(
-        child,
-        test
-      )
-    );
-    child = child.sibling;
+  let node = currentParent;
+  let ret = [];
+  while (true) {
+    if (node.tag === HostComponent || node.tag === HostText ||
+        node.tag === ClassComponent || node.tag === FunctionalComponent) {
+      var publicInst = node.stateNode;
+      if (test(publicInst)) {
+        ret.push(publicInst);
+      }
+    }
+    if (node.child) {
+      // TODO: Coroutines need to visit the stateNode.
+      node.child.return = node;
+      node = node.child;
+      continue;
+    }
+    if (node === currentParent) {
+      return ret;
+    }
+    while (!node.sibling) {
+      if (!node.return || node.return === currentParent) {
+        return ret;
+      }
+      node = node.return;
+    }
+    node.sibling.return = node.return;
+    node = node.sibling;
   }
-  // TODO: visit stateNode for coroutines
-  return ret;
 }
 
 /**
@@ -222,21 +224,7 @@ var ReactTestUtils = {
     );
     var internalInstance = ReactInstanceMap.get(inst);
     if (internalInstance && typeof internalInstance.tag === 'number') {
-      var fiber = internalInstance;
-      if (!fiber.child && fiber.alternate) {
-        // When we don't have children, we try the alternate first to see if it
-        // has any current children first.
-        fiber = fiber.alternate;
-      }
-      var results = findAllInRenderedFiberTreeInternal(fiber, test);
-      if (results === null) {
-        // Null is a sentinel that indicates that this was the wrong fiber.
-        results = findAllInRenderedFiberTreeInternal(fiber.alternate, test);
-        if (results === null) {
-          throw new Error('This should never happen.');
-        }
-      }
-      return results;
+      return findAllInRenderedFiberTreeInternal(internalInstance, test);
     } else {
       return findAllInRenderedStackTreeInternal(internalInstance, test);
     }