Make sure top-level callback has correct context

Closes #5125 -- incorporates essentially the same fix with a much simpler test case. Thanks @yiminghe for reporting!

src/renderers/dom/client/ReactMount.js

@@ -610,12 +610,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

@@ -299,4 +299,44 @@ describe('ReactMount', function() {
     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);
+  });
 });