Bailout of consumer updates using bitmask

The context type defines an optional function that compares two context
values, returning a bitfield. A consumer may specify the bits it needs
for rendering. If a provider's context changes, and the consumer's bits
do not intersect with the changed bits, we can skip the consumer.

This is similar to how selectors are used in Redux but fast enough to do
while scanning the tree. The only user code involved is the function
that computes the changed bits. But that's only called once per provider
update, not for every consumer.

Author: acdlite

packages/react-dom/src/server/ReactPartialRenderer.js

@@ -30,7 +30,7 @@ import {
   REACT_RETURN_TYPE,
   REACT_PORTAL_TYPE,
   REACT_PROVIDER_TYPE,
-  REACT_CONSUMER_TYPE,
+  REACT_CONTEXT_TYPE,
 } from 'shared/ReactSymbols';
 
 import {
@@ -755,15 +755,15 @@ class ReactDOMServerRenderer {
             this.stack.push(frame);
             return '';
           }
-          case REACT_CONSUMER_TYPE: {
+          case REACT_CONTEXT_TYPE: {
             const consumer: ReactConsumer<any> = (nextChild: any);
             const nextProps = consumer.props;
 
-            const provider = consumer.type.context.currentProvider;
+            const provider = consumer.type.currentProvider;
             let nextValue;
             if (provider === null) {
               // Detached consumer
-              nextValue = consumer.type.context.defaultValue;
+              nextValue = consumer.type.defaultValue;
             } else {
               nextValue = provider.props.value;
             }

packages/react-reconciler/src/ReactFiber.js

@@ -38,7 +38,7 @@ import {
   REACT_RETURN_TYPE,
   REACT_CALL_TYPE,
   REACT_PROVIDER_TYPE,
-  REACT_CONSUMER_TYPE,
+  REACT_CONTEXT_TYPE,
 } from 'shared/ReactSymbols';
 
 let hasBadMapPolyfill;
@@ -339,7 +339,8 @@ export function createFiberFromElement(
             case REACT_PROVIDER_TYPE:
               fiberTag = ProviderComponent;
               break;
-            case REACT_CONSUMER_TYPE:
+            case REACT_CONTEXT_TYPE:
+              // This is a consumer
               fiberTag = ConsumerComponent;
               break;
             default:

packages/react-reconciler/src/ReactFiberBeginWork.js

@@ -8,11 +8,7 @@
  */
 
 import type {HostConfig} from 'react-reconciler';
-import type {
-  ReactProviderType,
-  ReactConsumerType,
-  ReactContext,
-} from 'shared/ReactTypes';
+import type {ReactProviderType, ReactContext} from 'shared/ReactTypes';
 import type {Fiber} from 'react-reconciler/src/ReactFiber';
 import type {HostContext} from './ReactFiberHostContext';
 import type {HydrationContext} from './ReactFiberHydrationContext';
@@ -68,6 +64,7 @@ import {
 } from './ReactFiberContext';
 import {pushProvider} from './ReactFiberNewContext';
 import {NoWork, Never} from './ReactFiberExpirationTime';
+import MAX_SIGNED_32_BIT_INT from './maxSigned32BitInt';
 
 let warnedAboutStatelessRefs;
 
@@ -616,6 +613,7 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
   function propagateContextChange<V>(
     workInProgress: Fiber,
     context: ReactContext<V>,
+    changedBits: number,
     renderExpirationTime: ExpirationTime,
   ): void {
     if (enableNewContextAPI) {
@@ -626,7 +624,8 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
         switch (fiber.tag) {
           case ConsumerComponent:
             // Check if the context matches.
-            if (fiber.type.context === context) {
+            const bits = fiber.stateNode;
+            if (fiber.type === context && (bits & changedBits) !== 0) {
               // Update the expiration time of all the ancestors, including
               // the alternates.
               let node = fiber;
@@ -668,7 +667,7 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
             break;
           case ProviderComponent:
             // Don't scan deeper if this is a matching provider
-            nextFiber = fiber.type === context ? null : fiber.child;
+            nextFiber = fiber.type === workInProgress.type ? null : fiber.child;
             break;
           default:
             // Traverse down.
@@ -724,12 +723,33 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
       workInProgress.memoizedProps = newProps;
 
       const newValue = newProps.value;
-      const oldValue = oldProps !== null ? oldProps.value : null;
 
-      // Use Object.is to compare the new context value to the old value.
-      if (!is(newValue, oldValue)) {
-        propagateContextChange(workInProgress, context, renderExpirationTime);
+      let changedBits: number;
+      if (oldProps === null) {
+        // Initial render
+        changedBits = MAX_SIGNED_32_BIT_INT;
+      } else {
+        const oldValue = oldProps.value;
+        // Use Object.is to compare the new context value to the old value.
+        if (!is(newValue, oldValue)) {
+          changedBits =
+            context.calculateChangedBits !== null
+              ? context.calculateChangedBits(oldValue, newValue)
+              : MAX_SIGNED_32_BIT_INT;
+          if (changedBits !== 0) {
+            propagateContextChange(
+              workInProgress,
+              context,
+              changedBits,
+              renderExpirationTime,
+            );
+          }
+        } else {
+          // No change.
+          changedBits = 0;
+        }
       }
+      workInProgress.stateNode = changedBits;
 
       if (oldProps !== null && oldProps.children === newProps.children) {
         return bailoutOnAlreadyFinishedWork(current, workInProgress);
@@ -748,8 +768,7 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
     renderExpirationTime,
   ) {
     if (enableNewContextAPI) {
-      const consumerType: ReactConsumerType<any> = workInProgress.type;
-      const context: ReactContext<any> = consumerType.context;
+      const context: ReactContext<any> = workInProgress.type;
 
       const newProps = workInProgress.pendingProps;
       const oldProps = workInProgress.memoizedProps;
@@ -758,12 +777,12 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
       const providerFiber: Fiber | null = context.currentProvider;
 
       let newValue;
-      let valueDidChange;
+      let changedBits;
       if (providerFiber === null) {
         // This is a detached consumer (has no provider). Use the default
         // context value.
         newValue = context.defaultValue;
-        valueDidChange = false;
+        changedBits = 0;
       } else {
         const provider = providerFiber.pendingProps;
         invariant(
@@ -772,35 +791,31 @@ export default function<T, P, I, TI, HI, PI, C, CC, CX, PL>(
             'a bug in React. Please file an issue.',
         );
         newValue = provider.value;
-
-        // Context change propagation stops at matching consumers, for time-
-        // slicing. Continue the propagation here.
-        if (oldProps === null) {
-          valueDidChange = true;
-          propagateContextChange(workInProgress, context, renderExpirationTime);
-        } else {
-          const oldValue = oldProps !== null ? oldProps.__memoizedValue : null;
-          // Use Object.is to compare the new context value to the old value.
-          if (!is(newValue, oldValue)) {
-            valueDidChange = true;
-            propagateContextChange(
-              workInProgress,
-              context,
-              renderExpirationTime,
-            );
-          }
+        changedBits = providerFiber.stateNode;
+        if (changedBits !== 0) {
+          // Context change propagation stops at matching consumers, for time-
+          // slicing. Continue the propagation here.
+          propagateContextChange(
+            workInProgress,
+            context,
+            changedBits,
+            renderExpirationTime,
+          );
         }
       }
 
-      // The old context value is stored on the consumer object. We can't use the
-      // provider's memoizedProps because those have already been updated by the
-      // time we get here, in the provider's begin phase.
-      newProps.__memoizedValue = newValue;
+      // Store the bits on the fiber's stateNode for quick access.
+      let bits = newProps.bits;
+      if (bits === undefined || bits === null) {
+        // Subscribe to all changes by default
+        bits = MAX_SIGNED_32_BIT_INT;
+      }
+      workInProgress.stateNode = bits;
 
       if (hasLegacyContextChanged()) {
         // Normally we can bail out on props equality but if context has changed
         // we don't do the bailout and we have to reuse existing props instead.
-      } else if (newProps === oldProps && !valueDidChange) {
+      } else if (newProps === oldProps && changedBits === 0) {
         return bailoutOnAlreadyFinishedWork(current, workInProgress);
       }
       const newChildren = newProps.render(newValue);

packages/react-reconciler/src/ReactFiberExpirationTime.js

@@ -7,12 +7,14 @@
  * @flow
  */
 
+import MAX_SIGNED_32_BIT_INT from './maxSigned32BitInt';
+
 // TODO: Use an opaque type once ESLint et al support the syntax
 export type ExpirationTime = number;
 
 export const NoWork = 0;
 export const Sync = 1;
-export const Never = 2147483647; // Max int32: Math.pow(2, 31) - 1
+export const Never = MAX_SIGNED_32_BIT_INT;
 
 const UNIT_SIZE = 10;
 const MAGIC_NUMBER_OFFSET = 2;

packages/react-reconciler/src/__tests__/ReactNewContext-test.internal.js

@@ -425,6 +425,80 @@ describe('ReactNewContext', () => {
     ]);
   });
 
+  it('can skip consumers with bitmask', () => {
+    const Context = React.unstable_createContext({foo: 0, bar: 0}, (a, b) => {
+      let result = 0;
+      if (a.foo !== b.foo) {
+        result |= 0b01;
+      }
+      if (a.bar !== b.bar) {
+        result |= 0b10;
+      }
+      return result;
+    });
+
+    function Provider(props) {
+      return Context.provide({foo: props.foo, bar: props.bar}, props.children);
+    }
+
+    function Foo() {
+      return Context.consume(value => {
+        ReactNoop.yield('Foo');
+        return <span prop={'Foo: ' + value.foo} />;
+      }, 0b01);
+    }
+
+    function Bar() {
+      return Context.consume(value => {
+        ReactNoop.yield('Bar');
+        return <span prop={'Bar: ' + value.bar} />;
+      }, 0b10);
+    }
+
+    class Indirection extends React.Component {
+      shouldComponentUpdate() {
+        return false;
+      }
+      render() {
+        return this.props.children;
+      }
+    }
+
+    function App(props) {
+      return (
+        <Provider foo={props.foo} bar={props.bar}>
+          <Indirection>
+            <Indirection>
+              <Foo />
+            </Indirection>
+            <Indirection>
+              <Bar />
+            </Indirection>
+          </Indirection>
+        </Provider>
+      );
+    }
+
+    ReactNoop.render(<App foo={1} bar={1} />);
+    expect(ReactNoop.flush()).toEqual(['Foo', 'Bar']);
+    expect(ReactNoop.getChildren()).toEqual([span('Foo: 1'), span('Bar: 1')]);
+
+    // Update only foo
+    ReactNoop.render(<App foo={2} bar={1} />);
+    expect(ReactNoop.flush()).toEqual(['Foo']);
+    expect(ReactNoop.getChildren()).toEqual([span('Foo: 2'), span('Bar: 1')]);
+
+    // Update only bar
+    ReactNoop.render(<App foo={2} bar={2} />);
+    expect(ReactNoop.flush()).toEqual(['Bar']);
+    expect(ReactNoop.getChildren()).toEqual([span('Foo: 2'), span('Bar: 2')]);
+
+    // Update both
+    ReactNoop.render(<App foo={3} bar={3} />);
+    expect(ReactNoop.flush()).toEqual(['Foo', 'Bar']);
+    expect(ReactNoop.getChildren()).toEqual([span('Foo: 3'), span('Bar: 3')]);
+  });
+
   describe('fuzz test', () => {
     const contextKeys = ['A', 'B', 'C', 'D', 'E', 'F', 'G'];
     const contexts = new Map(
@@ -521,6 +595,7 @@ describe('ReactNewContext', () => {
                 />
               </Fragment>
             ),
+            null,
             i,
           );
         });

packages/react-reconciler/src/maxSigned32BitInt.js

@@ -0,0 +1,13 @@
+/**
+ * Copyright (c) 2013-present, Facebook, Inc.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+// The maximum safe integer for bitwise operations.
+// Math.pow(2, 31) - 1
+// 0b1111111111111111111111111111111
+export default 2147483647;