Move ref code to ReactCompositeComponent

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;