Completion callbacks resolve synchronously if tree is already complete

More unit tests. These completion callbacks (as I'm calling them) have
some interesting properties.

src/renderers/noop/ReactNoopEntry.js

@@ -26,6 +26,7 @@ var ReactFiberInstrumentation = require('ReactFiberInstrumentation');
 var ReactFiberReconciler = require('ReactFiberReconciler');
 var ReactInstanceMap = require('ReactInstanceMap');
 var emptyObject = require('fbjs/lib/emptyObject');
+var invariant = require('fbjs/lib/invariant');
 
 var expect = require('jest-matchers');
 
@@ -285,6 +286,27 @@ var ReactNoop = {
     }
   },
 
+  create(rootID: string) {
+    rootID = typeof rootID === 'string' ? rootID : DEFAULT_ROOT_ID;
+    invariant(
+      !roots.has(rootID),
+      'Root with id %s already exists. Choose a different id.',
+      rootID,
+    );
+    const container = {rootID: rootID, children: []};
+    rootContainers.set(rootID, container);
+    const root = NoopRenderer.createContainer(container);
+    roots.set(rootID, root);
+    return {
+      prerender(children: any) {
+        return NoopRenderer.updateRoot(children, root, null);
+      },
+      getChildren() {
+        return ReactNoop.getChildren(rootID);
+      },
+    };
+  },
+
   findInstance(
     componentOrElement: Element | ?React$Component<any, any>,
   ): null | Instance | TextInstance {

src/renderers/shared/fiber/ReactFiberReconciler.js

@@ -14,7 +14,6 @@
 
 import type {Fiber} from 'ReactFiber';
 import type {FiberRoot} from 'ReactFiberRoot';
-import type {PriorityLevel} from 'ReactPriorityLevel';
 import type {ExpirationTime} from 'ReactFiberExpirationTime';
 import type {ReactNodeList} from 'ReactTypes';
 
@@ -218,6 +217,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
 
   var {
     scheduleUpdate,
+    scheduleCompletionCallback,
     getPriorityContext,
     getExpirationTimeForPriority,
     recalculateCurrentTime,
@@ -315,7 +315,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
       if (root.blockers === null) {
         root.blockers = createUpdateQueue();
       }
-      const blockUpdate = {
+      const block = {
         priorityLevel: null,
         expirationTime,
         partialState: null,
@@ -325,7 +325,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
         isTopLevelUnmount: false,
         next: null,
       };
-      insertUpdateIntoQueue(root.blockers, blockUpdate, currentTime);
+      insertUpdateIntoQueue(root.blockers, block, currentTime);
     }
 
     scheduleUpdate(current, expirationTime);
@@ -349,25 +349,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
   WorkNode.prototype.then = function(callback) {
     const root = this._reactRootContainer;
     const expirationTime = this._expirationTime;
-
-    // Add callback to queue of callbacks on the root. It will be called once
-    // the root completes at the corresponding expiration time.
-    const update = {
-      priorityLevel: null,
-      expirationTime,
-      partialState: null,
-      callback,
-      isReplace: false,
-      isForced: false,
-      isTopLevelUnmount: false,
-      next: null,
-    };
-    const currentTime = recalculateCurrentTime();
-    if (root.completionCallbacks === null) {
-      root.completionCallbacks = createUpdateQueue();
-    }
-    insertUpdateIntoQueue(root.completionCallbacks, update, currentTime);
-    scheduleUpdate(root.current, expirationTime);
+    scheduleCompletionCallback(root, callback, expirationTime);
   };
 
   return {

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -98,6 +98,8 @@ var {
 var {
   getUpdateExpirationTime,
   processUpdateQueue,
+  createUpdateQueue,
+  insertUpdateIntoQueue,
 } = require('ReactFiberUpdateQueue');
 
 var {resetContext} = require('ReactFiberContext');
@@ -369,13 +371,28 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
 
     if (completedAt !== Done) {
       // The root completed but was blocked from committing.
-
       if (expirationTime < completedAt) {
-        // We have work that expires earlier than the completed root. Regardless
-        // of whether the root is blocked, we should work on it.
+        // We have work that expires earlier than the completed root.
         return expirationTime;
       }
 
+      // If the expiration time of the pending work is equal to the time at
+      // which we completed the work-in-progress, it's possible additional
+      // work was scheduled that happens to fall within the same expiration
+      // bucket. We need to check the work-in-progress fiber.
+      if (expirationTime === completedAt) {
+        const workInProgress = root.current.alternate;
+        if (
+          workInProgress !== null &&
+          (workInProgress.expirationTime !== Done &&
+            workInProgress.expirationTime <= expirationTime)
+        ) {
+          // We have more work. Restart the completed tree.
+          root.completedAt = Done;
+          return expirationTime;
+        }
+      }
+
       // There have been no higher priority updates since we completed the root.
       // If it's still blocked, return Done, as if it has no more work. If it's
       // no longer blocked, return the time at which it completed so that we
@@ -797,7 +814,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
         if (isRootBlocked(root, nextRenderExpirationTime)) {
           // The root is blocked from committing. Mark it as complete so we
           // know we can commit it later without starting new work.
-          root.completedAt = workInProgress.expirationTime = nextRenderExpirationTime;
+          root.completedAt = nextRenderExpirationTime;
         } else {
           // The root is not blocked, so we can commit it now.
           pendingCommit = workInProgress;
@@ -1612,6 +1629,37 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
     }
   }
 
+  function scheduleCompletionCallback(
+    root: FiberRoot,
+    callback: () => mixed,
+    expirationTime: ExpirationTime,
+  ) {
+    // Add callback to queue of callbacks on the root. It will be called once
+    // the root completes at the corresponding expiration time.
+    const update = {
+      priorityLevel: null,
+      expirationTime,
+      partialState: null,
+      callback,
+      isReplace: false,
+      isForced: false,
+      isTopLevelUnmount: false,
+      next: null,
+    };
+    const currentTime = recalculateCurrentTime();
+    if (root.completionCallbacks === null) {
+      root.completionCallbacks = createUpdateQueue();
+    }
+    insertUpdateIntoQueue(root.completionCallbacks, update, currentTime);
+    if (expirationTime === root.completedAt) {
+      // The tree already completed at this expiration time. Resolve the
+      // callback synchronously.
+      performWork(TaskPriority, null);
+    } else {
+      scheduleUpdate(root.current, expirationTime);
+    }
+  }
+
   function getPriorityContext(
     fiber: Fiber,
     forceAsync: boolean,
@@ -1751,6 +1799,7 @@ module.exports = function<T, P, I, TI, PI, C, CX, PL>(
 
   return {
     scheduleUpdate: scheduleUpdate,
+    scheduleCompletionCallback: scheduleCompletionCallback,
     getPriorityContext: getPriorityContext,
     recalculateCurrentTime: recalculateCurrentTime,
     getExpirationTimeForPriority: getExpirationTimeForPriority,

src/renderers/shared/fiber/__tests__/ReactIncrementalRoot-test.js

@@ -0,0 +1,101 @@
+/**
+ * Copyright 2013-present, Facebook, Inc.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree. An additional grant
+ * of patent rights can be found in the PATENTS file in the same directory.
+ *
+ * @emails react-core
+ */
+
+'use strict';
+
+var React;
+var ReactNoop;
+
+describe('ReactIncrementalRoot', () => {
+  beforeEach(() => {
+    jest.resetModules();
+    React = require('react');
+    ReactNoop = require('react-noop-renderer');
+  });
+
+  function span(prop) {
+    return {type: 'span', children: [], prop};
+  }
+
+  it('prerenders roots', () => {
+    const root = ReactNoop.create();
+    const work = root.prerender(<span prop="A" />);
+    expect(root.getChildren()).toEqual([]);
+    work.commit();
+    expect(root.getChildren()).toEqual([span('A')]);
+  });
+
+  it('resolves `then` callback synchronously if tree is already completed', () => {
+    const root = ReactNoop.create();
+    const work = root.prerender(<span prop="A" />);
+    ReactNoop.flush();
+    let wasCalled = false;
+    work.then(() => {
+      wasCalled = true;
+    });
+    expect(wasCalled).toBe(true);
+  });
+
+  it('does not restart a completed tree if there were no additional updates', () => {
+    let ops = [];
+    function Foo(props) {
+      ops.push('Foo');
+      return <span prop={props.children} />;
+    }
+    const root = ReactNoop.create();
+    const work = root.prerender(<Foo>Hi</Foo>);
+
+    ReactNoop.flush();
+    expect(ops).toEqual(['Foo']);
+    expect(root.getChildren([]));
+
+    work.then(() => {
+      ops.push('Root completed');
+      work.commit();
+      ops.push('Root committed');
+    });
+
+    expect(ops).toEqual([
+      'Foo',
+      'Root completed',
+      // Should not re-render Foo
+      'Root committed',
+    ]);
+    expect(root.getChildren([span('Hi')]));
+  });
+
+  it('works on a blocked tree if the expiration time is less than or equal to the blocked update', () => {
+    let ops = [];
+    function Foo(props) {
+      ops.push('Foo: ' + props.children);
+      return <span prop={props.children} />;
+    }
+    const root = ReactNoop.create();
+    root.prerender(<Foo>A</Foo>);
+    ReactNoop.flush();
+
+    expect(ops).toEqual(['Foo: A']);
+    expect(root.getChildren([]));
+
+    // workA and workB have the same expiration time
+    root.prerender(<Foo>B</Foo>);
+    ReactNoop.flush();
+
+    // Should have re-rendered the root, even though it's blocked
+    // from committing.
+    expect(ops).toEqual(['Foo: A', 'Foo: B']);
+    expect(root.getChildren([]));
+  });
+
+  it(
+    'does not work on on a blocked tree if the expiration time is greater than the blocked update',
+  );
+});