diff --git a/packages/react-reconciler/src/ReactFiberScope.js b/packages/react-reconciler/src/ReactFiberScope.js index 34eb261f24..06f5a22329 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, + searchNode: null | Set, ): 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, + searchNode: null | Set, ): 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, ): null | Array { 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 4e3737c9f3..8a62c2decf 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 ? ( + +
DIV
+ SPAN + A +
+ ) : ( + + A +
DIV
+ SPAN +
+ ); + } + + ReactDOM.render(, 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(, 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 97033ba917..aaca80d695 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, ): null | Array, queryFirstNode( (type: string | Object, props: Object) => boolean,