API for prerendering a top-level update and deferring the commit

packages/react-dom/src/__tests__/ReactDOMRoot-test.js

@@ -13,6 +13,8 @@ var React = require('react');
 var ReactDOM = require('react-dom');
 var ReactDOMServer = require('react-dom/server');
 
+const AsyncComponent = React.unstable_AsyncComponent;
+
 describe('ReactDOMRoot', () => {
   let container;
 
@@ -68,4 +70,114 @@ describe('ReactDOMRoot', () => {
     root.render(<div><span>d</span><span>c</span></div>);
     expect(container.textContent).toEqual('abdc');
   });
+
+  it('can defer commit using prerender', () => {
+    const root = ReactDOM.createRoot(container);
+    const work = root.prerender(<div>Hi</div>);
+    // Hasn't updated yet
+    expect(container.textContent).toEqual('');
+    // Flush work
+    work.commit();
+    expect(container.textContent).toEqual('Hi');
+  });
+
+  it("does not restart a blocked root that wasn't updated", () => {
+    let ops = [];
+    function Foo(props) {
+      ops.push('Foo');
+      return props.children;
+    }
+    const root = ReactDOM.createRoot(container);
+    const work = root.prerender(<Foo>Hi</Foo>);
+    expect(ops).toEqual(['Foo']);
+    // Hasn't updated yet
+    expect(container.textContent).toEqual('');
+
+    ops = [];
+
+    // Flush work. Shouldn't re-render Foo.
+    work.commit();
+    expect(ops).toEqual([]);
+    expect(container.textContent).toEqual('Hi');
+  });
+
+  it('can wait for prerender to finish', () => {
+    const Async = React.unstable_AsyncComponent;
+    const root = ReactDOM.createRoot(container);
+    const work = root.prerender(<Async>Foo</Async>);
+
+    // Hasn't updated yet
+    expect(container.textContent).toEqual('');
+
+    let ops = [];
+    work.then(() => {
+      // Still hasn't updated
+      ops.push(container.textContent);
+      // Should synchronously commit
+      work.commit();
+      ops.push(container.textContent);
+    });
+    // Flush async work
+    jest.runAllTimers();
+    expect(ops).toEqual(['', 'Foo']);
+  });
+
+  it('resolves `then` callback synchronously if update is sync', () => {
+    const root = ReactDOM.createRoot(container);
+    const work = root.prerender(<div>Hi</div>);
+
+    let ops = [];
+    work.then(() => {
+      work.commit();
+      ops.push(container.textContent);
+      expect(container.textContent).toEqual('Hi');
+    });
+    // `then` should have synchronously resolved
+    expect(ops).toEqual(['Hi']);
+  });
+
+  it('resolves `then` callback if tree already completed', () => {
+    const root = ReactDOM.createRoot(container);
+    const work = root.prerender(<div>Hi</div>);
+
+    let ops = [];
+    work.then(() => {
+      work.commit();
+      ops.push(container.textContent);
+      expect(container.textContent).toEqual('Hi');
+    });
+
+    work.then(() => {
+      ops.push('Second callback');
+    });
+
+    // `then` should have synchronously resolved
+    expect(ops).toEqual(['Hi', 'Second callback']);
+  });
+
+  it('commits an earlier time without unblocking a later time', () => {
+    const root = ReactDOM.createRoot(container);
+    // Sync update
+    const work1 = root.prerender(<div>a</div>);
+    // Async update
+    const work2 = root.prerender(<AsyncComponent>b</AsyncComponent>);
+    // Flush only the sync update
+    work1.commit();
+    jest.runAllTimers();
+    expect(container.textContent).toBe('a');
+    // Now flush the async update
+    work2.commit();
+    expect(container.textContent).toBe('b');
+  });
+
+  it('render returns a work object, too', () => {
+    const root = ReactDOM.createRoot(container);
+    const work = root.render(<div>Hello</div>);
+    let ops = [];
+    work.then(() => {
+      // Work already committed.
+      ops.push(container.textContent);
+    });
+    expect(container.textContent).toEqual('Hello');
+  });
 });

packages/react-dom/src/client/ReactDOM.js

@@ -10,6 +10,17 @@
 'use strict';
 
 import type {ReactNodeList} from 'shared/ReactTypes';
+// TODO: This type is shared between the reconciler and ReactDOM, but will
+// eventually be lifted out to the renderer.
+import type {
+  FiberRoot,
+  WorkNode as FiberRootWorkNode,
+} from 'react-reconciler/src/ReactFiberRoot';
+// TODO: ExpirationTime (or whatever we end up calling it) should be a public
+// type that renderers can consume.
+import type {
+  ExpirationTime,
+} from 'react-reconciler/src/ReactFiberExpirationTime';
 
 require('../shared/checkReact');
 
@@ -759,9 +770,90 @@ function createPortal(
   return ReactPortal.createPortal(children, container, null, key);
 }
 
+type WorkNode = FiberRootWorkNode & {
+  then(onComplete: () => mixed): void,
+  commit(): void,
+
+  _reactRootContainer: FiberRoot,
+
+  _completionCallbacks: Array<() => mixed> | null,
+  _didComplete: boolean,
+};
+
+function insertWork(root: FiberRoot, work: FiberRootWorkNode) {
+  // Insert into root's list of work nodes.
+  const expirationTime = work._expirationTime;
+  const firstNode = root.firstTopLevelWork;
+  if (firstNode === null) {
+    root.firstTopLevelWork = work;
+    work._next = null;
+  } else {
+    // Insert sorted by expiration time then insertion order
+    let insertAfter = null;
+    let insertBefore = firstNode;
+    while (
+      insertBefore !== null &&
+      insertBefore._expirationTime <= expirationTime
+    ) {
+      insertAfter = insertBefore;
+      insertBefore = insertBefore._next;
+    }
+    work._next = insertBefore;
+    if (insertAfter !== null) {
+      insertAfter._next = work;
+    }
+  }
+}
+
+function Work(root: FiberRoot, defer: boolean, expirationTime: ExpirationTime) {
+  this._reactRootContainer = root;
+
+  this._defer = defer;
+  this._expirationTime = expirationTime;
+  this._next = null;
+
+  this._completionCallbacks = null;
+  this._didComplete = false;
+}
+Work.prototype._onComplete = function() {
+  this._didComplete = true;
+  // Set this to null so it isn't called again.
+  this._onComplete = null;
+  const callbacks = this._completionCallbacks;
+  if (callbacks === null) {
+    return;
+  }
+  for (let i = 0; i < callbacks.length; i++) {
+    // TODO: Error handling
+    const callback = callbacks[i];
+    callback();
+  }
+};
+Work.prototype.then = function(cb: () => mixed) {
+  if (this._didComplete) {
+    cb();
+    return;
+  }
+  let completionCallbacks = this._completionCallbacks;
+  if (completionCallbacks === null) {
+    completionCallbacks = this._completionCallbacks = [];
+  }
+  completionCallbacks.push(cb);
+
+  const root = this._reactRootContainer;
+  const expirationTime = this._expirationTime;
+  DOMRenderer.requestWork(root, expirationTime);
+};
+Work.prototype.commit = function() {
+  this._defer = false;
+  const root = this._reactRootContainer;
+  const expirationTime = this._expirationTime;
+  DOMRenderer.flushRoot(root, expirationTime);
+};
+
 type ReactRootNode = {
-  render(children: ReactNodeList, callback: ?() => mixed): void,
-  unmount(callback: ?() => mixed): void,
+  render(children: ReactNodeList, callback: ?() => mixed): WorkNode,
+  unmount(callback: ?() => mixed): WorkNode,
 
   _reactRootContainer: *,
 };
@@ -777,13 +869,51 @@ function ReactRoot(container: Container, hydrate: boolean) {
 ReactRoot.prototype.render = function(
   children: ReactNodeList,
   callback: ?() => mixed,
-): void {
+): WorkNode {
+  const root = this._reactRootContainer;
+  // TODO: Wrapping in batchedUpdates is needed to prevent work on the root from
+  // starting until after the work object is inserted. Remove it once
+  // root scheduling is lifted into the renderer.
+  return DOMRenderer.batchedUpdates(() => {
+    const expirationTime = DOMRenderer.updateContainer(
+      children,
+      root,
+      null,
+      null,
+    );
+    const work = new Work(root, false, expirationTime);
+    insertWork(root, work);
+    return work;
+  });
+};
+ReactRoot.prototype.prerender = function(children: ReactNodeList): WorkNode {
   const root = this._reactRootContainer;
-  DOMRenderer.updateContainer(children, root, null, callback);
+  // TODO: Wrapping in batchedUpdates is needed to prevent work on the root from
+  // starting until after the work object is inserted. Remove it once
+  // root scheduling is lifted into the renderer.
+  return DOMRenderer.batchedUpdates(() => {
+    const expirationTime = DOMRenderer.updateContainer(
+      children,
+      root,
+      null,
+      null,
+    );
+    const work = new Work(root, true, expirationTime);
+    insertWork(root, work);
+    return work;
+  });
 };
-ReactRoot.prototype.unmount = function(callback) {
+ReactRoot.prototype.unmount = function(): WorkNode {
   const root = this._reactRootContainer;
-  DOMRenderer.updateContainer(null, root, null, callback);
+  // TODO: Wrapping in batchedUpdates is needed to prevent work on the root from
+  // starting until after the work object is inserted. Remove it once
+  // root scheduling is lifted into the renderer.
+  return DOMRenderer.batchedUpdates(() => {
+    const expirationTime = DOMRenderer.updateContainer(null, root, null, null);
+    const work = new Work(root, false, expirationTime);
+    insertWork(root, work);
+    return work;
+  });
 };
 
 var ReactDOM = {

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -47,7 +47,7 @@ var {
   reconcileChildFibersInPlace,
   cloneChildFibers,
 } = require('./ReactChildFiber');
-var {processUpdateQueue} = require('./ReactFiberUpdateQueue');
+var {processFiberUpdateQueue} = require('./ReactFiberUpdateQueue');
 var {
   getMaskedContext,
   getUnmaskedContext,
@@ -323,7 +323,7 @@ module.exports = function<T, P, I, TI, PI, C, CC, CX, PL>(
     const updateQueue = workInProgress.updateQueue;
     if (updateQueue !== null) {
       const prevState = workInProgress.memoizedState;
-      const state = processUpdateQueue(
+      const state = processFiberUpdateQueue(
         current,
         workInProgress,
         updateQueue,
@@ -714,7 +714,7 @@ module.exports = function<T, P, I, TI, PI, C, CC, CX, PL>(
   function memoizeState(workInProgress: Fiber, nextState: any) {
     workInProgress.memoizedState = nextState;
     // Don't reset the updateQueue, in case there are pending updates. Resetting
-    // is handled by processUpdateQueue.
+    // is handled by processFiberUpdateQueue.
   }
 
   function beginWork(

packages/react-reconciler/src/ReactFiberClassComponent.js

@@ -30,7 +30,7 @@ var {
 } = require('./ReactFiberContext');
 var {
   insertUpdateIntoFiber,
-  processUpdateQueue,
+  processFiberUpdateQueue,
 } = require('./ReactFiberUpdateQueue');
 var {hasContextChanged} = require('./ReactFiberContext');
 
@@ -448,7 +448,7 @@ module.exports = function(
       // process them now.
       const updateQueue = workInProgress.updateQueue;
       if (updateQueue !== null) {
-        instance.state = processUpdateQueue(
+        instance.state = processFiberUpdateQueue(
           current,
           workInProgress,
           updateQueue,
@@ -505,7 +505,7 @@ module.exports = function(
   //   // Process the update queue before calling shouldComponentUpdate
   //   const updateQueue = workInProgress.updateQueue;
   //   if (updateQueue !== null) {
-  //     newState = processUpdateQueue(
+  //     newState = processFiberUpdateQueue(
   //       workInProgress,
   //       updateQueue,
   //       instance,
@@ -548,7 +548,7 @@ module.exports = function(
   //     // componentWillMount may have called setState. Process the update queue.
   //     const newUpdateQueue = workInProgress.updateQueue;
   //     if (newUpdateQueue !== null) {
-  //       newState = processUpdateQueue(
+  //       newState = processFiberUpdateQueue(
   //         workInProgress,
   //         newUpdateQueue,
   //         instance,
@@ -614,7 +614,7 @@ module.exports = function(
     // TODO: Previous state can be null.
     let newState;
     if (workInProgress.updateQueue !== null) {
-      newState = processUpdateQueue(
+      newState = processFiberUpdateQueue(
         current,
         workInProgress,
         workInProgress.updateQueue,