setState for Fiber

Updates are scheduled by setting a work priority on the fiber and bubbling it to
the root. Because the instance does not know which tree is current at any given
time, the update is scheduled on both fiber alternates.

Need to add more unit tests to cover edge cases.

Author: acdlite

src/renderers/shared/fiber/ReactFiber.js

@@ -78,6 +78,10 @@ export type Fiber = Instance & {
   pendingProps: any, // This type will be more specific once we overload the tag.
   // TODO: I think that there is a way to merge pendingProps and memoizedProps.
   memoizedProps: any, // The props used to create the output.
+  // Local state for class components. May need better naming to disambiguate
+  // from stateNode.
+  pendingState: any,
+  memoizedState: any, // The state used to create the output.
   // Output is the return value of this fiber, or a linked list of return values
   // if this returns multiple values. Such as a fragment.
   output: any, // This type will be more specific once we overload the tag.
@@ -142,6 +146,8 @@ var createFiber = function(tag : TypeOfWork, key : null | string) : Fiber {
 
     pendingProps: null,
     memoizedProps: null,
+    pendingState: null,
+    memoizedState: null,
     output: null,
 
     nextEffect: null,
@@ -176,6 +182,7 @@ exports.cloneFiber = function(fiber : Fiber, priorityLevel : PriorityLevel) : Fi
     alt.sibling = fiber.sibling;
     alt.ref = alt.ref;
     alt.pendingProps = fiber.pendingProps;
+    alt.pendingState = fiber.pendingState;
     alt.pendingWorkPriority = priorityLevel;
 
     // Whenever we clone, we do so to get a new work in progress.
@@ -196,6 +203,7 @@ exports.cloneFiber = function(fiber : Fiber, priorityLevel : PriorityLevel) : Fi
   alt.ref = alt.ref;
   // pendingProps is here for symmetry but is unnecessary in practice for now.
   alt.pendingProps = fiber.pendingProps;
+  alt.pendingState = fiber.pendingState;
   alt.pendingWorkPriority = priorityLevel;
 
   alt.alternate = fiber;
@@ -232,7 +240,7 @@ function createFiberFromElementType(type : mixed, key : null | string) {
     throw new Error('Unknown component type: ' + typeof type);
   }
   return fiber;
-};
+}
 
 exports.createFiberFromElementType = createFiberFromElementType;
 

src/renderers/shared/fiber/ReactFiberBeginWork.js

@@ -14,9 +14,13 @@
 
 import type { ReactCoroutine } from 'ReactCoroutine';
 import type { Fiber } from 'ReactFiber';
+import type { FiberRoot } from 'ReactFiberRoot';
 import type { HostConfig } from 'ReactFiberReconciler';
+import type { Scheduler } from 'ReactFiberScheduler';
+import type { PriorityLevel } from 'ReactPriorityLevel';
 
 var { reconcileChildFibers } = require('ReactChildFiber');
+var { LowPriority } = require('ReactPriorityLevel');
 var ReactTypeOfWork = require('ReactTypeOfWork');
 var {
   IndeterminateComponent,
@@ -34,8 +38,7 @@ var {
 } = require('ReactPriorityLevel');
 var { findNextUnitOfWorkAtPriority } = require('ReactFiberPendingWork');
 
-module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
-
+module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>, getScheduler: () => Scheduler) {
   function reconcileChildren(current, workInProgress, nextChildren) {
     const priority = workInProgress.pendingWorkPriority;
     workInProgress.child = reconcileChildFibers(
@@ -54,31 +57,88 @@ module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
     workInProgress.pendingWorkPriority = NoWork;
   }
 
+  function updateFiber(fiber: Fiber, state: any, priorityLevel : PriorityLevel): void {
+    const { scheduleLowPriWork } = getScheduler();
+    fiber.pendingState = state;
+
+    while (true) {
+      if (fiber.pendingWorkPriority === NoWork ||
+          fiber.pendingWorkPriority >= priorityLevel) {
+        fiber.pendingWorkPriority = priorityLevel;
+      }
+      // Duck type root
+      if (fiber.stateNode && fiber.stateNode.containerInfo) {
+        const root : FiberRoot = (fiber.stateNode : any);
+        scheduleLowPriWork(root, priorityLevel);
+        return;
+      }
+      if (!fiber.return) {
+        throw new Error('No root!');
+      }
+      fiber = fiber.return;
+    }
+  }
+
+  // Class component state updater
+  const updater = {
+    enqueueSetState(instance, partialState) {
+      const fiber = instance._fiber;
+
+      const prevState = fiber.pendingState || fiber.memoizedState;
+      const state = Object.assign({}, prevState, partialState);
+
+      // Must schedule an update on both alternates, because we don't know tree
+      // is current.
+      updateFiber(fiber, state, LowPriority);
+      if (fiber.alternate) {
+        updateFiber(fiber.alternate, state, LowPriority);
+      }
+    },
+  };
+
   function updateClassComponent(current : ?Fiber, workInProgress : Fiber) {
+    // A class component update is the result of either new props or new state.
+    // Account for the possibly of missing pending props or state by falling
+    // back to the most recent props or state.
     var props = workInProgress.pendingProps;
+    var state = workInProgress.pendingState;
+    if (!props && current) {
+      props = current.memoizedProps;
+    }
+    if (!state && current) {
+      state = current.memoizedState;
+    }
+
     var instance = workInProgress.stateNode;
     if (!instance) {
       var ctor = workInProgress.type;
       workInProgress.stateNode = instance = new ctor(props);
+      state = workInProgress.pendingState = instance.state || null;
+      // The instance needs access to the fiber so that it can schedule updates
+      instance._fiber = workInProgress;
+      instance.updater = updater;
     } else if (typeof instance.shouldComponentUpdate === 'function') {
       if (current && current.memoizedProps) {
-        // Revert to the last flushed props, incase we aborted an update.
+        // Revert to the last flushed props and state, incase we aborted an update.
         instance.props = current.memoizedProps;
-        if (!instance.shouldComponentUpdate(props)) {
+        instance.state = current.memoizedState;
+        if (!instance.shouldComponentUpdate(props, state)) {
           return bailoutOnCurrent(current, workInProgress);
         }
       }
       if (!workInProgress.hasWorkInProgress && workInProgress.memoizedProps) {
-        // Reset the props, in case this is a ping-pong case rather than a
-        // completed update case. For the completed update case, the instance
-        // props will already be the memoizedProps.
+        // Reset the props and state, in case this is a ping-pong case rather
+        // than a completed update case. For the completed update case, the
+        // instance props and state will already be the memoized props and state.
         instance.props = workInProgress.memoizedProps;
-        if (!instance.shouldComponentUpdate(props)) {
+        instance.state = workInProgress.memoizedState;
+        if (!instance.shouldComponentUpdate(props, state)) {
           return bailoutOnAlreadyFinishedWork(workInProgress);
         }
       }
     }
     instance.props = props;
+    instance.state = state;
     var nextChildren = instance.render();
     reconcileChildren(current, workInProgress, nextChildren);
     workInProgress.pendingWorkPriority = NoWork;
@@ -188,9 +248,11 @@ module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
     // the same props as the new one. In that case, we can just copy the output
     // and children from that node.
     workInProgress.memoizedProps = workInProgress.pendingProps;
+    workInProgress.memoizedState = workInProgress.pendingState;
     workInProgress.output = current.output;
     const priorityLevel = workInProgress.pendingWorkPriority;
     workInProgress.pendingProps = null;
+    workInProgress.pendingState = current.pendingState = null;
     workInProgress.pendingWorkPriority = NoWork;
     workInProgress.stateNode = current.stateNode;
     if (current.child) {
@@ -219,6 +281,10 @@ module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
     // looking for. In that case, we should be able to just bail out.
     const priorityLevel = workInProgress.pendingWorkPriority;
     workInProgress.pendingProps = null;
+    workInProgress.pendingState = null;
+    if (workInProgress.alternate) {
+      workInProgress.alternate.pendingState = null;
+    }
     workInProgress.pendingWorkPriority = NoWork;
 
     workInProgress.firstEffect = null;
@@ -252,12 +318,17 @@ module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
     // Ideally nothing should rely on this, but relying on it here
     // means that we don't need an additional field on the work in
     // progress.
-    if (current && workInProgress.pendingProps === current.memoizedProps) {
+    if (current &&
+        workInProgress.pendingProps === current.memoizedProps &&
+        workInProgress.pendingState === current.memoizedState
+    ) {
       return bailoutOnCurrent(current, workInProgress);
     }
 
     if (!workInProgress.hasWorkInProgress &&
-        workInProgress.pendingProps === workInProgress.memoizedProps) {
+        workInProgress.pendingProps === workInProgress.memoizedProps &&
+        workInProgress.pendingState === workInProgress.memoizedState
+    ) {
       return bailoutOnAlreadyFinishedWork(workInProgress);
     }
 

src/renderers/shared/fiber/ReactFiberCompleteWork.js

@@ -67,6 +67,7 @@ module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
     // the linked list of fibers that has the individual output values.
     returnFiber.output = (child && !child.sibling) ? child.output : child;
     returnFiber.memoizedProps = returnFiber.pendingProps;
+    returnFiber.memoizedState = returnFiber.pendingState;
   }
 
   function recursivelyFillYields(yields, output : ?Fiber | ?ReifiedYield) {

src/renderers/shared/fiber/ReactFiberPendingWork.js

@@ -40,13 +40,14 @@ exports.findNextUnitOfWorkAtPriority = function(currentRoot : Fiber, priorityLev
     if (current.pendingWorkPriority !== NoWork &&
         current.pendingWorkPriority <= priorityLevel) {
       // This node has work to do that fits our priority level criteria.
-      if (current.pendingProps !== null) {
+      if (current.pendingProps !== null || current.pendingState !== null) {
         // We found some work to do. We need to return the "work in progress"
         // of this node which will be the alternate.
         const workInProgress = current.alternate;
         if (!workInProgress) {
           throw new Error('Should have wip now');
         }
+        // If current has pendingState, it's already set on the workInProgress.
         workInProgress.pendingProps = current.pendingProps;
         return workInProgress;
       }

src/renderers/shared/fiber/ReactFiberScheduler.js

@@ -33,9 +33,19 @@ var {
 
 var timeHeuristicForUnitOfWork = 1;
 
+export type Scheduler = {
+  scheduleLowPriWork: (root : FiberRoot, priority : PriorityLevel) => void
+};
+
 module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
+  // Use a closure to circumvent the circular dependency between the scheduler
+  // and ReactFiberBeginWork. Don't know if there's a better way to do this.
+  let scheduler;
+  function getScheduler(): Scheduler {
+    return scheduler;
+  }
 
-  const { beginWork } = ReactFiberBeginWork(config);
+  const { beginWork } = ReactFiberBeginWork(config, getScheduler);
   const { completeWork } = ReactFiberCompleteWork(config);
   const { commitWork } = ReactFiberCommitWork(config);
 
@@ -119,6 +129,10 @@ module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
       // The work is now done. We don't need this anymore. This flags
       // to the system not to redo any work here.
       workInProgress.pendingProps = null;
+      workInProgress.pendingState = null;
+      if (current) {
+        current.pendingState = null;
+      }
       if (workInProgress.pendingWorkPriority === NoWork) {
         workInProgress.hasWorkInProgress = false;
       }
@@ -179,7 +193,7 @@ module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
 
   function performUnitOfWork(workInProgress : Fiber) : ?Fiber {
     // Ignore work if there is nothing to do.
-    if (workInProgress.pendingProps === null) {
+    if (workInProgress.pendingProps === null && workInProgress.pendingState === null) {
       return completeUnitOfWork(workInProgress);
     }
     // The current, flushed, state of this fiber is the alternate.
@@ -250,7 +264,8 @@ module.exports = function<T, P, I, C>(config : HostConfig<T, P, I, C>) {
   }
   */
 
-  return {
+  scheduler = {
     scheduleLowPriWork: scheduleLowPriWork,
   };
+  return scheduler;
 };

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

@@ -539,4 +539,36 @@ describe('ReactIncremental', function() {
     expect(ops).toEqual(['Content', 'Bar', 'Middle']);
 
   });
+
+  it('can update in the middle of a tree using setState', () => {
+    let instance;
+    let states = [];
+
+    class Bar extends React.Component {
+      constructor() {
+        super();
+        this.state = { string: 'a' };
+        instance = this;
+      }
+      render() {
+        states.push(this.state.string);
+        return <div>{this.props.children}</div>;
+      }
+    }
+
+    function Foo() {
+      return (
+        <div>
+          <Bar />
+        </div>
+      );
+    }
+
+    ReactNoop.render(<Foo />);
+    ReactNoop.flush();
+    expect(states).toEqual(['a']);
+    instance.setState({ string: 'b' });
+    ReactNoop.flush();
+    expect(states).toEqual(['a', 'b']);
+  });
 });