useMergeRefs: call when dependency changes after ref change (#29892)

packages/compose/src/hooks/use-merge-refs/index.js

@@ -47,9 +47,14 @@ export default function useMergeRefs( refs ) {
 		} );
 
 		previousRefs.current = refs;
-		didElementChange.current = false;
 	}, refs );
 
+	// No dependencies, must be reset after every render so ref callbacks are
+	// correctly called after a ref change.
+	useLayoutEffect( () => {
+		didElementChange.current = false;
+	} );
+
 	// There should be no dependencies so that `callback` is only called when
 	// the node changes.
 	return useCallback( ( value ) => {

packages/compose/src/hooks/use-merge-refs/test/index.js

@@ -218,4 +218,59 @@ describe( 'useMergeRefs', () => {
 			[ [], [ newElement, null ] ],
 		] );
 	} );
+
+	it( 'should work for dependency change after node change', () => {
+		const rootElement = document.getElementById( 'root' );
+
+		ReactDOM.render( <MergedRefs />, rootElement );
+
+		const originalElement = rootElement.firstElementChild;
+
+		ReactDOM.render( <MergedRefs tagName="button" />, rootElement );
+
+		const newElement = rootElement.firstElementChild;
+
+		// After a render with the original element and a second render with the
+		// new element, expect the initial callback functions to be called with
+		// the original element, then null, then the new element.
+		// Again, the new callback functions should not be called! There has
+		// been no dependency change.
+		expect( renderCallback.history ).toEqual( [
+			[
+				[ originalElement, null, newElement ],
+				[ originalElement, null, newElement ],
+			],
+			[ [], [] ],
+		] );
+
+		ReactDOM.render(
+			<MergedRefs tagName="button" count={ 1 } />,
+			rootElement
+		);
+
+		// After a third render with a dependency change, expect the inital
+		// callback function to be called with null and the new callback
+		// function to be called with the new element. Note that for callback
+		// one no dependencies have changed.
+		expect( renderCallback.history ).toEqual( [
+			[
+				[ originalElement, null, newElement ],
+				[ originalElement, null, newElement, null ],
+			],
+			[ [], [] ],
+			[ [], [ newElement ] ],
+		] );
+
+		ReactDOM.render( null, rootElement );
+
+		// Unmount: current callback functions should be called with null.
+		expect( renderCallback.history ).toEqual( [
+			[
+				[ originalElement, null, newElement, null ],
+				[ originalElement, null, newElement, null ],
+			],
+			[ [], [] ],
+			[ [], [ newElement, null ] ],
+		] );
+	} );
 } );