Batch nested updates when in sync mode

Almost working...

src/renderers/dom/fiber/ReactDOMFiber.js

@@ -112,6 +112,8 @@ var DOMRenderer = ReactFiberReconciler({
 
   scheduleDeferredCallback: window.requestIdleCallback,
 
+  useSyncScheduling: true,
+
 });
 
 var warned = false;

src/renderers/shared/fiber/ReactFiberClassComponent.js

@@ -130,7 +130,7 @@ module.exports = function(scheduleUpdate : (fiber: Fiber, priorityLevel : ?Prior
       // process them now.
       const updateQueue = workInProgress.updateQueue;
       if (updateQueue) {
-        instance.state = mergeUpdateQueue(updateQueue, state, props);
+        instance.state = mergeUpdateQueue(updateQueue, instance, state, props);
       }
     }
   }
@@ -174,7 +174,7 @@ module.exports = function(scheduleUpdate : (fiber: Fiber, priorityLevel : ?Prior
       // process them now.
       const newUpdateQueue = workInProgress.updateQueue;
       if (newUpdateQueue) {
-        newInstance.state = mergeUpdateQueue(newUpdateQueue, newState, newProps);
+        newInstance.state = mergeUpdateQueue(newUpdateQueue, newInstance, newState, newProps);
       }
     }
     return true;
@@ -211,7 +211,7 @@ module.exports = function(scheduleUpdate : (fiber: Fiber, priorityLevel : ?Prior
     // TODO: Previous state can be null.
     let newState;
     if (updateQueue) {
-      newState = mergeUpdateQueue(updateQueue, previousState, newProps);
+      newState = mergeUpdateQueue(updateQueue, instance, previousState, newProps);
     } else {
       newState = previousState;
     }

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -71,6 +71,9 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
     SynchronousPriority :
     LowPriority;
 
+  // Whether updates should be batched. Only applies when using sync scheduling.
+  let shouldBatchUpdates : boolean = false;
+
   // The next work in progress fiber that we're currently working on.
   let nextUnitOfWork : ?Fiber = null;
   let nextPriorityLevel : PriorityLevel = NoWork;
@@ -482,9 +485,10 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
     while (nextUnitOfWork &&
            nextPriorityLevel === SynchronousPriority) {
       nextUnitOfWork = performUnitOfWork(nextUnitOfWork, false);
-      // If there's no nextUnitOfWork, we don't need to search for more
-      // because it shouldn't be possible to schedule sync work without
-      // immediately performing it
+
+      if (!nextUnitOfWork && shouldBatchUpdates) {
+        nextUnitOfWork = findNextUnitOfWork();
+      }
     }
     if (nextUnitOfWork) {
       if (nextPriorityLevel > AnimationPriority) {
@@ -497,6 +501,16 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
 
   function performSynchronousWork() {
     performAndHandleErrors(performSynchronousWorkUnsafe);
+    shouldBatchUpdates = false;
+  }
+
+  function scheduleSynchronousWork(root : FiberRoot) {
+    root.current.pendingWorkPriority = SynchronousPriority;
+    // Unless in batch mode, perform the work immediately. Otherwise, the work
+    // will be flushed at the end of the batch.
+    if (!shouldBatchUpdates) {
+      performSynchronousWork();
+    }
   }
 
   function performAndHandleErrors<A>(fn: (a: A) => void, a: A) {
@@ -552,10 +566,9 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
 
       // We will process an update caused by each error boundary synchronously.
       affectedBoundaries.forEach(boundary => {
-        // FIXME: We only specify LowPriority here so that setState() calls from the error
-        // boundaries are respected. Instead we should set default priority level or something
-        // like this. Reconsider this piece when synchronous scheduling is in place.
-        const priority = LowPriority;
+        const priority = priorityContext !== null ?
+          priorityContext :
+          defaultPriorityContext;
         const root = scheduleErrorBoundaryWork(boundary, priority);
         // This should use findNextUnitOfWork() when synchronous scheduling is implemented.
         let fiber = cloneFiber(root.current, priority);
@@ -594,8 +607,7 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
     }
 
     if (priorityLevel === SynchronousPriority) {
-      root.current.pendingWorkPriority = SynchronousPriority;
-      performSynchronousWork();
+      scheduleSynchronousWork(root);
     }
 
     if (priorityLevel === NoWork) {
@@ -616,13 +628,9 @@ module.exports = function<T, P, I, TI, C>(config : HostConfig<T, P, I, TI, C>) {
         defaultPriorityContext;
     }
 
-    // Don't bother bubbling the priority to the root if it is synchronous. Just
-    // perform it now.
-    if (priorityLevel === SynchronousPriority) {
-      fiber.pendingWorkPriority = SynchronousPriority;
-      nextUnitOfWork = fiber;
-      performSynchronousWork();
-      return;
+    if (priorityLevel === SynchronousPriority && useSyncScheduling) {
+      // Start batching updates
+      shouldBatchUpdates = true;
     }
 
     while (true) {

src/renderers/shared/fiber/ReactFiberUpdateQueue.js

@@ -73,14 +73,14 @@ exports.callCallbacks = function(queue : UpdateQueue, context : any) {
   }
 };
 
-exports.mergeUpdateQueue = function(queue : UpdateQueue, prevState : any, props : any) : any {
+exports.mergeUpdateQueue = function(queue : UpdateQueue, instance : any, prevState : any, props : any) : any {
   let node : ?UpdateQueueNode = queue;
   let state = queue.isReplace ? null : Object.assign({}, prevState);
   while (node) {
     let partialState;
     if (typeof node.partialState === 'function') {
       const updateFn = node.partialState;
-      partialState = updateFn(state, props);
+      partialState = updateFn.call(instance, state, props);
     } else {
       partialState = node.partialState;
     }