Fixed descrepancy between host and class component refs (facebook#12178)

When a ref is removed from a class component, React now calls the previous ref-setter (if there was one) with null. Previously this was the case only for host component refs.

A new test has been added.

Author: bvaughn

packages/react-dom/src/__tests__/refs-destruction-test.js

@@ -23,15 +23,31 @@ describe('refs-destruction', () => {
     ReactDOM = require('react-dom');
     ReactTestUtils = require('react-dom/test-utils');
 
+    class ClassComponent extends React.Component {
+      render() {
+        return null;
+      }
+    }
+
     TestComponent = class extends React.Component {
       render() {
-        return (
-          <div>
-            {this.props.destroy ? null : (
-              <div ref="theInnerDiv">Lets try to destroy this.</div>
-            )}
-          </div>
-        );
+        if (this.props.destroy) {
+          return <div />;
+        } else if (this.props.removeRef) {
+          return (
+            <div>
+              <div />
+              <ClassComponent />
+            </div>
+          );
+        } else {
+          return (
+            <div>
+              <div ref="theInnerDiv" />
+              <ClassComponent ref="theInnerClassComponent" />
+            </div>
+          );
+        }
       }
     };
   });
@@ -45,7 +61,7 @@ describe('refs-destruction', () => {
     expect(
       Object.keys(testInstance.refs || {}).filter(key => testInstance.refs[key])
         .length,
-    ).toEqual(1);
+    ).toEqual(2);
     ReactDOM.unmountComponentAtNode(container);
     expect(
       Object.keys(testInstance.refs || {}).filter(key => testInstance.refs[key])
@@ -62,14 +78,31 @@ describe('refs-destruction', () => {
     expect(
       Object.keys(testInstance.refs || {}).filter(key => testInstance.refs[key])
         .length,
-    ).toEqual(1);
+    ).toEqual(2);
     ReactDOM.render(<TestComponent destroy={true} />, container);
     expect(
       Object.keys(testInstance.refs || {}).filter(key => testInstance.refs[key])
         .length,
     ).toEqual(0);
   });
 
+  it('should remove refs when removing the child ref attribute', () => {
+    const container = document.createElement('div');
+    const testInstance = ReactDOM.render(<TestComponent />, container);
+    expect(ReactTestUtils.isDOMComponent(testInstance.refs.theInnerDiv)).toBe(
+      true,
+    );
+    expect(
+      Object.keys(testInstance.refs || {}).filter(key => testInstance.refs[key])
+        .length,
+    ).toEqual(2);
+    ReactDOM.render(<TestComponent removeRef={true} />, container);
+    expect(
+      Object.keys(testInstance.refs || {}).filter(key => testInstance.refs[key])
+        .length,
+    ).toEqual(0);
+  });
+
   it('should not error when destroying child with ref asynchronously', () => {
     class Modal extends React.Component {
       componentDidMount() {

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -184,7 +184,10 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
 
   function markRef(current: Fiber | null, workInProgress: Fiber) {
     const ref = workInProgress.ref;
-    if (ref !== null && (!current || current.ref !== ref)) {
+    if (
+      (current === null && ref !== null) ||
+      (current !== null && current.ref !== ref)
+    ) {
       // Schedule a Ref effect
       workInProgress.effectTag |= Ref;
     }