Merge pull request facebook#5237 from spicyj/facebookgh-5125

Make sure top-level callback has correct context
(cherry picked from commit b0a7a00)

Author: sophiebits

src/renderers/dom/client/ReactMount.js

@@ -635,12 +635,17 @@ var ReactMount = {
       var prevWrappedElement = prevComponent._currentElement;
       var prevElement = prevWrappedElement.props;
       if (shouldUpdateReactComponent(prevElement, nextElement)) {
-        return ReactMount._updateRootComponent(
+        var publicInst = prevComponent._renderedComponent.getPublicInstance();
+        var updatedCallback = callback && function() {
+          callback.call(publicInst);
+        };
+        ReactMount._updateRootComponent(
           prevComponent,
           nextWrappedElement,
           container,
-          callback
-        )._renderedComponent.getPublicInstance();
+          updatedCallback
+        );
+        return publicInst;
       } else {
         ReactMount.unmountComponentAtNode(container);
       }

src/renderers/dom/client/__tests__/ReactMount-test.js

@@ -245,4 +245,99 @@ describe('ReactMount', function() {
       'render the new components instead of calling ReactDOM.render.'
     );
   });
+
+  it('should not crash in node cache when unmounting', function() {
+    var Component = React.createClass({
+      render: function() {
+        // Add refs to some nodes so that they get traversed and cached
+        return (
+          <div ref="a">
+            <div ref="b">b</div>
+            {this.props.showC && <div>c</div>}
+          </div>
+        );
+      },
+    });
+
+    var container = document.createElement('container');
+
+    ReactDOM.render(<div><Component showC={false} /></div>, container);
+
+    // Right now, A and B are in the cache. When we add C, it won't get added to
+    // the cache (assuming markup-string mode).
+    ReactDOM.render(<div><Component showC={true} /></div>, container);
+
+    // Remove A, B, and C. Unmounting C shouldn't cause B to get recached.
+    ReactDOM.render(<div></div>, container);
+
+    // Add them back -- this shouldn't cause a cached node collision.
+    ReactDOM.render(<div><Component showC={true} /></div>, container);
+
+    ReactDOM.unmountComponentAtNode(container);
+  });
+
+  it('should not crash in node cache when unmounting, case 2', function() {
+    var A = React.createClass({
+      render: function() {
+        return <a key={this.props.innerKey}>{this.props.innerKey}</a>;
+      },
+    });
+    var Component = React.createClass({
+      render: function() {
+        return (
+          <b>
+            <i>{this.props.step === 1 && <q />}</i>
+            {this.props.step === 1 && <A innerKey={this.props.step} />}
+          </b>
+        );
+      },
+    });
+
+    var container = document.createElement('container');
+
+    ReactDOM.render(<Component step={1} />, container);
+    ReactDOM.render(<Component step={2} />, container);
+    ReactDOM.render(<Component step={1} />, container);
+    ReactMount.getID(container.querySelector('a'));
+  });
+
+  it('passes the correct callback context', function() {
+    var container = document.createElement('div');
+    var calls = 0;
+
+    ReactDOM.render(<div />, container, function() {
+      expect(this.nodeName).toBe('DIV');
+      calls++;
+    });
+
+    // Update, no type change
+    ReactDOM.render(<div />, container, function() {
+      expect(this.nodeName).toBe('DIV');
+      calls++;
+    });
+
+    // Update, type change
+    ReactDOM.render(<span />, container, function() {
+      expect(this.nodeName).toBe('SPAN');
+      calls++;
+    });
+
+    // Batched update, no type change
+    ReactDOM.unstable_batchedUpdates(function() {
+      ReactDOM.render(<span />, container, function() {
+        expect(this.nodeName).toBe('SPAN');
+        calls++;
+      });
+    });
+
+    // Batched update, type change
+    ReactDOM.unstable_batchedUpdates(function() {
+      ReactDOM.render(<article />, container, function() {
+        expect(this.nodeName).toBe('ARTICLE');
+        calls++;
+      });
+    });
+
+    expect(calls).toBe(5);
+  });
 });