[Fizz] Implement New Context

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

@@ -727,4 +727,70 @@ describe('ReactDOMFizzServer', () => {
       </div>,
     );
   });
+
+  // @gate experimental
+  it('should resume the context from where it left off', async () => {
+    const ContextA = React.createContext('A0');
+    const ContextB = React.createContext('B0');
+
+    function PrintA() {
+      return (
+        <ContextA.Consumer>{value => <Text text={value} />}</ContextA.Consumer>
+      );
+    }
+
+    class PrintB extends React.Component {
+      static contextType = ContextB;
+      render() {
+        return <Text text={this.context} />;
+      }
+    }
+
+    function AsyncParent({text, children}) {
+      return (
+        <>
+          <AsyncText text={text} />
+          <b>{children}</b>
+        </>
+      );
+    }
+
+    await act(async () => {
+      const {startWriting} = ReactDOMFizzServer.pipeToNodeWritable(
+        <div>
+          <PrintA />
+          <div>
+            <ContextA.Provider value="A0.1">
+              <Suspense fallback={<Text text="Loading..." />}>
+                <AsyncParent text="Child:">
+                  <PrintA />
+                </AsyncParent>
+                <PrintB />
+              </Suspense>
+            </ContextA.Provider>
+          </div>
+          <PrintA />
+        </div>,
+        writable,
+      );
+      startWriting();
+    });
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        A0<div>Loading...</div>A0
+      </div>,
+    );
+    await act(async () => {
+      resolveText('Child:');
+    });
+    expect(getVisibleChildren(container)).toEqual(
+      <div>
+        A0
+        <div>
+          Child:<b>A0.1</b>B0
+        </div>
+        A0
+      </div>,
+    );
+  });
 });

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

@@ -48,6 +48,10 @@ import hasOwnProperty from 'shared/hasOwnProperty';
 import sanitizeURL from '../shared/sanitizeURL';
 import isArray from 'shared/isArray';
 
+// Used to distinguish these contexts from ones used in other renderers.
+// E.g. this can be used to distinguish legacy renderers from this modern one.
+export const isPrimaryRenderer = true;
+
 // Per response, global state that is not contextual to the rendering subtree.
 export type ResponseState = {
   placeholderPrefix: PrecomputedChunk,

packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js

@@ -23,6 +23,8 @@ import {
 
 import invariant from 'shared/invariant';
 
+export const isPrimaryRenderer = true;
+
 // Every list of children or string is null terminated.
 const END_TAG = 0;
 // Tree node tags.

packages/react-server/src/ReactFizzClassComponent.js

@@ -8,6 +8,7 @@
  */
 
 import {emptyContextObject} from './ReactFizzContext';
+import {readContext} from './ReactFizzNewContext';
 
 import {disableLegacyContext} from 'shared/ReactFeatureFlags';
 import {get as getInstance, set as setInstance} from 'shared/ReactInstanceMap';
@@ -211,9 +212,7 @@ export function constructClassInstance(
   }
 
   if (typeof contextType === 'object' && contextType !== null) {
-    // TODO: Implement Context.
-    // context = readContext((contextType: any));
-    throw new Error('Context is not yet implemented.');
+    context = readContext((contextType: any));
   } else if (!disableLegacyContext) {
     context = maskedLegacyContext;
   }
@@ -617,9 +616,7 @@ export function mountClassInstance(
 
   const contextType = ctor.contextType;
   if (typeof contextType === 'object' && contextType !== null) {
-    // TODO: Implement Context.
-    // instance.context = readContext(contextType);
-    throw new Error('Context is not yet implemented.');
+    instance.context = readContext(contextType);
   } else if (disableLegacyContext) {
     instance.context = emptyContextObject;
   } else {

packages/react-server/src/ReactFizzNewContext.js

@@ -0,0 +1,278 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file in the root directory of this source tree.
+ *
+ * @flow
+ */
+
+import type {ReactContext} from 'shared/ReactTypes';
+
+import {isPrimaryRenderer} from './ReactServerFormatConfig';
+
+import invariant from 'shared/invariant';
+
+let rendererSigil;
+if (__DEV__) {
+  // Use this to detect multiple renderers using the same context
+  rendererSigil = {};
+}
+
+// Used to store the parent path of all context overrides in a shared linked list.
+// Forming a reverse tree.
+type ContextNode<T> = {
+  parent: null | ContextNode<any>,
+  depth: number, // Short hand to compute the depth of the tree at this node.
+  context: ReactContext<T>,
+  parentValue: T,
+  value: T,
+};
+
+// The structure of a context snapshot is an implementation of this file.
+// Currently, it's implemented as tracking the current active node.
+export opaque type ContextSnapshot = null | ContextNode<any>;
+
+export const rootContextSnapshot: ContextSnapshot = null;
+
+// We assume that this runtime owns the "current" field on all ReactContext instances.
+// This global (actually thread local) state represents what state all those "current",
+// fields are currently in.
+let currentActiveSnapshot: ContextSnapshot = null;
+
+function popNode(prev: ContextNode<any>): void {
+  if (isPrimaryRenderer) {
+    prev.context._currentValue = prev.parentValue;
+  } else {
+    prev.context._currentValue2 = prev.parentValue;
+  }
+}
+
+function pushNode(next: ContextNode<any>): void {
+  if (isPrimaryRenderer) {
+    next.context._currentValue = next.value;
+  } else {
+    next.context._currentValue2 = next.value;
+  }
+}
+
+function popToNearestCommonAncestor(
+  prev: ContextNode<any>,
+  next: ContextNode<any>,
+): void {
+  if (prev === next) {
+    // We've found a shared ancestor. We don't need to pop nor reapply this one or anything above.
+  } else {
+    popNode(prev);
+    const parentPrev = prev.parent;
+    const parentNext = next.parent;
+    if (parentPrev === null) {
+      invariant(
+        parentNext === null,
+        'The stacks must reach the root at the same time. This is a bug in React.',
+      );
+    } else {
+      invariant(
+        parentNext !== null,
+        'The stacks must reach the root at the same time. This is a bug in React.',
+      );
+      popToNearestCommonAncestor(parentPrev, parentNext);
+      // On the way back, we push the new ones that weren't common.
+      pushNode(next);
+    }
+  }
+}
+
+function popAllPrevious(prev: ContextNode<any>): void {
+  popNode(prev);
+  const parentPrev = prev.parent;
+  if (parentPrev !== null) {
+    popAllPrevious(parentPrev);
+  }
+}
+
+function pushAllNext(next: ContextNode<any>): void {
+  const parentNext = next.parent;
+  if (parentNext !== null) {
+    pushAllNext(parentNext);
+  }
+  pushNode(next);
+}
+
+function popPreviousToCommonLevel(
+  prev: ContextNode<any>,
+  next: ContextNode<any>,
+): void {
+  popNode(prev);
+  const parentPrev = prev.parent;
+  invariant(
+    parentPrev !== null,
+    'The depth must equal at least at zero before reaching the root. This is a bug in React.',
+  );
+  if (parentPrev.depth === next.depth) {
+    // We found the same level. Now we just need to find a shared ancestor.
+    popToNearestCommonAncestor(parentPrev, next);
+  } else {
+    // We must still be deeper.
+    popPreviousToCommonLevel(parentPrev, next);
+  }
+}
+
+function popNextToCommonLevel(
+  prev: ContextNode<any>,
+  next: ContextNode<any>,
+): void {
+  const parentNext = next.parent;
+  invariant(
+    parentNext !== null,
+    'The depth must equal at least at zero before reaching the root. This is a bug in React.',
+  );
+  if (prev.depth === parentNext.depth) {
+    // We found the same level. Now we just need to find a shared ancestor.
+    popToNearestCommonAncestor(prev, parentNext);
+  } else {
+    // We must still be deeper.
+    popNextToCommonLevel(prev, parentNext);
+  }
+  pushNode(next);
+}
+
+// Perform context switching to the new snapshot.
+// To make it cheap to read many contexts, while not suspending, we make the switch eagerly by
+// updating all the context's current values. That way reads, always just read the current value.
+// At the cost of updating contexts even if they're never read by this subtree.
+export function switchContext(newSnapshot: ContextSnapshot): void {
+  // The basic algorithm we need to do is to pop back any contexts that are no longer on the stack.
+  // We also need to update any new contexts that are now on the stack with the deepest value.
+  // The easiest way to update new contexts is to just reapply them in reverse order from the
+  // perspective of the backpointers. To avoid allocating a lot when switching, we use the stack
+  // for that. Therefore this algorithm is recursive.
+  // 1) First we pop which ever snapshot tree was deepest. Popping old contexts as we go.
+  // 2) Then we find the nearest common ancestor from there. Popping old contexts as we go.
+  // 3) Then we reapply new contexts on the way back up the stack.
+  const prev = currentActiveSnapshot;
+  const next = newSnapshot;
+  if (prev !== next) {
+    if (prev === null) {
+      // $FlowFixMe: This has to be non-null since it's not equal to prev.
+      pushAllNext(next);
+    } else if (next === null) {
+      popAllPrevious(prev);
+    } else if (prev.depth === next.depth) {
+      popToNearestCommonAncestor(prev, next);
+    } else if (prev.depth > next.depth) {
+      popPreviousToCommonLevel(prev, next);
+    } else {
+      popNextToCommonLevel(prev, next);
+    }
+    currentActiveSnapshot = next;
+  }
+}
+
+export function pushProvider<T>(
+  context: ReactContext<T>,
+  nextValue: T,
+): ContextSnapshot {
+  let prevValue;
+  if (isPrimaryRenderer) {
+    prevValue = context._currentValue;
+    context._currentValue = nextValue;
+    if (__DEV__) {
+      if (
+        context._currentRenderer !== undefined &&
+        context._currentRenderer !== null &&
+        context._currentRenderer !== rendererSigil
+      ) {
+        console.error(
+          'Detected multiple renderers concurrently rendering the ' +
+            'same context provider. This is currently unsupported.',
+        );
+      }
+      context._currentRenderer = rendererSigil;
+    }
+  } else {
+    prevValue = context._currentValue2;
+    context._currentValue2 = nextValue;
+    if (__DEV__) {
+      if (
+        context._currentRenderer2 !== undefined &&
+        context._currentRenderer2 !== null &&
+        context._currentRenderer2 !== rendererSigil
+      ) {
+        console.error(
+          'Detected multiple renderers concurrently rendering the ' +
+            'same context provider. This is currently unsupported.',
+        );
+      }
+      context._currentRenderer2 = rendererSigil;
+    }
+  }
+  const prevNode = currentActiveSnapshot;
+  const newNode: ContextNode<T> = {
+    parent: prevNode,
+    depth: prevNode === null ? 0 : prevNode.depth + 1,
+    context: context,
+    parentValue: prevValue,
+    value: nextValue,
+  };
+  currentActiveSnapshot = newNode;
+  return newNode;
+}
+
+export function popProvider<T>(context: ReactContext<T>): ContextSnapshot {
+  const prevSnapshot = currentActiveSnapshot;
+  invariant(
+    prevSnapshot !== null,
+    'Tried to pop a Context at the root of the app. This is a bug in React.',
+  );
+  if (__DEV__) {
+    if (prevSnapshot.context !== context) {
+      console.error(
+        'The parent context is not the expected context. This is probably a bug in React.',
+      );
+    }
+  }
+  if (isPrimaryRenderer) {
+    prevSnapshot.context._currentValue = prevSnapshot.parentValue;
+    if (__DEV__) {
+      if (
+        context._currentRenderer !== undefined &&
+        context._currentRenderer !== null &&
+        context._currentRenderer !== rendererSigil
+      ) {
+        console.error(
+          'Detected multiple renderers concurrently rendering the ' +
+            'same context provider. This is currently unsupported.',
+        );
+      }
+      context._currentRenderer = rendererSigil;
+    }
+  } else {
+    prevSnapshot.context._currentValue2 = prevSnapshot.parentValue;
+    if (__DEV__) {
+      if (
+        context._currentRenderer2 !== undefined &&
+        context._currentRenderer2 !== null &&
+        context._currentRenderer2 !== rendererSigil
+      ) {
+        console.error(
+          'Detected multiple renderers concurrently rendering the ' +
+            'same context provider. This is currently unsupported.',
+        );
+      }
+      context._currentRenderer2 = rendererSigil;
+    }
+  }
+  return (currentActiveSnapshot = prevSnapshot.parent);
+}
+
+export function getActiveContext(): ContextSnapshot {
+  return currentActiveSnapshot;
+}
+
+export function readContext<T>(context: ReactContext<T>): T {
+  const value = isPrimaryRenderer
+    ? context._currentValue
+    : context._currentValue2;
+  return value;
+}