From 0c52e24cb65a8f1c370184f58ee2d5601a3acd7f Mon Sep 17 00:00:00 2001
From: Brian Vaughn <bvaughn@fb.com>
Date: Mon, 10 Aug 2020 10:49:10 -0400
Subject: [PATCH] Support inner component _debugOwner in memo (#19556)

* Support inner component _debugOwner in memo

* test with devtool context

* remove memo test

* Merged master; tweaked test and snapshot

* Pass owner to createFiber fn when creating a memo component.

Co-authored-by: Theodore Han <tqhan317@gmail.com>
---
 .../ownersListContext-test.js.snap            | 27 +++++++++++++
 .../src/__tests__/ownersListContext-test.js   | 39 +++++++++++++++++++
 .../react-reconciler/src/ReactFiber.new.js    |  4 ++
 .../react-reconciler/src/ReactFiber.old.js    |  4 ++
 .../src/ReactFiberBeginWork.new.js            |  2 +-
 .../src/ReactFiberBeginWork.old.js            |  2 +-
 6 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/packages/react-devtools-shared/src/__tests__/__snapshots__/ownersListContext-test.js.snap b/packages/react-devtools-shared/src/__tests__/__snapshots__/ownersListContext-test.js.snap
index 0699eb5b18dc5..25e56daf6f44e 100644
--- a/packages/react-devtools-shared/src/__tests__/__snapshots__/ownersListContext-test.js.snap
+++ b/packages/react-devtools-shared/src/__tests__/__snapshots__/ownersListContext-test.js.snap
@@ -78,6 +78,33 @@ Array [
 ]
 `;
 
+exports[`OwnersListContext should include all owners for a component wrapped in react memo: owners for "InnerComponent" 1`] = `
+Array [
+  Object {
+    "displayName": "Grandparent",
+    "hocDisplayNames": null,
+    "id": 2,
+    "type": 5,
+  },
+  Object {
+    "displayName": "InnerComponent",
+    "hocDisplayNames": Array [
+      "Memo",
+    ],
+    "id": 3,
+    "type": 8,
+  },
+  Object {
+    "displayName": "InnerComponent",
+    "hocDisplayNames": Array [
+      "ForwardRef",
+    ],
+    "id": 4,
+    "type": 6,
+  },
+]
+`;
+
 exports[`OwnersListContext should include the current element even if there are no other owners: mount 1`] = `
 [root]
     <Grandparent>
diff --git a/packages/react-devtools-shared/src/__tests__/ownersListContext-test.js b/packages/react-devtools-shared/src/__tests__/ownersListContext-test.js
index b39da70632065..4e60bd6a55025 100644
--- a/packages/react-devtools-shared/src/__tests__/ownersListContext-test.js
+++ b/packages/react-devtools-shared/src/__tests__/ownersListContext-test.js
@@ -206,4 +206,43 @@ describe('OwnersListContext', () => {
 
     done();
   });
+
+  it('should include all owners for a component wrapped in react memo', async done => {
+    const InnerComponent = (props, ref) => <div ref={ref} />;
+    const ForwardRef = React.forwardRef(InnerComponent);
+    const Memo = React.memo(ForwardRef);
+    const Grandparent = () => {
+      const ref = React.createRef();
+      return <Memo ref={ref} />;
+    };
+
+    utils.act(() =>
+      ReactDOM.render(<Grandparent />, document.createElement('div')),
+    );
+
+    let didFinish = false;
+    function Suspender({owner}) {
+      const read = React.useContext(OwnersListContext);
+      const owners = read(owner.id);
+      didFinish = true;
+      expect(owners.length).toBe(3);
+      expect(owners).toMatchSnapshot(
+        `owners for "${(owner && owner.displayName) || ''}"`,
+      );
+      return null;
+    }
+
+    const wrapped = ((store.getElementAtIndex(2): any): Element);
+    await utils.actAsync(() =>
+      TestRenderer.create(
+        <Contexts defaultOwnerID={wrapped.id}>
+          <React.Suspense fallback={null}>
+            <Suspender owner={wrapped} />
+          </React.Suspense>
+        </Contexts>,
+      ),
+    );
+    expect(didFinish).toBe(true);
+    done();
+  });
 });
diff --git a/packages/react-reconciler/src/ReactFiber.new.js b/packages/react-reconciler/src/ReactFiber.new.js
index 586825a336209..6a90a307c222c 100644
--- a/packages/react-reconciler/src/ReactFiber.new.js
+++ b/packages/react-reconciler/src/ReactFiber.new.js
@@ -579,6 +579,10 @@ export function createFiberFromTypeAndProps(
   fiber.type = resolvedType;
   fiber.lanes = lanes;
 
+  if (__DEV__) {
+    fiber._debugOwner = owner;
+  }
+
   return fiber;
 }
 
diff --git a/packages/react-reconciler/src/ReactFiber.old.js b/packages/react-reconciler/src/ReactFiber.old.js
index 818f001fa3376..3d896efe8e851 100644
--- a/packages/react-reconciler/src/ReactFiber.old.js
+++ b/packages/react-reconciler/src/ReactFiber.old.js
@@ -573,6 +573,10 @@ export function createFiberFromTypeAndProps(
   fiber.type = resolvedType;
   fiber.lanes = lanes;
 
+  if (__DEV__) {
+    fiber._debugOwner = owner;
+  }
+
   return fiber;
 }
 
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.new.js b/packages/react-reconciler/src/ReactFiberBeginWork.new.js
index 54962f2365d9f..5dbb49e3a7c02 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.new.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.new.js
@@ -434,7 +434,7 @@ function updateMemoComponent(
       Component.type,
       null,
       nextProps,
-      null,
+      workInProgress,
       workInProgress.mode,
       renderLanes,
     );
diff --git a/packages/react-reconciler/src/ReactFiberBeginWork.old.js b/packages/react-reconciler/src/ReactFiberBeginWork.old.js
index 4d7732ab9c543..163aedb4d66bf 100644
--- a/packages/react-reconciler/src/ReactFiberBeginWork.old.js
+++ b/packages/react-reconciler/src/ReactFiberBeginWork.old.js
@@ -434,7 +434,7 @@ function updateMemoComponent(
       Component.type,
       null,
       nextProps,
-      null,
+      workInProgress,
       workInProgress.mode,
       renderLanes,
     );