Resolve ref callbacks

During the deletion phase we call detachRefs on any deleted refs.

During the insertion/update phase we call attachRef on class
and host components.

Unfortunately, when a ref switches, we are supposed to call all
the unmounts before doing any mounts. This means that we have to
expact the deletion phase to also include updates in case they
need to detach their ref.

src/renderers/shared/fiber/ReactChildFiber.js

@@ -221,11 +221,13 @@ function ChildReconciler(shouldClone, shouldTrackSideEffects) {
     if (current == null || current.type !== element.type) {
       // Insert
       const created = createFiberFromElement(element, priority);
+      created.ref = element.ref;
       created.return = returnFiber;
       return created;
     } else {
       // Move based on index
       const existing = useFiber(current, priority);
+      existing.ref = element.ref;
       existing.pendingProps = element.props;
       existing.return = returnFiber;
       return existing;
@@ -317,6 +319,7 @@ function ChildReconciler(shouldClone, shouldTrackSideEffects) {
       switch (newChild.$$typeof) {
         case REACT_ELEMENT_TYPE: {
           const created = createFiberFromElement(newChild, priority);
+          created.ref = newChild.ref;
           created.return = returnFiber;
           return created;
         }
@@ -643,6 +646,7 @@ function ChildReconciler(shouldClone, shouldTrackSideEffects) {
         if (child.type === element.type) {
           deleteRemainingChildren(returnFiber, child.sibling);
           const existing = useFiber(child, priority);
+          existing.ref = element.ref;
           existing.pendingProps = element.props;
           existing.return = returnFiber;
           return existing;
@@ -657,6 +661,7 @@ function ChildReconciler(shouldClone, shouldTrackSideEffects) {
     }
 
     const created = createFiberFromElement(element, priority);
+    created.ref = element.ref;
     created.return = returnFiber;
     return created;
   }

src/renderers/shared/fiber/ReactFiberCommitWork.js

@@ -40,6 +40,29 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
   const insertBefore = config.insertBefore;
   const removeChild = config.removeChild;
 
+  function detachRef(current : Fiber) {
+    const ref = current.ref;
+    if (ref) {
+      ref(null);
+    }
+  }
+
+  function attachRef(current : ?Fiber, finishedWork : Fiber, instance : any) {
+    const ref = finishedWork.ref;
+    if (current) {
+      const currentRef = current.ref;
+      if (currentRef && currentRef !== ref) {
+        // TODO: This needs to be done in a separate pass before any other refs
+        // gets resolved. Otherwise we might invoke them in the wrong order
+        // when the same ref switches between two components.
+        currentRef(null);
+      }
+    }
+    if (ref) {
+      ref(instance);
+    }
+  }
+
   function getHostParent(fiber : Fiber) : ?I {
     let parent = fiber.return;
     while (parent) {
@@ -163,9 +186,6 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
     // Recursively delete all host nodes from the parent.
     // TODO: Error handling.
     const parent = getHostParent(current);
-    if (!parent) {
-      return;
-    }
     // We only have the top Fiber that was inserted but we need recurse down its
     // children to find all the terminal nodes.
     // TODO: Call componentWillUnmount on all classes as needed. Recurse down
@@ -177,7 +197,9 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
         commitNestedUnmounts(node);
         // After all the children have unmounted, it is now safe to remove the
         // node from the tree.
-        removeChild(parent, node.stateNode);
+        if (parent) {
+          removeChild(parent, node.stateNode);
+        }
       } else {
         commitUnmount(node);
         if (node.child) {
@@ -202,10 +224,16 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
   function commitUnmount(current : Fiber) : void {
     switch (current.tag) {
       case ClassComponent: {
+        detachRef(current);
         const instance = current.stateNode;
         if (typeof instance.componentWillUnmount === 'function') {
           instance.componentWillUnmount();
         }
+        return;
+      }
+      case HostComponent: {
+        detachRef(current);
+        return;
       }
     }
   }
@@ -239,7 +267,7 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
           finishedWork.callbackList = null;
           callCallbacks(callbackList, instance);
         }
-        // TODO: Fire update refs
+        attachRef(current, finishedWork, instance);
         return;
       }
       case HostContainer: {
@@ -251,14 +279,14 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
         return;
       }
       case HostComponent: {
-        if (finishedWork.stateNode == null || !current) {
-          throw new Error('This should only be done during updates.');
-        }
-        // Commit the work prepared earlier.
-        const newProps = finishedWork.memoizedProps;
-        const oldProps = current.memoizedProps;
         const instance : I = finishedWork.stateNode;
-        commitUpdate(instance, oldProps, newProps);
+        if (instance != null && current) {
+          // Commit the work prepared earlier.
+          const newProps = finishedWork.memoizedProps;
+          const oldProps = current.memoizedProps;
+          commitUpdate(instance, oldProps, newProps);
+        }
+        attachRef(current, finishedWork, instance);
         return;
       }
       case HostText: {

src/renderers/shared/fiber/ReactFiberCompleteWork.js

@@ -175,6 +175,10 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
           // TODO: This seems like unnecessary duplication.
           workInProgress.stateNode = instance;
           workInProgress.output = instance;
+          if (workInProgress.ref) {
+            // If there is a ref on a host node we need to schedule a callback
+            markUpdate(workInProgress);
+          }
         }
         workInProgress.memoizedProps = newProps;
         return null;

src/renderers/shared/fiber/__tests__/ReactIncrementalSideEffects-test.js

@@ -637,6 +637,72 @@ describe('ReactIncrementalSideEffects', () => {
 
   });
 
+  it('invokes ref callbacks after insertion/update/unmount', () => {
+
+    var classInstance = null;
+
+    var ops = [];
+
+    class ClassComponent extends React.Component {
+      render() {
+        classInstance = this;
+        return <span />;
+      }
+    }
+
+    function FunctionalComponent(props) {
+      return <span />;
+    }
+
+    function Foo(props) {
+      return (
+        props.show ?
+        <div>
+          <ClassComponent ref={n => ops.push(n)} />
+          <FunctionalComponent ref={n => ops.push(n)} />
+          <div ref={n => ops.push(n)} />
+        </div> :
+        null
+      );
+    }
+
+    ReactNoop.render(<Foo show={true} />);
+    ReactNoop.flush();
+    expect(ops).toEqual([
+      classInstance,
+      // no call for functional components
+      div(),
+    ]);
+
+    ops = [];
+
+    // Refs that switch function instances get reinvoked
+    ReactNoop.render(<Foo show={true} />);
+    ReactNoop.flush();
+    expect(ops).toEqual([
+      // TODO: All detach should happen first. Currently they're interleaved.
+      // detach
+      null,
+      // reattach
+      classInstance,
+      // detach
+      null,
+      // reattach
+      div(),
+    ]);
+
+    ops = [];
+
+    ReactNoop.render(<Foo show={false} />);
+    ReactNoop.flush();
+    expect(ops).toEqual([
+      // unmount
+      null,
+      null,
+    ]);
+
+  });
+
   // TODO: Test that mounts, updates, refs, unmounts and deletions happen in the
   // expected way for aborted and resumed render life-cycles.