Support Context as renderable node (#25641)

## Based on #25634

Like promises, this adds support for Context as a React node.

In this initial implementation, the context dependency is added to the
parent of child node. This allows the parent to re-reconcile its
children when the context updates, so that it can delete the old node if
the identity of the child has changed (i.e. if the key or type of an
element has changed). But it also means that the parent will replay its
entire begin phase. Ideally React would delete the old node and mount
the new node without reconciling all the children. I'll leave this for a
future optimization.

Author: acdlite

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]