Move ref code to ReactCompositeComponent

You can only get a ref to a ReactCompositeComponent, so move the ref code here which gives us more flexibility to put it at the correct time in the lifecycle.

There should be no behavior change in this commit.

Test Plan: jest

src/core/ReactComponent.js

@@ -12,27 +12,9 @@
 'use strict';
 
 var ReactElementValidator = require('ReactElementValidator');
-var ReactOwner = require('ReactOwner');
-var ReactRef = require('ReactRef');
 
 var invariant = require('invariant');
 
-function attachRef(ref, component, owner) {
-  if (ref instanceof ReactRef) {
-    ReactRef.attachRef(ref, component);
-  } else {
-    ReactOwner.addComponentAsRefTo(component, ref, owner);
-  }
-}
-
-function detachRef(ref, component, owner) {
-  if (ref instanceof ReactRef) {
-    ReactRef.detachRef(ref, component);
-  } else {
-    ReactOwner.removeComponentAsRefFrom(component, ref, owner);
-  }
-}
-
 /**
  * Components are the basic units of composition in React.
  *
@@ -114,12 +96,6 @@ var ReactComponent = {
       if (__DEV__) {
         ReactElementValidator.checkAndWarnForMutatedProps(this._currentElement);
       }
-
-      var ref = this._currentElement.ref;
-      if (ref != null) {
-        var owner = this._currentElement._owner;
-        attachRef(ref, this, owner);
-      }
       // Effectively: return '';
     },
 
@@ -134,10 +110,6 @@ var ReactComponent = {
      * @internal
      */
     unmountComponent: function() {
-      var ref = this._currentElement.ref;
-      if (ref != null) {
-        detachRef(ref, this, this._currentElement._owner);
-      }
     },
 
     /**
@@ -152,29 +124,6 @@ var ReactComponent = {
       if (__DEV__) {
         ReactElementValidator.checkAndWarnForMutatedProps(nextElement);
       }
-
-      // If either the owner or a `ref` has changed, make sure the newest owner
-      // has stored a reference to `this`, and the previous owner (if different)
-      // has forgotten the reference to `this`. We use the element instead
-      // of the public this.props because the post processing cannot determine
-      // a ref. The ref conceptually lives on the element.
-
-      // TODO: Should this even be possible? The owner cannot change because
-      // it's forbidden by shouldUpdateReactComponent. The ref can change
-      // if you swap the keys of but not the refs. Reconsider where this check
-      // is made. It probably belongs where the key checking and
-      // instantiateReactComponent is done.
-
-      if (nextElement._owner !== prevElement._owner ||
-          nextElement.ref !== prevElement.ref) {
-        if (prevElement.ref != null) {
-          detachRef(prevElement.ref, this, prevElement._owner);
-        }
-        // Correct, even if the owner is the same, and only the ref has changed.
-        if (nextElement.ref != null) {
-          attachRef(nextElement.ref, this, nextElement._owner);
-        }
-      }
     },
 
     /**

src/core/ReactCompositeComponent.js

@@ -20,6 +20,7 @@ var ReactInstanceMap = require('ReactInstanceMap');
 var ReactPerf = require('ReactPerf');
 var ReactPropTypeLocations = require('ReactPropTypeLocations');
 var ReactPropTypeLocationNames = require('ReactPropTypeLocationNames');
+var ReactRef = require('ReactRef');
 var ReactUpdates = require('ReactUpdates');
 
 var assign = require('Object.assign');
@@ -173,6 +174,8 @@ var ReactCompositeComponentMixin = assign({},
       context
     );
 
+    ReactRef.attachRefs(this, this._currentElement);
+
     this._context = context;
     this._mountOrder = nextMountID++;
     this._rootNodeID = rootID;
@@ -302,6 +305,8 @@ var ReactCompositeComponentMixin = assign({},
     this._pendingCallbacks = null;
     this._pendingElement = null;
 
+    ReactRef.detachRefs(this, this._currentElement);
+
     ReactComponent.Mixin.unmountComponent.call(this);
 
     ReactComponentEnvironment.unmountIDFromEnvironment(this._rootNodeID);
@@ -702,7 +707,6 @@ var ReactCompositeComponentMixin = assign({},
     prevUnmaskedContext,
     nextUnmaskedContext
   ) {
-    // Update refs regardless of what shouldComponentUpdate returns
     ReactComponent.Mixin.updateComponent.call(
       this,
       transaction,
@@ -712,6 +716,9 @@ var ReactCompositeComponentMixin = assign({},
       nextUnmaskedContext
     );
 
+    // Update refs regardless of what shouldComponentUpdate returns
+    ReactRef.updateRefs(this, prevParentElement, nextParentElement);
+
     var inst = this._instance;
 
     var prevContext = inst.context;

src/core/ReactRef.js

@@ -11,6 +11,7 @@
 
 'use strict';
 
+var ReactOwner = require('ReactOwner');
 var ReactUpdates = require('ReactUpdates');
 
 var accumulate = require('accumulate');
@@ -81,16 +82,71 @@ assign(ReactRef.prototype, {
   }
 });
 
-ReactRef.attachRef = function(ref, value) {
+function attachFirstClassRef(ref, value) {
   ref._value = value.getPublicInstance();
-};
+}
 
-ReactRef.detachRef = function(ref, value) {
+function detachFirstClassRef(ref, value) {
   // Check that `component` is still the current ref because we do not want to
   // detach the ref if another component stole it.
   if (ref._value === value) {
     ref._value = null;
   }
+}
+
+function attachRef(ref, component, owner) {
+  if (ref instanceof ReactRef) {
+    attachFirstClassRef(ref, component);
+  } else {
+    ReactOwner.addComponentAsRefTo(component, ref, owner);
+  }
+}
+
+function detachRef(ref, component, owner) {
+  if (ref instanceof ReactRef) {
+    detachFirstClassRef(ref, component);
+  } else {
+    ReactOwner.removeComponentAsRefFrom(component, ref, owner);
+  }
+}
+
+ReactRef.attachRefs = function(instance, element) {
+  var ref = element.ref;
+  if (ref != null) {
+    attachRef(ref, instance, element._owner);
+  }
+};
+
+ReactRef.updateRefs = function(instance, prevElement, nextElement) {
+  // If either the owner or a `ref` has changed, make sure the newest owner
+  // has stored a reference to `this`, and the previous owner (if different)
+  // has forgotten the reference to `this`. We use the element instead
+  // of the public this.props because the post processing cannot determine
+  // a ref. The ref conceptually lives on the element.
+
+  // TODO: Should this even be possible? The owner cannot change because
+  // it's forbidden by shouldUpdateReactComponent. The ref can change
+  // if you swap the keys of but not the refs. Reconsider where this check
+  // is made. It probably belongs where the key checking and
+  // instantiateReactComponent is done.
+
+  if (nextElement._owner !== prevElement._owner ||
+      nextElement.ref !== prevElement.ref) {
+    if (prevElement.ref != null) {
+      detachRef(prevElement.ref, instance, prevElement._owner);
+    }
+    // Correct, even if the owner is the same, and only the ref has changed.
+    if (nextElement.ref != null) {
+      attachRef(nextElement.ref, instance, nextElement._owner);
+    }
+  }
+};
+
+ReactRef.detachRefs = function(instance, element) {
+  var ref = element.ref;
+  if (ref != null) {
+    detachRef(ref, instance, element._owner);
+  }
 };
 
 module.exports = ReactRef;