Call setState callbacks enqueued in cWM

src/renderers/dom/stack/client/ReactMount.js

@@ -357,7 +357,8 @@ var ReactMount = {
     nextElement,
     container,
     shouldReuseMarkup,
-    context
+    context,
+    callback
   ) {
     // Various parts of our code (such as ReactCompositeComponent's
     // _renderValidatedComponent) assume that calls to render aren't nested;
@@ -380,6 +381,12 @@ var ReactMount = {
     ReactBrowserEventEmitter.ensureScrollValueMonitoring();
     var componentInstance = instantiateReactComponent(nextElement, false);
 
+    if (callback) {
+      componentInstance._pendingCallbacks = [function() {
+        callback.call(componentInstance._renderedComponent.getPublicInstance());
+      }];
+    }
+
     // The initial render is synchronous but any updates that happen during
     // rendering, in componentWillMount or componentDidMount, will be batched
     // according to the current batching strategy.
@@ -529,11 +536,9 @@ var ReactMount = {
       nextWrappedElement,
       container,
       shouldReuseMarkup,
-      nextContext
+      nextContext,
+      callback
     )._renderedComponent.getPublicInstance();
-    if (callback) {
-      callback.call(component);
-    }
     return component;
   },
 

src/renderers/native/ReactNativeMount.js

@@ -134,6 +134,12 @@ var ReactNativeMount = {
     var instance = instantiateReactComponent(nextWrappedElement, false);
     ReactNativeMount._instancesByContainerID[containerTag] = instance;
 
+    if (callback) {
+      instance._pendingCallbacks = [function() {
+        callback.call(instance._renderedComponent.getPublicInstance());
+      }];
+    }
+
     // The initial render is synchronous but any updates that happen during
     // rendering, in componentWillMount or componentDidMount, will be batched
     // according to the current batching strategy.
@@ -143,10 +149,7 @@ var ReactNativeMount = {
       instance,
       containerTag
     );
-    var component = instance.getPublicInstance();
-    if (callback) {
-      callback.call(component);
-    }
+    var component = instance._renderedComponent.getPublicInstance();
     return component;
   },
 

src/renderers/native/__tests__/ReactNativeMount-test.js

@@ -58,4 +58,25 @@ describe('ReactNative', () => {
     expect(UIManager.updateView).toBeCalledWith(3, 'View', { foo: 'bar' });
   });
 
+  it('should be able to create and update a native component', () => {
+    var View = createReactNativeComponentClass({
+      validAttributes: { foo: true },
+      uiViewClassName: 'View',
+    });
+
+    var a;
+    var b;
+    var c = ReactNative.render(
+      <View foo="foo" ref={(v) => a = v} />,
+      11,
+      function() {
+        b = this;
+      }
+    );
+
+    expect(a).toBeTruthy();
+    expect(a).toBe(b);
+    expect(a).toBe(c);
+  });
+
 });

src/renderers/shared/stack/reconciler/ReactCompositeComponent.js

@@ -383,6 +383,18 @@ var ReactCompositeComponent = {
       }
     }
 
+    // setState callbacks during willMount should end up here
+    const callbacks = this._pendingCallbacks;
+    if (callbacks) {
+      this._pendingCallbacks = null;
+      for (let i = 0; i < callbacks.length; i++) {
+        transaction.getReactMountReady().enqueue(
+          callbacks[i],
+          inst
+        );
+      }
+    }
+
     return markup;
   },
 

src/renderers/shared/stack/reconciler/__tests__/ReactCompositeComponentState-test.js

@@ -11,8 +11,6 @@
 
 'use strict';
 
-var ReactDOMFeatureFlags = require('ReactDOMFeatureFlags');
-
 var React;
 var ReactDOM;
 
@@ -171,40 +169,22 @@ describe('ReactCompositeComponent-state', () => {
       ['componentDidMount-end', 'orange'],
     ];
 
-    if (ReactDOMFeatureFlags.useFiber) {
-      // The setState callbacks in componentWillMount, and the initial callback
-      // passed to ReactDOM.render, should be flushed right after component
-      // did mount:
-      expected.push(
-        ['setState-sunrise', 'orange'], // 1
-        ['setState-orange', 'orange'], // 2
-        ['initial-callback', 'orange'], // 3
-        ['shouldComponentUpdate-currentState', 'orange'],
-        ['shouldComponentUpdate-nextState', 'yellow'],
-        ['componentWillUpdate-currentState', 'orange'],
-        ['componentWillUpdate-nextState', 'yellow'],
-        ['render', 'yellow'],
-        ['componentDidUpdate-currentState', 'yellow'],
-        ['componentDidUpdate-prevState', 'orange'],
-        ['setState-yellow', 'yellow'],
-      );
-    } else {
-      // There is a bug in the stack reconciler where those callbacks are
-      // enqueued, but aren't called until the next flush.
-      expected.push(
-        ['shouldComponentUpdate-currentState', 'orange'],
-        ['shouldComponentUpdate-nextState', 'yellow'],
-        ['componentWillUpdate-currentState', 'orange'],
-        ['componentWillUpdate-nextState', 'yellow'],
-        ['render', 'yellow'],
-        ['componentDidUpdate-currentState', 'yellow'],
-        ['componentDidUpdate-prevState', 'orange'],
-        ['setState-sunrise', 'yellow'], // 1
-        ['setState-orange', 'yellow'], // 2
-        ['setState-yellow', 'yellow'],
-        ['initial-callback', 'yellow'] // 3
-      );
-    }
+    // The setState callbacks in componentWillMount, and the initial callback
+    // passed to ReactDOM.render, should be flushed right after component
+    // did mount:
+    expected.push(
+      ['setState-sunrise', 'orange'], // 1
+      ['setState-orange', 'orange'], // 2
+      ['initial-callback', 'orange'], // 3
+      ['shouldComponentUpdate-currentState', 'orange'],
+      ['shouldComponentUpdate-nextState', 'yellow'],
+      ['componentWillUpdate-currentState', 'orange'],
+      ['componentWillUpdate-nextState', 'yellow'],
+      ['render', 'yellow'],
+      ['componentDidUpdate-currentState', 'yellow'],
+      ['componentDidUpdate-prevState', 'orange'],
+      ['setState-yellow', 'yellow'],
+    );
 
     expected.push(
       ['componentWillReceiveProps-start', 'yellow'],