From cbc2d43407d2009ed16dfa6561c55ce641c0f68d Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Mon, 12 Apr 2021 12:27:15 -0400
Subject: [PATCH 1/9] Add NewContext module

This implements a reverse linked list tree containing the previous
contexts.
---
 .../react-server/src/ReactFizzNewContext.js   | 148 ++++++++++++++++++
 packages/react-server/src/ReactFizzServer.js  |  32 ++++
 2 files changed, 180 insertions(+)
 create mode 100644 packages/react-server/src/ReactFizzNewContext.js

diff --git a/packages/react-server/src/ReactFizzNewContext.js b/packages/react-server/src/ReactFizzNewContext.js
new file mode 100644
index 0000000000000..ee5f50a17e482
--- /dev/null
+++ b/packages/react-server/src/ReactFizzNewContext.js
@@ -0,0 +1,148 @@
+/**
+ * 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 invariant from 'shared/invariant';
+
+// TODO: Move to format config.
+const isPrimaryRenderer = true;
+
+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;
+
+// Perform context switching to the new snapshot.
+export function switchContext(newSnapshot: ContextSnapshot): void {
+  // TODO: Switch the context.
+}
+
+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>): void {
+  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;
+    }
+  }
+  currentActiveSnapshot = prevSnapshot.parent;
+}
+
+export function getActiveContext(): ContextSnapshot {
+  return currentActiveSnapshot;
+}
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index a3b1babdf5463..65f42ecd897c9 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -19,6 +19,7 @@ import type {
   ResponseState,
   FormatContext,
 } from './ReactServerFormatConfig';
+import type {ContextSnapshot} from './ReactFizzNewContext';
 
 import {
   scheduleWork,
@@ -56,6 +57,12 @@ import {
   processChildContext,
   emptyContextObject,
 } from './ReactFizzContext';
+import {
+  rootContextSnapshot,
+  switchContext,
+  getActiveContext,
+} from './ReactFizzNewContext';
+
 import {
   getIteratorFn,
   REACT_ELEMENT_TYPE,
@@ -104,6 +111,7 @@ type Task = {
   blockedSegment: Segment, // the segment we'll write to
   abortSet: Set<Task>, // the abortable set that this task belongs to
   legacyContext: LegacyContext, // the current legacy context that this task is executing in
+  context: ContextSnapshot, // the current new context that this task is executing in
   assignID: null | SuspenseBoundaryID, // id to assign to the content
 };
 
@@ -220,6 +228,7 @@ export function createRequest(
     rootSegment,
     abortSet,
     emptyContextObject,
+    rootContextSnapshot,
     null,
   );
   pingedTasks.push(rootTask);
@@ -257,6 +266,7 @@ function createTask(
   blockedSegment: Segment,
   abortSet: Set<Task>,
   legacyContext: LegacyContext,
+  context: ContextSnapshot,
   assignID: null | SuspenseBoundaryID,
 ): Task {
   request.allPendingTasks++;
@@ -272,6 +282,7 @@ function createTask(
     blockedSegment,
     abortSet,
     legacyContext,
+    context,
     assignID,
   };
   abortSet.add(task);
@@ -394,6 +405,7 @@ function renderSuspenseBoundary(
     boundarySegment,
     fallbackAbortSet,
     task.legacyContext,
+    task.context,
     newBoundary.id, // This is the ID we want to give this fallback so we can replace it later.
   );
   // TODO: This should be queued at a separate lower priority queue so that we only work
@@ -918,6 +930,7 @@ function spawnNewSuspendedTask(
     newSegment,
     task.abortSet,
     task.legacyContext,
+    task.context,
     task.assignID,
   );
   // We've delegated the assignment.
@@ -936,6 +949,7 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
   // process.
   const previousFormatContext = task.blockedSegment.formatContext;
   const previousLegacyContext = task.legacyContext;
+  const previousContext = task.context;
   try {
     return renderNodeDestructive(request, task, node);
   } catch (x) {
@@ -945,6 +959,9 @@ function renderNode(request: Request, task: Task, node: ReactNodeList): void {
       // functions in case nothing throws so we don't use "finally" here.
       task.blockedSegment.formatContext = previousFormatContext;
       task.legacyContext = previousLegacyContext;
+      task.context = previousContext;
+      // Restore all active ReactContexts to what they were before.
+      switchContext(previousContext);
     } else {
       // We assume that we don't need the correct context.
       // Let's terminate the rest of the tree and don't render any siblings.
@@ -1104,6 +1121,10 @@ function retryTask(request: Request, task: Task): void {
     // We completed this by other means before we had a chance to retry it.
     return;
   }
+  // We restore the context to what it was when we suspended.
+  // We don't restore it after we leave because it's likely that we'll end up
+  // needing a very similar context soon again.
+  switchContext(task.context);
   try {
     // We call the destructive form that mutates this task. That way if something
     // suspends again, we can reuse the same task instead of spawning a new one.
@@ -1129,6 +1150,7 @@ function performWork(request: Request): void {
   if (request.status === CLOSED) {
     return;
   }
+  const prevContext = getActiveContext();
   const prevDispatcher = ReactCurrentDispatcher.current;
   ReactCurrentDispatcher.current = Dispatcher;
 
@@ -1148,6 +1170,16 @@ function performWork(request: Request): void {
     fatalError(request, error);
   } finally {
     ReactCurrentDispatcher.current = prevDispatcher;
+    if (prevDispatcher === Dispatcher) {
+      // This means that we were in a reentrant work loop. This could happen
+      // in a renderer that supports synchronous work like renderToString,
+      // when it's called from within another renderer.
+      // Normally we don't bother switching the contexts to their root/default
+      // values when leaving because we'll likely need the same or similar
+      // context again. However, when we're inside a synchronous loop like this
+      // we'll to restore the context to what it was before returning.
+      switchContext(prevContext);
+    }
   }
 }
 

From cfbc68efc2246c07f75d9e5dc3962fad77698ac1 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Tue, 13 Apr 2021 12:22:37 -0400
Subject: [PATCH 2/9] Implement recursive algorithm

This algorithm pops the contexts back to a shared ancestor on the way down
the stack and then pushes new contexts in reverse order up the stack.
---
 .../react-server/src/ReactFizzNewContext.js   | 126 +++++++++++++++++-
 1 file changed, 125 insertions(+), 1 deletion(-)

diff --git a/packages/react-server/src/ReactFizzNewContext.js b/packages/react-server/src/ReactFizzNewContext.js
index ee5f50a17e482..fc5c3e6119902 100644
--- a/packages/react-server/src/ReactFizzNewContext.js
+++ b/packages/react-server/src/ReactFizzNewContext.js
@@ -41,9 +41,133 @@ export const rootContextSnapshot: ContextSnapshot = null;
 // 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 {
-  // TODO: Switch the context.
+  // 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>(

From beb08e3c33eb7314cba9ae9396992ec2df220670 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Tue, 13 Apr 2021 12:30:00 -0400
Subject: [PATCH 3/9] Move isPrimaryRenderer to ServerFormatConfig

This is primarily intended to be used to support renderToString with a
separate build than the main one. This allows them to be nested.
---
 packages/react-dom/src/server/ReactDOMServerFormatConfig.js  | 4 ++++
 .../src/server/ReactNativeServerFormatConfig.js              | 2 ++
 packages/react-server/src/ReactFizzNewContext.js             | 5 ++---
 .../react-server/src/forks/ReactServerFormatConfig.custom.js | 2 ++
 4 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/packages/react-dom/src/server/ReactDOMServerFormatConfig.js b/packages/react-dom/src/server/ReactDOMServerFormatConfig.js
index 45cd0fecb5e96..521fe04439c11 100644
--- a/packages/react-dom/src/server/ReactDOMServerFormatConfig.js
+++ b/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,
diff --git a/packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js b/packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js
index 40ab6b84b26af..223734c0483de 100644
--- a/packages/react-native-renderer/src/server/ReactNativeServerFormatConfig.js
+++ b/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.
diff --git a/packages/react-server/src/ReactFizzNewContext.js b/packages/react-server/src/ReactFizzNewContext.js
index fc5c3e6119902..9c547803cf9a9 100644
--- a/packages/react-server/src/ReactFizzNewContext.js
+++ b/packages/react-server/src/ReactFizzNewContext.js
@@ -9,10 +9,9 @@
 
 import type {ReactContext} from 'shared/ReactTypes';
 
-import invariant from 'shared/invariant';
+import {isPrimaryRenderer} from './ReactServerFormatConfig';
 
-// TODO: Move to format config.
-const isPrimaryRenderer = true;
+import invariant from 'shared/invariant';
 
 let rendererSigil;
 if (__DEV__) {
diff --git a/packages/react-server/src/forks/ReactServerFormatConfig.custom.js b/packages/react-server/src/forks/ReactServerFormatConfig.custom.js
index 58f2d504b46b8..1ab6f95090822 100644
--- a/packages/react-server/src/forks/ReactServerFormatConfig.custom.js
+++ b/packages/react-server/src/forks/ReactServerFormatConfig.custom.js
@@ -29,6 +29,8 @@ export opaque type ResponseState = mixed;
 export opaque type FormatContext = mixed;
 export opaque type SuspenseBoundaryID = mixed;
 
+export const isPrimaryRenderer = false;
+
 export const getChildFormatContext = $$$hostConfig.getChildFormatContext;
 export const createSuspenseBoundaryID = $$$hostConfig.createSuspenseBoundaryID;
 export const pushEmpty = $$$hostConfig.pushEmpty;

From ea8f0dd212fbfe1793bd4a55a2a410552c655148 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Tue, 13 Apr 2021 12:55:09 -0400
Subject: [PATCH 4/9] Wire up more element type matchers

---
 packages/react-server/src/ReactFizzServer.js | 142 +++++++++++++++----
 1 file changed, 115 insertions(+), 27 deletions(-)

diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index 65f42ecd897c9..269d25436fb19 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -13,7 +13,12 @@ import type {
   Chunk,
   PrecomputedChunk,
 } from './ReactServerStreamConfig';
-import type {ReactNodeList} from 'shared/ReactTypes';
+import type {
+  ReactNodeList,
+  ReactContext,
+  ReactProviderType,
+} from 'shared/ReactTypes';
+import type {LazyComponent as LazyComponentType} from 'react/src/ReactLazy';
 import type {
   SuspenseBoundaryID,
   ResponseState,
@@ -75,6 +80,10 @@ import {
   REACT_PROFILER_TYPE,
   REACT_SUSPENSE_LIST_TYPE,
   REACT_FRAGMENT_TYPE,
+  REACT_FORWARD_REF_TYPE,
+  REACT_MEMO_TYPE,
+  REACT_PROVIDER_TYPE,
+  REACT_CONTEXT_TYPE,
 } from 'shared/ReactSymbols';
 import ReactSharedInternals from 'shared/ReactSharedInternals';
 import {
@@ -706,49 +715,128 @@ function validateFunctionComponentInDev(Component: any): void {
   }
 }
 
+function renderForwardRef(
+  request: Request,
+  task: Task,
+  type: any,
+  props: Object,
+): void {
+  throw new Error('Not yet implemented element type.');
+}
+
+function renderMemo(
+  request: Request,
+  task: Task,
+  type: any,
+  props: Object,
+): void {
+  throw new Error('Not yet implemented element type.');
+}
+
+function renderContextConsumer(
+  request: Request,
+  task: Task,
+  type: ReactContext<any>,
+  props: Object,
+): void {
+  throw new Error('Not yet implemented element type.');
+}
+
+function renderContextProvider(
+  request: Request,
+  task: Task,
+  type: ReactProviderType<any>,
+  props: Object,
+): void {
+  throw new Error('Not yet implemented element type.');
+}
+
+function renderLazyComponent(
+  request: Request,
+  task: Task,
+  type: LazyComponentType<any, any>,
+  props: Object,
+): void {
+  throw new Error('Not yet implemented element type.');
+}
+
 function renderElement(
   request: Request,
   task: Task,
   type: any,
   props: Object,
-  node: ReactNodeList,
 ): void {
   if (typeof type === 'function') {
     if (shouldConstruct(type)) {
       renderClassComponent(request, task, type, props);
+      return;
     } else {
       renderIndeterminateComponent(request, task, type, props);
+      return;
     }
-  } else if (typeof type === 'string') {
+  }
+  if (typeof type === 'string') {
     renderHostElement(request, task, type, props);
-  } else {
-    switch (type) {
-      // TODO: LegacyHidden acts the same as a fragment. This only works
-      // because we currently assume that every instance of LegacyHidden is
-      // accompanied by a host component wrapper. In the hidden mode, the host
-      // component is given a `hidden` attribute, which ensures that the
-      // initial HTML is not visible. To support the use of LegacyHidden as a
-      // true fragment, without an extra DOM node, we would have to hide the
-      // initial HTML in some other way.
-      // TODO: Add REACT_OFFSCREEN_TYPE here too with the same capability.
-      case REACT_LEGACY_HIDDEN_TYPE:
-      case REACT_DEBUG_TRACING_MODE_TYPE:
-      case REACT_STRICT_MODE_TYPE:
-      case REACT_PROFILER_TYPE:
-      case REACT_SUSPENSE_LIST_TYPE: // TODO: SuspenseList should control the boundaries.
-      case REACT_FRAGMENT_TYPE: {
-        renderNodeDestructive(request, task, props.children);
-        break;
+    return;
+  }
+
+  switch (type) {
+    // TODO: LegacyHidden acts the same as a fragment. This only works
+    // because we currently assume that every instance of LegacyHidden is
+    // accompanied by a host component wrapper. In the hidden mode, the host
+    // component is given a `hidden` attribute, which ensures that the
+    // initial HTML is not visible. To support the use of LegacyHidden as a
+    // true fragment, without an extra DOM node, we would have to hide the
+    // initial HTML in some other way.
+    // TODO: Add REACT_OFFSCREEN_TYPE here too with the same capability.
+    case REACT_LEGACY_HIDDEN_TYPE:
+    case REACT_DEBUG_TRACING_MODE_TYPE:
+    case REACT_STRICT_MODE_TYPE:
+    case REACT_PROFILER_TYPE:
+    case REACT_SUSPENSE_LIST_TYPE: // TODO: SuspenseList should control the boundaries.
+    case REACT_FRAGMENT_TYPE: {
+      renderNodeDestructive(request, task, props.children);
+      return;
+    }
+    case REACT_SUSPENSE_TYPE: {
+      renderSuspenseBoundary(request, task, props);
+      return;
+    }
+  }
+
+  if (typeof type === 'object' && type !== null) {
+    switch (type.$$typeof) {
+      case REACT_FORWARD_REF_TYPE: {
+        renderForwardRef(request, task, type, props);
+        return;
+      }
+      case REACT_MEMO_TYPE: {
+        renderMemo(request, task, type, props);
+        return;
+      }
+      case REACT_PROVIDER_TYPE: {
+        renderContextProvider(request, task, type, props);
+        return;
       }
-      case REACT_SUSPENSE_TYPE: {
-        renderSuspenseBoundary(request, task, props);
-        break;
+      case REACT_CONTEXT_TYPE: {
+        renderContextConsumer(request, task, type, props);
+        return;
       }
-      default: {
-        throw new Error('Not yet implemented element type.');
+      case REACT_LAZY_TYPE: {
+        renderLazyComponent(request, task, type, props);
+        return;
       }
     }
   }
+
+  invariant(
+    false,
+    'Element type is invalid: expected a string (for built-in ' +
+      'components) or a class/function (for composite components) ' +
+      'but got: %s.%s',
+    type == null ? type : typeof type,
+    '',
+  );
 }
 
 function validateIterable(iterable, iteratorFn: Function): void {
@@ -803,7 +891,7 @@ function renderNodeDestructive(
         const element: React$Element<any> = (node: any);
         const type = element.type;
         const props = element.props;
-        renderElement(request, task, type, props, node);
+        renderElement(request, task, type, props);
         return;
       }
       case REACT_PORTAL_TYPE:

From 0459b9b0834e35330429def10f9f62732fa52f7d Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Tue, 13 Apr 2021 13:09:48 -0400
Subject: [PATCH 5/9] Wire up Context Provider type

---
 .../react-server/src/ReactFizzNewContext.js   |  4 ++--
 packages/react-server/src/ReactFizzServer.js  | 20 ++++++++++++++++++-
 2 files changed, 21 insertions(+), 3 deletions(-)

diff --git a/packages/react-server/src/ReactFizzNewContext.js b/packages/react-server/src/ReactFizzNewContext.js
index 9c547803cf9a9..aed6d3b1e4574 100644
--- a/packages/react-server/src/ReactFizzNewContext.js
+++ b/packages/react-server/src/ReactFizzNewContext.js
@@ -219,7 +219,7 @@ export function pushProvider<T>(
   return newNode;
 }
 
-export function popProvider<T>(context: ReactContext<T>): void {
+export function popProvider<T>(context: ReactContext<T>): ContextSnapshot {
   const prevSnapshot = currentActiveSnapshot;
   invariant(
     prevSnapshot !== null,
@@ -263,7 +263,7 @@ export function popProvider<T>(context: ReactContext<T>): void {
       context._currentRenderer2 = rendererSigil;
     }
   }
-  currentActiveSnapshot = prevSnapshot.parent;
+  return (currentActiveSnapshot = prevSnapshot.parent);
 }
 
 export function getActiveContext(): ContextSnapshot {
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index 269d25436fb19..ecfb3b8d82987 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -66,6 +66,8 @@ import {
   rootContextSnapshot,
   switchContext,
   getActiveContext,
+  pushProvider,
+  popProvider,
 } from './ReactFizzNewContext';
 
 import {
@@ -748,7 +750,23 @@ function renderContextProvider(
   type: ReactProviderType<any>,
   props: Object,
 ): void {
-  throw new Error('Not yet implemented element type.');
+  const context = type._context;
+  const value = props.value;
+  const children = props.children;
+  let prevSnapshot;
+  if (__DEV__) {
+    prevSnapshot = task.context;
+  }
+  task.context = pushProvider(context, value);
+  renderNodeDestructive(request, task, children);
+  task.context = popProvider(context);
+  if (__DEV__) {
+    if (prevSnapshot !== task.context) {
+      console.error(
+        'Popping the context provider did not return back to the original snapshot. This is a bug in React.',
+      );
+    }
+  }
 }
 
 function renderLazyComponent(

From 19fb829675e293909e6eaf69191f0a2c80c7d1b9 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Tue, 13 Apr 2021 13:17:26 -0400
Subject: [PATCH 6/9] Wire up Context Consumer

---
 .../react-server/src/ReactFizzNewContext.js   |  7 +++
 packages/react-server/src/ReactFizzServer.js  | 47 ++++++++++++++++++-
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/packages/react-server/src/ReactFizzNewContext.js b/packages/react-server/src/ReactFizzNewContext.js
index aed6d3b1e4574..43f364717d81f 100644
--- a/packages/react-server/src/ReactFizzNewContext.js
+++ b/packages/react-server/src/ReactFizzNewContext.js
@@ -269,3 +269,10 @@ export function popProvider<T>(context: ReactContext<T>): ContextSnapshot {
 export function getActiveContext(): ContextSnapshot {
   return currentActiveSnapshot;
 }
+
+export function readContext<T>(context: ReactContext<T>): T {
+  const value = isPrimaryRenderer
+    ? context._currentValue
+    : context._currentValue2;
+  return value;
+}
diff --git a/packages/react-server/src/ReactFizzServer.js b/packages/react-server/src/ReactFizzServer.js
index ecfb3b8d82987..c1d51c3b2422e 100644
--- a/packages/react-server/src/ReactFizzServer.js
+++ b/packages/react-server/src/ReactFizzServer.js
@@ -63,6 +63,7 @@ import {
   emptyContextObject,
 } from './ReactFizzContext';
 import {
+  readContext,
   rootContextSnapshot,
   switchContext,
   getActiveContext,
@@ -552,6 +553,7 @@ let didWarnAboutReassigningProps = false;
 const didWarnAboutDefaultPropsOnFunctionComponent = {};
 let didWarnAboutGenerators = false;
 let didWarnAboutMaps = false;
+let hasWarnedAboutUsingContextAsConsumer = false;
 
 // This would typically be a function component but we still support module pattern
 // components for some reason.
@@ -738,10 +740,51 @@ function renderMemo(
 function renderContextConsumer(
   request: Request,
   task: Task,
-  type: ReactContext<any>,
+  context: ReactContext<any>,
   props: Object,
 ): void {
-  throw new Error('Not yet implemented element type.');
+  // The logic below for Context differs depending on PROD or DEV mode. In
+  // DEV mode, we create a separate object for Context.Consumer that acts
+  // like a proxy to Context. This proxy object adds unnecessary code in PROD
+  // so we use the old behaviour (Context.Consumer references Context) to
+  // reduce size and overhead. The separate object references context via
+  // a property called "_context", which also gives us the ability to check
+  // in DEV mode if this property exists or not and warn if it does not.
+  if (__DEV__) {
+    if ((context: any)._context === undefined) {
+      // This may be because it's a Context (rather than a Consumer).
+      // Or it may be because it's older React where they're the same thing.
+      // We only want to warn if we're sure it's a new React.
+      if (context !== context.Consumer) {
+        if (!hasWarnedAboutUsingContextAsConsumer) {
+          hasWarnedAboutUsingContextAsConsumer = true;
+          console.error(
+            'Rendering <Context> directly is not supported and will be removed in ' +
+              'a future major release. Did you mean to render <Context.Consumer> instead?',
+          );
+        }
+      }
+    } else {
+      context = (context: any)._context;
+    }
+  }
+  const render = props.children;
+
+  if (__DEV__) {
+    if (typeof render !== 'function') {
+      console.error(
+        'A context consumer was rendered with multiple children, or a child ' +
+          "that isn't a function. A context consumer expects a single child " +
+          'that is a function. If you did pass a function, make sure there ' +
+          'is no trailing or leading whitespace around it.',
+      );
+    }
+  }
+
+  const newValue = readContext(context);
+  const newChildren = render(newValue);
+
+  renderNodeDestructive(request, task, newChildren);
 }
 
 function renderContextProvider(

From 99600c8844a50ee297afb742cd5222f8cb7068f1 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Tue, 13 Apr 2021 16:13:35 -0400
Subject: [PATCH 7/9] Test

---
 .../src/__tests__/ReactDOMFizzServer-test.js  | 66 +++++++++++++++++++
 1 file changed, 66 insertions(+)

diff --git a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js b/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
index 9346f20f0442e..edbb2e9f2e970 100644
--- a/packages/react-dom/src/__tests__/ReactDOMFizzServer-test.js
+++ b/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>,
+    );
+  });
 });

From 82bdc8bf7b52f7651bde093323a2520b4100bf7e Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Tue, 13 Apr 2021 16:51:27 -0400
Subject: [PATCH 8/9] Implement reader in class

---
 packages/react-server/src/ReactFizzClassComponent.js | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/packages/react-server/src/ReactFizzClassComponent.js b/packages/react-server/src/ReactFizzClassComponent.js
index d56bc633b8498..a7385d54aa781 100644
--- a/packages/react-server/src/ReactFizzClassComponent.js
+++ b/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 {

From 5acef5f076b48f459f77808db99dc967e8701671 Mon Sep 17 00:00:00 2001
From: Sebastian Markbage <sema@fb.com>
Date: Tue, 13 Apr 2021 18:09:21 -0400
Subject: [PATCH 9/9] Update error codez

---
 scripts/error-codes/codes.json | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/scripts/error-codes/codes.json b/scripts/error-codes/codes.json
index 17c28c3bc7018..4dbb95f55e809 100644
--- a/scripts/error-codes/codes.json
+++ b/scripts/error-codes/codes.json
@@ -388,5 +388,8 @@
   "397": "Unknown insertion mode. This is a bug in React.",
   "398": "`dangerouslySetInnerHTML` does not work on <option>.",
   "399": "%s is a self-closing tag and must neither have `children` nor use `dangerouslySetInnerHTML`.",
-  "400": "menuitems cannot have `children` nor `dangerouslySetInnerHTML`."
+  "400": "menuitems cannot have `children` nor `dangerouslySetInnerHTML`.",
+  "401": "The stacks must reach the root at the same time. This is a bug in React.",
+  "402": "The depth must equal at least at zero before reaching the root. This is a bug in React.",
+  "403": "Tried to pop a Context at the root of the app. This is a bug in React."
 }