Allow assigning to this.state inside componentWillMount, with a warning

scripts/fiber/tests-passing.txt

@@ -506,6 +506,7 @@ src/renderers/__tests__/ReactCompositeComponentState-test.js
 * should update state when called from child cWRP
 * should merge state when sCU returns false
 * should treat assigning to this.state inside cWRP as a replaceState, with a warning
+* should treat assigning to this.state inside cWM as a replaceState, with a warning
 
 src/renderers/__tests__/ReactEmptyComponent-test.js
 * should not produce child DOM nodes for null and false

src/renderers/__tests__/ReactCompositeComponentState-test.js

@@ -447,4 +447,44 @@ describe('ReactCompositeComponent-state', () => {
         'Use setState instead.',
     );
   });
+
+  it('should treat assigning to this.state inside cWM as a replaceState, with a warning', () => {
+    spyOn(console, 'error');
+
+    let ops = [];
+    class Test extends React.Component {
+      state = {step: 1, extra: true};
+      componentWillMount() {
+        this.setState({step: 2}, () => {
+          // Tests that earlier setState callbacks are not dropped
+          ops.push(
+            `callback -- step: ${this.state.step}, extra: ${!!this.state.extra}`,
+          );
+        });
+        // Treat like replaceState
+        this.state = {step: 3};
+      }
+      render() {
+        ops.push(
+          `render -- step: ${this.state.step}, extra: ${!!this.state.extra}`,
+        );
+        return null;
+      }
+    }
+
+    // Mount
+    const container = document.createElement('div');
+    ReactDOM.render(<Test />, container);
+
+    expect(ops).toEqual([
+      'render -- step: 3, extra: false',
+      'callback -- step: 3, extra: false',
+    ]);
+    expect(console.error.calls.count()).toEqual(1);
+    expect(console.error.calls.argsFor(0)[0]).toEqual(
+      'Warning: Test.componentWillMount(): Assigning directly to ' +
+        "this.state is deprecated (except inside a component's constructor). " +
+        'Use setState instead.',
+    );
+  });
 });

src/renderers/shared/fiber/ReactFiberClassComponent.js

@@ -304,6 +304,59 @@ module.exports = function(
     return instance;
   }
 
+  function callComponentWillMount(workInProgress, instance) {
+    if (__DEV__) {
+      startPhaseTimer(workInProgress, 'componentWillMount');
+    }
+    const oldState = instance.state;
+    instance.componentWillMount();
+    if (__DEV__) {
+      stopPhaseTimer();
+    }
+
+    if (oldState !== instance.state) {
+      if (__DEV__) {
+        warning(
+          false,
+          '%s.componentWillMount(): Assigning directly to this.state is ' +
+            "deprecated (except inside a component's " +
+            'constructor). Use setState instead.',
+          getComponentName(workInProgress),
+        );
+      }
+      updater.enqueueReplaceState(instance, instance.state, null);
+    }
+  }
+
+  function callComponentWillReceiveProps(
+    workInProgress,
+    instance,
+    newProps,
+    newContext,
+  ) {
+    if (__DEV__) {
+      startPhaseTimer(workInProgress, 'componentWillReceiveProps');
+    }
+    const oldState = instance.state;
+    instance.componentWillReceiveProps(newProps, newContext);
+    if (__DEV__) {
+      stopPhaseTimer();
+    }
+
+    if (instance.state !== oldState) {
+      if (__DEV__) {
+        warning(
+          false,
+          '%s.componentWillReceiveProps(): Assigning directly to ' +
+            "this.state is deprecated (except inside a component's " +
+            'constructor). Use setState instead.',
+          getComponentName(workInProgress),
+        );
+      }
+      updater.enqueueReplaceState(instance, instance.state, null);
+    }
+  }
+
   // Invokes the mount life-cycles on a previously never rendered instance.
   function mountClassInstance(
     workInProgress: Fiber,
@@ -339,13 +392,7 @@ module.exports = function(
     }
 
     if (typeof instance.componentWillMount === 'function') {
-      if (__DEV__) {
-        startPhaseTimer(workInProgress, 'componentWillMount');
-      }
-      instance.componentWillMount();
-      if (__DEV__) {
-        stopPhaseTimer();
-      }
+      callComponentWillMount(workInProgress, instance);
       // If we had additional state updates during this life-cycle, let's
       // process them now.
       const updateQueue = workInProgress.updateQueue;
@@ -365,36 +412,6 @@ module.exports = function(
     }
   }
 
-  function callComponentWillReceiveProps(
-    workInProgress,
-    instance,
-    newProps,
-    newContext,
-  ) {
-    if (typeof instance.componentWillReceiveProps === 'function') {
-      if (__DEV__) {
-        startPhaseTimer(workInProgress, 'componentWillReceiveProps');
-      }
-      instance.componentWillReceiveProps(newProps, newContext);
-      if (__DEV__) {
-        stopPhaseTimer();
-      }
-
-      if (instance.state !== workInProgress.memoizedState) {
-        if (__DEV__) {
-          warning(
-            false,
-            '%s.componentWillReceiveProps(): Assigning directly to ' +
-              "this.state is deprecated (except inside a component's " +
-              'constructor). Use setState instead.',
-            getComponentName(workInProgress),
-          );
-        }
-        updater.enqueueReplaceState(instance, instance.state, null);
-      }
-    }
-  }
-
   // Called on a preexisting class instance. Returns false if a resumed render
   // could be reused.
   function resumeMountClassInstance(
@@ -421,7 +438,11 @@ module.exports = function(
 
     const oldContext = instance.context;
     const oldProps = workInProgress.memoizedProps;
-    if (oldProps !== newProps || oldContext !== newContext) {
+
+    if (
+      typeof instance.componentWillReceiveProps === 'function' &&
+      (oldProps !== newProps || oldContext !== newContext)
+    ) {
       callComponentWillReceiveProps(
         workInProgress,
         instance,
@@ -430,6 +451,19 @@ module.exports = function(
       );
     }
 
+    // Process the update queue before calling shouldComponentUpdate
+    const updateQueue = workInProgress.updateQueue;
+    if (updateQueue !== null) {
+      newState = beginUpdateQueue(
+        workInProgress,
+        updateQueue,
+        instance,
+        newState,
+        newProps,
+        priorityLevel,
+      );
+    }
+
     // TODO: Should we deal with a setState that happened after the last
     // componentWillMount and before this componentWillMount? Probably
     // unsupported anyway.
@@ -452,19 +486,6 @@ module.exports = function(
       return false;
     }
 
-    // componentWillMount may have called setState. Process the update queue.
-    let newUpdateQueue = workInProgress.updateQueue;
-    if (newUpdateQueue !== null) {
-      newState = beginUpdateQueue(
-        workInProgress,
-        newUpdateQueue,
-        instance,
-        newState,
-        newProps,
-        priorityLevel,
-      );
-    }
-
     // Update the input pointers now so that they are correct when we call
     // componentWillMount
     instance.props = newProps;
@@ -472,27 +493,21 @@ module.exports = function(
     instance.context = newContext;
 
     if (typeof instance.componentWillMount === 'function') {
-      if (__DEV__) {
-        startPhaseTimer(workInProgress, 'componentWillMount');
-      }
-      instance.componentWillMount();
-      if (__DEV__) {
-        stopPhaseTimer();
+      callComponentWillMount(workInProgress, instance);
+      // componentWillMount may have called setState. Process the update queue.
+      const newUpdateQueue = workInProgress.updateQueue;
+      if (newUpdateQueue !== null) {
+        newState = beginUpdateQueue(
+          workInProgress,
+          updateQueue,
+          instance,
+          newState,
+          newProps,
+          priorityLevel,
+        );
       }
     }
 
-    // componentWillMount may have called setState. Process the update queue.
-    newUpdateQueue = workInProgress.updateQueue;
-    if (newUpdateQueue !== null) {
-      newState = beginUpdateQueue(
-        workInProgress,
-        newUpdateQueue,
-        instance,
-        newState,
-        newProps,
-        priorityLevel,
-      );
-    }
     if (typeof instance.componentDidMount === 'function') {
       workInProgress.effectTag |= Update;
     }
@@ -531,7 +546,10 @@ module.exports = function(
     // ever the previously attempted to render - not the "current". However,
     // during componentDidUpdate we pass the "current" props.
 
-    if (oldProps !== newProps || oldContext !== newContext) {
+    if (
+      typeof instance.componentWillReceiveProps === 'function' &&
+      (oldProps !== newProps || oldContext !== newContext)
+    ) {
       callComponentWillReceiveProps(
         workInProgress,
         instance,