Support Context as renderable node

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

@@ -5420,6 +5420,34 @@ describe('ReactDOMFizzServer', () => {
 
       expect(getVisibleChildren(container)).toEqual('Hi');
     });
+
+    it('context as node', async () => {
+      const Context = React.createContext('Hi');
+      await act(async () => {
+        const {pipe} = renderToPipeableStream(Context);
+        pipe(writable);
+      });
+      expect(getVisibleChildren(container)).toEqual('Hi');
+    });
+
+    it('recursive Usable as node', async () => {
+      const Context = React.createContext('Hi');
+      const promiseForContext = Promise.resolve(Context);
+      await act(async () => {
+        const {pipe} = renderToPipeableStream(promiseForContext);
+        pipe(writable);
+      });
+
+      // TODO: The `act` implementation in this file doesn't unwrap microtasks
+      // automatically. We can't use the same `act` we use for Fiber tests
+      // because that relies on the mock Scheduler. Doesn't affect any public
+      // API but we might want to fix this for our own internal tests.
+      await act(async () => {
+        await promiseForContext;
+      });
+
+      expect(getVisibleChildren(container)).toEqual('Hi');
+    });
   });
 
   describe('useEffectEvent', () => {

packages/react-reconciler/src/ReactChildFiber.js

@@ -8,7 +8,7 @@
  */
 
 import type {ReactElement} from 'shared/ReactElementType';
-import type {ReactPortal, Thenable} from 'shared/ReactTypes';
+import type {ReactPortal, Thenable, ReactContext} from 'shared/ReactTypes';
 import type {Fiber} from './ReactInternalTypes';
 import type {Lanes} from './ReactFiberLane';
 import type {ThenableState} from './ReactFiberThenable';
@@ -45,6 +45,7 @@ import {isCompatibleFamilyForHotReloading} from './ReactFiberHotReloading';
 import {getIsHydrating} from './ReactFiberHydrationContext';
 import {pushTreeFork} from './ReactFiberTreeContext';
 import {createThenableState, trackUsedThenable} from './ReactFiberThenable';
+import {readContextDuringReconcilation} from './ReactFiberNewContext';
 
 // This tracks the thenables that are unwrapped during reconcilation.
 let thenableState: ThenableState | null = null;
@@ -580,7 +581,12 @@ function createChildReconciler(
         newChild.$$typeof === REACT_CONTEXT_TYPE ||
         newChild.$$typeof === REACT_SERVER_CONTEXT_TYPE
       ) {
-        // TODO: Implement Context as child type.
+        const context: ReactContext<mixed> = (newChild: any);
+        return createChild(
+          returnFiber,
+          readContextDuringReconcilation(returnFiber, context, lanes),
+          lanes,
+        );
       }
 
       throwOnInvalidObjectType(returnFiber, newChild);
@@ -665,7 +671,13 @@ function createChildReconciler(
         newChild.$$typeof === REACT_CONTEXT_TYPE ||
         newChild.$$typeof === REACT_SERVER_CONTEXT_TYPE
       ) {
-        // TODO: Implement Context as child type.
+        const context: ReactContext<mixed> = (newChild: any);
+        return updateSlot(
+          returnFiber,
+          oldFiber,
+          readContextDuringReconcilation(returnFiber, context, lanes),
+          lanes,
+        );
       }
 
       throwOnInvalidObjectType(returnFiber, newChild);
@@ -748,7 +760,14 @@ function createChildReconciler(
         newChild.$$typeof === REACT_CONTEXT_TYPE ||
         newChild.$$typeof === REACT_SERVER_CONTEXT_TYPE
       ) {
-        // TODO: Implement Context as child type.
+        const context: ReactContext<mixed> = (newChild: any);
+        return updateFromMap(
+          existingChildren,
+          returnFiber,
+          newIdx,
+          readContextDuringReconcilation(returnFiber, context, lanes),
+          lanes,
+        );
       }
 
       throwOnInvalidObjectType(returnFiber, newChild);
@@ -1427,7 +1446,13 @@ function createChildReconciler(
         newChild.$$typeof === REACT_CONTEXT_TYPE ||
         newChild.$$typeof === REACT_SERVER_CONTEXT_TYPE
       ) {
-        // TODO: Implement Context as child type.
+        const context: ReactContext<mixed> = (newChild: any);
+        return reconcileChildFibersImpl(
+          returnFiber,
+          currentFirstChild,
+          readContextDuringReconcilation(returnFiber, context, lanes),
+          lanes,
+        );
       }
 
       throwOnInvalidObjectType(returnFiber, newChild);

packages/react-reconciler/src/ReactFiberNewContext.js

@@ -688,7 +688,24 @@ export function readContext<T>(context: ReactContext<T>): T {
       );
     }
   }
+  return readContextForConsumer(currentlyRenderingFiber, context);
+}
 
+export function readContextDuringReconcilation<T>(
+  consumer: Fiber,
+  context: ReactContext<T>,
+  renderLanes: Lanes,
+): T {
+  if (currentlyRenderingFiber === null) {
+    prepareToReadContext(consumer, renderLanes);
+  }
+  return readContextForConsumer(consumer, context);
+}
+
+function readContextForConsumer<T>(
+  consumer: Fiber | null,
+  context: ReactContext<T>,
+): T {
   const value = isPrimaryRenderer
     ? context._currentValue
     : context._currentValue2;
@@ -703,7 +720,7 @@ export function readContext<T>(context: ReactContext<T>): T {
     };
 
     if (lastContextDependency === null) {
-      if (currentlyRenderingFiber === null) {
+      if (consumer === null) {
         throw new Error(
           'Context can only be read while React is rendering. ' +
             'In classes, you can read it in the render method or getDerivedStateFromProps. ' +
@@ -714,12 +731,12 @@ export function readContext<T>(context: ReactContext<T>): T {
 
       // This is the first dependency for this component. Create a new list.
       lastContextDependency = contextItem;
-      currentlyRenderingFiber.dependencies = {
+      consumer.dependencies = {
         lanes: NoLanes,
         firstContext: contextItem,
       };
       if (enableLazyContextPropagation) {
-        currentlyRenderingFiber.flags |= NeedsPropagation;
+        consumer.flags |= NeedsPropagation;
       }
     } else {
       // Append a new context item.

packages/react-reconciler/src/__tests__/ReactUse-test.js

@@ -1381,4 +1381,111 @@ describe('ReactUse', () => {
     assertLog(['B', 'A', 'C']);
     expect(root).toMatchRenderedOutput('BAC');
   });
+
+  test('basic Context as node', async () => {
+    const Context = React.createContext(null);
+
+    function Indirection({children}) {
+      Scheduler.log('Indirection');
+      return children;
+    }
+
+    function ParentOfContextNode() {
+      Scheduler.log('ParentOfContextNode');
+      return Context;
+    }
+
+    function Child({text}) {
+      useEffect(() => {
+        Scheduler.log('Mount');
+        return () => {
+          Scheduler.log('Unmount');
+        };
+      }, []);
+      return <Text text={text} />;
+    }
+
+    function App({contextValue, children}) {
+      const memoizedChildren = useMemo(
+        () => (
+          <Indirection>
+            <ParentOfContextNode />
+          </Indirection>
+        ),
+        [children],
+      );
+      return (
+        <Context.Provider value={contextValue}>
+          {memoizedChildren}
+        </Context.Provider>
+      );
+    }
+
+    // Initial render
+    const root = ReactNoop.createRoot();
+    await act(() => {
+      root.render(<App contextValue={<Child text="A" />} />);
+    });
+    assertLog(['Indirection', 'ParentOfContextNode', 'A', 'Mount']);
+    expect(root).toMatchRenderedOutput('A');
+
+    // Update the child to a new value
+    await act(async () => {
+      root.render(<App contextValue={<Child text="B" />} />);
+    });
+    assertLog([
+      // Notice that the <Indirection /> did not rerender, because the
+      // update was sent via Context.
+
+      // TODO: We shouldn't have to re-render the parent of the context node.
+      // This happens because we need to reconcile the parent's children again.
+      // However, we should be able to skip directly to reconcilation without
+      // evaluating the component. One way to do this might be to mark the
+      // context dependency with a flag that says it was added
+      // during reconcilation.
+      'ParentOfContextNode',
+
+      // Notice that this was an update, not a remount.
+      'B',
+    ]);
+    expect(root).toMatchRenderedOutput('B');
+
+    // Delete the old child and replace it with a new one, by changing the key
+    await act(async () => {
+      root.render(<App contextValue={<Child key="C" text="C" />} />);
+    });
+    assertLog([
+      'ParentOfContextNode',
+
+      // A new instance is mounted
+      'C',
+      'Unmount',
+      'Mount',
+    ]);
+  });
+
+  test('context as node, at the root', async () => {
+    const Context = React.createContext(<Text text="Hi" />);
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      startTransition(() => {
+        root.render(Context);
+      });
+    });
+    assertLog(['Hi']);
+    expect(root).toMatchRenderedOutput('Hi');
+  });
+
+  test('promises that resolves to a context, rendered as a node', async () => {
+    const Context = React.createContext(<Text text="Hi" />);
+    const promise = Promise.resolve(Context);
+    const root = ReactNoop.createRoot();
+    await act(async () => {
+      startTransition(() => {
+        root.render(promise);
+      });
+    });
+    assertLog(['Hi']);
+    expect(root).toMatchRenderedOutput('Hi');
+  });
 });

packages/react-server/src/ReactFizzServer.js

@@ -1467,7 +1467,13 @@ function renderNodeDestructiveImpl(
       maybeUsable.$$typeof === REACT_CONTEXT_TYPE ||
       maybeUsable.$$typeof === REACT_SERVER_CONTEXT_TYPE
     ) {
-      // TODO: Implement Context as child type.
+      const context: ReactContext<ReactNodeList> = (maybeUsable: any);
+      return renderNodeDestructiveImpl(
+        request,
+        task,
+        null,
+        readContext(context),
+      );
     }
 
     // $FlowFixMe[method-unbinding]