Schedule state and callback at the same time

src/isomorphic/classic/class/ReactClass.js

@@ -728,10 +728,7 @@ var ReactClassMixin = {
    * type signature and the only use case for this, is to avoid that.
    */
   replaceState: function(newState, callback) {
-    this.updater.enqueueReplaceState(this, newState);
-    if (callback) {
-      this.updater.enqueueCallback(this, callback, 'replaceState');
-    }
+    this.updater.enqueueReplaceState(this, newState, callback, 'replaceState');
   },
 
   /**

src/isomorphic/modern/class/ReactComponent.js

@@ -65,10 +65,7 @@ ReactComponent.prototype.setState = function(partialState, callback) {
     'setState(...): takes an object of state variables to update or a ' +
     'function which returns an object of state variables.'
   );
-  this.updater.enqueueSetState(this, partialState);
-  if (callback) {
-    this.updater.enqueueCallback(this, callback, 'setState');
-  }
+  this.updater.enqueueSetState(this, partialState, callback, 'setState');
 };
 
 /**
@@ -86,10 +83,7 @@ ReactComponent.prototype.setState = function(partialState, callback) {
  * @protected
  */
 ReactComponent.prototype.forceUpdate = function(callback) {
-  this.updater.enqueueForceUpdate(this);
-  if (callback) {
-    this.updater.enqueueCallback(this, callback, 'forceUpdate');
-  }
+  this.updater.enqueueForceUpdate(this, callback, 'forceUpdate');
 };
 
 /**

src/isomorphic/modern/class/ReactNoopUpdateQueue.js

@@ -44,16 +44,6 @@ var ReactNoopUpdateQueue = {
     return false;
   },
 
-  /**
-   * Enqueue a callback that will be executed after all the pending updates
-   * have processed.
-   *
-   * @param {ReactClass} publicInstance The instance to use as `this` context.
-   * @param {?function} callback Called after state is updated.
-   * @internal
-   */
-  enqueueCallback: function(publicInstance, callback) { },
-
   /**
    * Forces an update. This should only be invoked when it is known with
    * certainty that we are **not** in a DOM transaction.
@@ -65,9 +55,11 @@ var ReactNoopUpdateQueue = {
    * `componentWillUpdate` and `componentDidUpdate`.
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueForceUpdate: function(publicInstance) {
+  enqueueForceUpdate: function(publicInstance, callback, callerName) {
     warnNoop(publicInstance, 'forceUpdate');
   },
 
@@ -80,9 +72,11 @@ var ReactNoopUpdateQueue = {
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
    * @param {object} completeState Next state.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueReplaceState: function(publicInstance, completeState) {
+  enqueueReplaceState: function(publicInstance, completeState, callback, callerName) {
     warnNoop(publicInstance, 'replaceState');
   },
 
@@ -94,9 +88,11 @@ var ReactNoopUpdateQueue = {
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
    * @param {object} partialState Next partial state to be merged with state.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueSetState: function(publicInstance, partialState) {
+  enqueueSetState: function(publicInstance, partialState, callback, callerName) {
     warnNoop(publicInstance, 'setState');
   },
 };

src/renderers/dom/stack/server/ReactServerUpdateQueue.js

@@ -58,20 +58,6 @@ class ReactServerUpdateQueue {
     return false;
   }
 
-  /**
-   * Enqueue a callback that will be executed after all the pending updates
-   * have processed.
-   *
-   * @param {ReactClass} publicInstance The instance to use as `this` context.
-   * @param {?function} callback Called after state is updated.
-   * @internal
-   */
-  enqueueCallback(publicInstance: ReactComponent<any, any, any>, callback?: Function, callerName?: string) {
-    if (this.transaction.isInTransaction()) {
-      ReactUpdateQueue.enqueueCallback(publicInstance, callback, callerName);
-    }
-  }
-
   /**
    * Forces an update. This should only be invoked when it is known with
    * certainty that we are **not** in a DOM transaction.
@@ -83,11 +69,13 @@ class ReactServerUpdateQueue {
    * `componentWillUpdate` and `componentDidUpdate`.
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueForceUpdate(publicInstance: ReactComponent<any, any, any>) {
+  enqueueForceUpdate(publicInstance: ReactComponent<any, any, any>, callback?: Function, callerName?: string) {
     if (this.transaction.isInTransaction()) {
-      ReactUpdateQueue.enqueueForceUpdate(publicInstance);
+      ReactUpdateQueue.enqueueForceUpdate(publicInstance, callback, callerName);
     } else {
       warnNoop(publicInstance, 'forceUpdate');
     }
@@ -102,11 +90,18 @@ class ReactServerUpdateQueue {
    *
    * @param {ReactClass} publicInstance The instance that should rerender.
    * @param {object|function} completeState Next state.
+   * @param {?function} callback Called after component is updated.
+   * @param {?string} Name of the calling function in the public API.
    * @internal
    */
-  enqueueReplaceState(publicInstance: ReactComponent<any, any, any>, completeState: Object|Function) {
+  enqueueReplaceState(
+    publicInstance: ReactComponent<any, any, any>,
+    completeState: Object|Function,
+    callback?: Function,
+    callerName?: string
+  ) {
     if (this.transaction.isInTransaction()) {
-      ReactUpdateQueue.enqueueReplaceState(publicInstance, completeState);
+      ReactUpdateQueue.enqueueReplaceState(publicInstance, completeState, callback, callerName);
     } else {
       warnNoop(publicInstance, 'replaceState');
     }
@@ -122,9 +117,14 @@ class ReactServerUpdateQueue {
    * @param {object|function} partialState Next partial state to be merged with state.
    * @internal
    */
-  enqueueSetState(publicInstance: ReactComponent<any, any, any>, partialState: Object|Function) {
+  enqueueSetState(
+    publicInstance: ReactComponent<any, any, any>,
+    partialState: Object|Function,
+    callback?: Function,
+    callerName?: string
+  ) {
     if (this.transaction.isInTransaction()) {
-      ReactUpdateQueue.enqueueSetState(publicInstance, partialState);
+      ReactUpdateQueue.enqueueSetState(publicInstance, partialState, callback, callerName);
     } else {
       warnNoop(publicInstance, 'setState');
     }

src/renderers/shared/fiber/ReactFiberBeginWork.js

@@ -71,10 +71,8 @@ if (__DEV__) {
 module.exports = function<T, P, I, TI, C, CX>(
   config : HostConfig<T, P, I, TI, C, CX>,
   hostContext : HostContext<C, CX>,
-  scheduleSetState: (fiber : Fiber, partialState : any) => void,
-  scheduleReplaceState: (fiber : Fiber, state : any) => void,
-  scheduleForceUpdate: (fiber : Fiber) => void,
-  scheduleUpdateCallback: (fiber : Fiber, callback : Function) => void,
+  scheduleUpdate : (fiber : Fiber, priorityLevel : PriorityLevel) => void,
+  getPriorityContext : () => PriorityLevel,
 ) {
 
   const { shouldSetTextContent } = config;
@@ -91,12 +89,7 @@ module.exports = function<T, P, I, TI, C, CX>(
     mountClassInstance,
     resumeMountClassInstance,
     updateClassInstance,
-  } = ReactFiberClassComponent(
-    scheduleSetState,
-    scheduleReplaceState,
-    scheduleForceUpdate,
-    scheduleUpdateCallback
-  );
+  } = ReactFiberClassComponent(scheduleUpdate, getPriorityContext);
 
   function markChildAsProgressed(current, workInProgress, priorityLevel) {
     // We now have clones. Let's store them as the currently progressed work.

src/renderers/shared/fiber/ReactFiberClassComponent.js

@@ -19,6 +19,9 @@ var {
   getMaskedContext,
 } = require('ReactFiberContext');
 var {
+  addUpdate,
+  addReplaceUpdate,
+  addForceUpdate,
   beginUpdateQueue,
 } = require('ReactFiberUpdateQueue');
 var { hasContextChanged } = require('ReactFiberContext');
@@ -31,30 +34,30 @@ var invariant = require('invariant');
 const isArray = Array.isArray;
 
 module.exports = function(
-  scheduleSetState: (fiber : Fiber, partialState : any) => void,
-  scheduleReplaceState: (fiber : Fiber, state : any) => void,
-  scheduleForceUpdate: (fiber : Fiber) => void,
-  scheduleUpdateCallback: (fiber : Fiber, callback : Function) => void,
+  scheduleUpdate : (fiber : Fiber, priorityLevel : PriorityLevel) => void,
+  getPriorityContext : () => PriorityLevel,
 ) {
 
   // Class component state updater
   const updater = {
     isMounted,
-    enqueueSetState(instance, partialState) {
+    enqueueSetState(instance, partialState, callback) {
       const fiber = ReactInstanceMap.get(instance);
-      scheduleSetState(fiber, partialState);
+      const priorityLevel = getPriorityContext();
+      addUpdate(fiber, partialState, callback || null, priorityLevel);
+      scheduleUpdate(fiber, priorityLevel);
     },
-    enqueueReplaceState(instance, state) {
+    enqueueReplaceState(instance, state, callback) {
       const fiber = ReactInstanceMap.get(instance);
-      scheduleReplaceState(fiber, state);
+      const priorityLevel = getPriorityContext();
+      addReplaceUpdate(fiber, state, callback || null, priorityLevel);
+      scheduleUpdate(fiber, priorityLevel);
     },
-    enqueueForceUpdate(instance) {
+    enqueueForceUpdate(instance, callback) {
       const fiber = ReactInstanceMap.get(instance);
-      scheduleForceUpdate(fiber);
-    },
-    enqueueCallback(instance, callback) {
-      const fiber = ReactInstanceMap.get(instance);
-      scheduleUpdateCallback(fiber, callback);
+      const priorityLevel = getPriorityContext();
+      addForceUpdate(fiber, callback || null, priorityLevel);
+      scheduleUpdate(fiber, priorityLevel);
     },
   };
 

src/renderers/shared/fiber/ReactFiberReconciler.js

@@ -17,6 +17,10 @@ import type { FiberRoot } from 'ReactFiberRoot';
 import type { PriorityLevel } from 'ReactPriorityLevel';
 import type { ReactNodeList } from 'ReactTypes';
 
+var {
+  addTopLevelUpdate,
+} = require('ReactFiberUpdateQueue');
+
 var {
   findCurrentUnmaskedContext,
   isContextProvider,
@@ -98,25 +102,29 @@ getContextForSubtree._injectFiber(function(fiber : Fiber) {
 module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C, CX>) : Reconciler<C, I, TI> {
 
   var {
-    scheduleTopLevelSetState,
-    scheduleUpdateCallback,
+    scheduleUpdate,
+    getPriorityContext,
     performWithPriority,
     batchedUpdates,
     syncUpdates,
     deferredUpdates,
   } = ReactFiberScheduler(config);
 
+  function scheduleTopLevelUpdate(current : Fiber, element : ReactNodeList, callback : ?Function) {
+    const priorityLevel = getPriorityContext();
+    const nextState = { element };
+    addTopLevelUpdate(current, nextState, callback || null, priorityLevel);
+    scheduleUpdate(current, priorityLevel);
+  }
+
   return {
 
     mountContainer(element : ReactNodeList, containerInfo : C, parentComponent : ?ReactComponent<any, any, any>, callback: ?Function) : OpaqueNode {
       const context = getContextForSubtree(parentComponent);
       const root = createFiberRoot(containerInfo, context);
       const current = root.current;
 
-      scheduleTopLevelSetState(current, { element });
-      if (callback) {
-        scheduleUpdateCallback(current, callback);
-      }
+      scheduleTopLevelUpdate(current, element, callback);
 
       if (__DEV__ && ReactFiberInstrumentation.debugTool) {
         ReactFiberInstrumentation.debugTool.onMountContainer(root);
@@ -135,10 +143,7 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
 
       root.pendingContext = getContextForSubtree(parentComponent);
 
-      scheduleTopLevelSetState(current, { element });
-      if (callback) {
-        scheduleUpdateCallback(current, callback);
-      }
+      scheduleTopLevelUpdate(current, element, callback);
 
       if (__DEV__) {
         if (ReactFiberInstrumentation.debugTool) {

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -55,11 +55,6 @@ var {
 
 var {
   getPendingPriority,
-  addUpdate,
-  addReplaceUpdate,
-  addForceUpdate,
-  addCallback,
-  addTopLevelUpdate,
 } = require('ReactFiberUpdateQueue');
 
 var {
@@ -79,10 +74,8 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
   const { beginWork, beginFailedWork } = ReactFiberBeginWork(
     config,
     hostContext,
-    scheduleSetState,
-    scheduleReplaceState,
-    scheduleForceUpdate,
-    scheduleUpdateCallback,
+    scheduleUpdate,
+    getPriorityContext,
   );
   const { completeWork } = ReactFiberCompleteWork(config, hostContext);
   const {
@@ -986,7 +979,7 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
     }
   }
 
-  function scheduleUpdateAtPriority(fiber : Fiber, priorityLevel : PriorityLevel) {
+  function scheduleUpdate(fiber : Fiber, priorityLevel : PriorityLevel) {
     // If we're in a batch, downgrade sync priority to task priority
     if (priorityLevel === SynchronousPriority && isPerformingWork) {
       priorityLevel = TaskPriority;
@@ -1049,34 +1042,15 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
     }
   }
 
-  function scheduleErrorRecovery(fiber : Fiber) {
-    scheduleUpdateAtPriority(fiber, TaskPriority);
-  }
-
-  function scheduleSetState(fiber : Fiber, partialState : any) {
-    addUpdate(fiber, partialState, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
-  }
-
-  function scheduleReplaceState(fiber : Fiber, state : any) {
-    addReplaceUpdate(fiber, state, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
-  }
-
-  function scheduleForceUpdate(fiber : Fiber) {
-    addForceUpdate(fiber, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
-  }
-
-  function scheduleUpdateCallback(fiber : Fiber, callback : Function) {
-    addCallback(fiber, callback, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
+  function getPriorityContext() : PriorityLevel {
+    if (priorityContext === SynchronousPriority && isPerformingWork) {
+      return TaskPriority;
+    }
+    return priorityContext;
   }
 
-  // TODO: This indirection will be removed as part of #8585
-  function scheduleTopLevelSetState(fiber : Fiber, partialState : any) {
-    addTopLevelUpdate(fiber, partialState, priorityContext);
-    scheduleUpdateAtPriority(fiber, priorityContext);
+  function scheduleErrorRecovery(fiber : Fiber) {
+    scheduleUpdate(fiber, TaskPriority);
   }
 
   function performWithPriority(priorityLevel : PriorityLevel, fn : Function) {
@@ -1126,8 +1100,8 @@ module.exports = function<T, P, I, TI, C, CX>(config : HostConfig<T, P, I, TI, C
   }
 
   return {
-    scheduleTopLevelSetState: scheduleTopLevelSetState,
-    scheduleUpdateCallback: scheduleUpdateCallback,
+    scheduleUpdate: scheduleUpdate,
+    getPriorityContext: getPriorityContext,
     performWithPriority: performWithPriority,
     batchedUpdates: batchedUpdates,
     syncUpdates: syncUpdates,