[react-interactions] Add optional searchNodes to Scope.queryAllNodes (#17293)

Author: trueadm

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;
  },

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();

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,