From 1fb16d0c5189dbc5c8962df8eb17a2b628d0c79a Mon Sep 17 00:00:00 2001
From: Sophie Alpert <git@sophiebits.com>
Date: Sun, 7 May 2023 23:45:15 -0700
Subject: [PATCH] Allow wrapping context provider in lazy

Is it OK to `export ContextProvider = Context.Provider;` from a `use client` file? If so, behold.
---
 packages/react-reconciler/src/ReactFiber.js   |  3 ++
 .../src/ReactFiberBeginWork.js                | 37 +++++++++++++++----
 .../src/__tests__/ReactLazy-test.internal.js  | 37 +++++++++++++++++++
 3 files changed, 69 insertions(+), 8 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiber.js b/packages/react-reconciler/src/ReactFiber.js
index 7eefefa1c4d1..9587bedbd98e 100644
--- a/packages/react-reconciler/src/ReactFiber.js
+++ b/packages/react-reconciler/src/ReactFiber.js
@@ -259,6 +259,9 @@ export function resolveLazyComponentTag(Component: Function): WorkTag {
     if ($$typeof === REACT_MEMO_TYPE) {
       return MemoComponent;
     }
+    if ($$typeof === REACT_PROVIDER_TYPE) {
+      return ContextProvider;
+    }
   }
   return IndeterminateComponent;
 }
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.js b/packages/react-reconciler/src/ReactFiberBeginWork.js
index f237a26f9070..112f2501c036 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.js
@@ -1843,6 +1843,15 @@ function mountLazyComponent(
       );
       return child;
     }
+    case ContextProvider: {
+      child = updateContextProvider(
+        null,
+        workInProgress,
+        resolvedProps,
+        renderLanes,
+      );
+      return child;
+    }
   }
   let hint = '';
   if (__DEV__) {
@@ -3487,18 +3496,18 @@ let hasWarnedAboutUsingNoValuePropOnContextProvider = false;
 function updateContextProvider(
   current: Fiber | null,
   workInProgress: Fiber,
+  nextProps: any,
   renderLanes: Lanes,
 ) {
   const providerType: ReactProviderType<any> = workInProgress.type;
   const context: ReactContext<any> = providerType._context;
 
-  const newProps = workInProgress.pendingProps;
   const oldProps = workInProgress.memoizedProps;
 
-  const newValue = newProps.value;
+  const newValue = nextProps.value;
 
   if (__DEV__) {
-    if (!('value' in newProps)) {
+    if (!('value' in nextProps)) {
       if (!hasWarnedAboutUsingNoValuePropOnContextProvider) {
         hasWarnedAboutUsingNoValuePropOnContextProvider = true;
         console.error(
@@ -3509,7 +3518,7 @@ function updateContextProvider(
     const providerPropTypes = workInProgress.type.propTypes;
 
     if (providerPropTypes) {
-      checkPropTypes(providerPropTypes, newProps, 'prop', 'Context.Provider');
+      checkPropTypes(providerPropTypes, nextProps, 'prop', 'Context.Provider');
     }
   }
 
@@ -3526,7 +3535,7 @@ function updateContextProvider(
       if (is(oldValue, newValue)) {
         // No change. Bailout early if children are the same.
         if (
-          oldProps.children === newProps.children &&
+          oldProps.children === nextProps.children &&
           !hasLegacyContextChanged()
         ) {
           return bailoutOnAlreadyFinishedWork(
@@ -3543,7 +3552,7 @@ function updateContextProvider(
     }
   }
 
-  const newChildren = newProps.children;
+  const newChildren = nextProps.children;
   reconcileChildren(current, workInProgress, newChildren, renderLanes);
   return workInProgress.child;
 }
@@ -4184,8 +4193,20 @@ function beginWork(
       return updateMode(current, workInProgress, renderLanes);
     case Profiler:
       return updateProfiler(current, workInProgress, renderLanes);
-    case ContextProvider:
-      return updateContextProvider(current, workInProgress, renderLanes);
+    case ContextProvider: {
+      const type = workInProgress.type;
+      const unresolvedProps = workInProgress.pendingProps;
+      const resolvedProps =
+        workInProgress.elementType === type
+          ? unresolvedProps
+          : resolveDefaultProps(type, unresolvedProps);
+      return updateContextProvider(
+        current,
+        workInProgress,
+        resolvedProps,
+        renderLanes,
+      );
+    }
     case ContextConsumer:
       return updateContextConsumer(current, workInProgress, renderLanes);
     case MemoComponent: {
diff --git a/packages/react-reconciler/src/__tests__/ReactLazy-test.internal.js b/packages/react-reconciler/src/__tests__/ReactLazy-test.internal.js
index a96c60bc26e8..68acc92bbf9c 100644
--- a/packages/react-reconciler/src/__tests__/ReactLazy-test.internal.js
+++ b/packages/react-reconciler/src/__tests__/ReactLazy-test.internal.js
@@ -120,6 +120,43 @@ describe('ReactLazy', () => {
     expect(root).toMatchRenderedOutput('Hi again');
   });
 
+  it('renders a lazy context provider', async () => {
+    const Context = React.createContext('default');
+    const ConsumerText = () => {
+      return <Text text={React.useContext(Context)} />;
+    };
+    const LazyProvider = lazy(() => fakeImport(Context.Provider));
+
+    const root = ReactTestRenderer.create(
+      <Suspense fallback={<Text text="Loading..." />}>
+        <LazyProvider value="Hi">
+          <ConsumerText />
+        </LazyProvider>
+      </Suspense>,
+      {
+        unstable_isConcurrent: true,
+      },
+    );
+
+    await waitForAll(['Loading...']);
+    expect(root).not.toMatchRenderedOutput('Hi');
+
+    await act(() => resolveFakeImport(Context.Provider));
+    assertLog(['Hi']);
+    expect(root).toMatchRenderedOutput('Hi');
+
+    // Should not suspend on update
+    root.update(
+      <Suspense fallback={<Text text="Loading..." />}>
+        <LazyProvider value="Hi again">
+          <ConsumerText />
+        </LazyProvider>
+      </Suspense>,
+    );
+    await waitForAll(['Hi again']);
+    expect(root).toMatchRenderedOutput('Hi again');
+  });
+
   it('can resolve synchronously without suspending', async () => {
     const LazyText = lazy(() => ({
       then(cb) {