From ce4b3e998191445b2934f6846ee7661578ef6444 Mon Sep 17 00:00:00 2001
From: Dominic Gannaway <trueadm@users.noreply.github.com>
Date: Wed, 6 Nov 2019 22:52:59 +0000
Subject: [PATCH] [react-interactions] Add optional searchNodes to
 Scope.queryAllNodes (#17293)

---
 .../react-reconciler/src/ReactFiberScope.js   | 18 +++++---
 .../src/__tests__/ReactScope-test.internal.js | 44 +++++++++++++++++++
 packages/shared/ReactTypes.js                 |  1 +
 3 files changed, 58 insertions(+), 5 deletions(-)

diff --git a/packages/react-reconciler/src/ReactFiberScope.js b/packages/react-reconciler/src/ReactFiberScope.js
index 34eb261f24d2e..06f5a223299a3 100644
--- a/packages/react-reconciler/src/ReactFiberScope.js
+++ b/packages/react-reconciler/src/ReactFiberScope.js
@@ -37,12 +37,17 @@ function collectScopedNodes(
   node: Fiber,
   fn: (type: string | Object, props: Object) => boolean,
   scopedNodes: Array<any>,
+  searchNode: null | Set<Object>,
 ): void {
   if (enableScopeAPI) {
     if (node.tag === HostComponent) {
+      const instance = getPublicInstance(node.stateNode);
       const {type, memoizedProps} = node;
-      if (fn(type, memoizedProps || emptyObject) === true) {
-        scopedNodes.push(getPublicInstance(node.stateNode));
+      if (
+        (searchNode !== null && searchNode.has(instance)) ||
+        fn(type, memoizedProps || emptyObject) === true
+      ) {
+        scopedNodes.push(instance);
       }
     }
     let child = node.child;
@@ -51,7 +56,7 @@ function collectScopedNodes(
       child = getSuspenseFallbackChild(node);
     }
     if (child !== null) {
-      collectScopedNodesFromChildren(child, fn, scopedNodes);
+      collectScopedNodesFromChildren(child, fn, scopedNodes, searchNode);
     }
   }
 }
@@ -83,10 +88,11 @@ function collectScopedNodesFromChildren(
   startingChild: Fiber,
   fn: (type: string | Object, props: Object) => boolean,
   scopedNodes: Array<any>,
+  searchNode: null | Set<Object>,
 ): void {
   let child = startingChild;
   while (child !== null) {
-    collectScopedNodes(child, fn, scopedNodes);
+    collectScopedNodes(child, fn, scopedNodes, searchNode);
     child = child.sibling;
   }
 }
@@ -192,12 +198,14 @@ export function createScopeMethods(
     },
     queryAllNodes(
       fn: (type: string | Object, props: Object) => boolean,
+      searchNodes?: Array<Object>,
     ): null | Array<Object> {
       const currentFiber = ((instance.fiber: any): Fiber);
       const child = currentFiber.child;
       const scopedNodes = [];
+      const searchNodeSet = searchNodes ? new Set(searchNodes) : null;
       if (child !== null) {
-        collectScopedNodesFromChildren(child, fn, scopedNodes);
+        collectScopedNodesFromChildren(child, fn, scopedNodes, searchNodeSet);
       }
       return scopedNodes.length === 0 ? null : scopedNodes;
     },
diff --git a/packages/react-reconciler/src/__tests__/ReactScope-test.internal.js b/packages/react-reconciler/src/__tests__/ReactScope-test.internal.js
index 4e3737c9f3f42..8a62c2decf38e 100644
--- a/packages/react-reconciler/src/__tests__/ReactScope-test.internal.js
+++ b/packages/react-reconciler/src/__tests__/ReactScope-test.internal.js
@@ -72,6 +72,50 @@ describe('ReactScope', () => {
       expect(scopeRef.current).toBe(null);
     });
 
+    it('queryAllNodes() works as intended with included nodes array', () => {
+      const testScopeQuery = (type, props) => type === 'div';
+      const TestScope = React.unstable_createScope();
+      const scopeRef = React.createRef();
+      const divRef = React.createRef();
+      const spanRef = React.createRef();
+      const aRef = React.createRef();
+
+      function Test({toggle}) {
+        return toggle ? (
+          <TestScope ref={scopeRef}>
+            <div ref={divRef}>DIV</div>
+            <span ref={spanRef}>SPAN</span>
+            <a ref={aRef}>A</a>
+          </TestScope>
+        ) : (
+          <TestScope ref={scopeRef}>
+            <a ref={aRef}>A</a>
+            <div ref={divRef}>DIV</div>
+            <span ref={spanRef}>SPAN</span>
+          </TestScope>
+        );
+      }
+
+      ReactDOM.render(<Test toggle={true} />, container);
+      let nodes = scopeRef.current.queryAllNodes(testScopeQuery);
+      expect(nodes).toEqual([divRef.current]);
+      nodes = scopeRef.current.queryAllNodes(testScopeQuery, [spanRef.current]);
+      expect(nodes).toEqual([divRef.current, spanRef.current]);
+      nodes = scopeRef.current.queryAllNodes(testScopeQuery, [
+        spanRef.current,
+        aRef.current,
+      ]);
+      expect(nodes).toEqual([divRef.current, spanRef.current, aRef.current]);
+      ReactDOM.render(<Test toggle={false} />, container);
+      nodes = scopeRef.current.queryAllNodes(testScopeQuery, [
+        spanRef.current,
+        aRef.current,
+      ]);
+      expect(nodes).toEqual([aRef.current, divRef.current, spanRef.current]);
+      ReactDOM.render(null, container);
+      expect(scopeRef.current).toBe(null);
+    });
+
     it('queryFirstNode() works as intended', () => {
       const testScopeQuery = (type, props) => true;
       const TestScope = React.unstable_createScope();
diff --git a/packages/shared/ReactTypes.js b/packages/shared/ReactTypes.js
index 97033ba917035..aaca80d695dc5 100644
--- a/packages/shared/ReactTypes.js
+++ b/packages/shared/ReactTypes.js
@@ -171,6 +171,7 @@ export type ReactScopeMethods = {|
   getProps(): Object,
   queryAllNodes(
     (type: string | Object, props: Object) => boolean,
+    searchNodes?: Array<Object>,
   ): null | Array<Object>,
   queryFirstNode(
     (type: string | Object, props: Object) => boolean,