From 460793740c47a16f46f875cf0c89b6de85cfb14c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mateusz=20Burzy=C5=84ski?= <mateuszburzynski@gmail.com>
Date: Fri, 15 Apr 2022 00:04:39 +0200
Subject: [PATCH] Fixed an issue with simple memo components being updated with
 new props during context change

---
 .../src/ReactFiberBeginWork.new.js            |  3 +-
 .../src/ReactFiberBeginWork.old.js            |  3 +-
 .../src/__tests__/ReactMemo-test.js           | 45 +++++++++++++++++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.new.js b/packages/react-reconciler/src/ReactFiberBeginWork.new.js
index 88876f05f7d8..ace4f1b4c708 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.new.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.new.js
@@ -592,8 +592,9 @@ function updateSimpleMemoComponent(
   }
   if (current !== null) {
     const prevProps = current.memoizedProps;
+    nextProps = shallowEqual(prevProps, nextProps) ? prevProps : nextProps;
     if (
-      shallowEqual(prevProps, nextProps) &&
+      prevProps === nextProps &&
       current.ref === workInProgress.ref &&
       // Prevent bailout if the implementation changed due to hot reload.
       (__DEV__ ? workInProgress.type === current.type : true)
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.old.js b/packages/react-reconciler/src/ReactFiberBeginWork.old.js
index 418b714117e8..a22e3b377ced 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.old.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.old.js
@@ -592,8 +592,9 @@ function updateSimpleMemoComponent(
   }
   if (current !== null) {
     const prevProps = current.memoizedProps;
+    nextProps = shallowEqual(prevProps, nextProps) ? prevProps : nextProps;
     if (
-      shallowEqual(prevProps, nextProps) &&
+      prevProps === nextProps &&
       current.ref === workInProgress.ref &&
       // Prevent bailout if the implementation changed due to hot reload.
       (__DEV__ ? workInProgress.type === current.type : true)
diff --git a/packages/react-reconciler/src/__tests__/ReactMemo-test.js b/packages/react-reconciler/src/__tests__/ReactMemo-test.js
index 0bbab43d33e0..7dd0b84609d9 100644
--- a/packages/react-reconciler/src/__tests__/ReactMemo-test.js
+++ b/packages/react-reconciler/src/__tests__/ReactMemo-test.js
@@ -179,6 +179,51 @@ describe('memo', () => {
         expect(ReactNoop.getChildren()).toEqual([span('Count: 1')]);
       });
 
+      it('bails out on props referential equality during context change', async () => {
+        const CountContext = React.createContext(0);
+        let renderCount = 0;
+
+        function Inner(props) {
+          React.useContext(CountContext);
+          return React.useMemo(() => {
+            const text = `Inner render count: ${++renderCount}`;
+            return <Text text={text} />;
+          }, [props]);
+        }
+
+        Inner = memo(Inner);
+
+        function Outer(props) {
+          React.useContext(CountContext);
+          return <Inner />;
+        }
+
+        function Parent({value}) {
+          return (
+            <Suspense fallback={<Text text="Loading..." />}>
+              <CountContext.Provider value={value}>
+                <Outer />
+              </CountContext.Provider>
+            </Suspense>
+          );
+        }
+
+        let ctxValue = 0;
+        ReactNoop.render(<Parent value={ctxValue++} />);
+        expect(Scheduler).toFlushAndYield(['Loading...']);
+        await Promise.resolve();
+        expect(Scheduler).toFlushAndYield(['Inner render count: 1']);
+        expect(ReactNoop.getChildren()).toEqual([
+          span('Inner render count: 1'),
+        ]);
+
+        ReactNoop.render(<Parent value={ctxValue++} />);
+        expect(Scheduler).toFlushAndYield([]);
+        expect(ReactNoop.getChildren()).toEqual([
+          span('Inner render count: 1'),
+        ]);
+      });
+
       it('accepts custom comparison function', async () => {
         function Counter({count}) {
           return <Text text={count} />;