RN: Support Ref Cleanup in useMergeRefs (#47040)

Summary:
Pull Request resolved: #47040

Updates `useMergeRefs` to support cleanup functions.

Changelog:
[General][Changed] - `useMergeRefs` and components using it (e.g. `Pressable`) now support ref cleanup functions.

Reviewed By: lunaleaps

Differential Revision: D64437947

fbshipit-source-id: b715abfa5b4236c1a7685ac023ff1d0384b6a3a4

packages/react-native/Libraries/Components/ScrollView/ScrollViewStickyHeader.js

@@ -65,8 +65,8 @@ const ScrollViewStickyHeaderWithForwardedRef: component(
     ref.setNextHeaderY = setNextHeaderLayoutY;
     setIsFabric(isFabricPublicInstance(ref));
   }, []);
-  const ref: (React.ElementRef<typeof Animated.View> | null) => void =
-    // $FlowFixMe[incompatible-type] - Ref is mutated by `callbackRef`.
+  const ref: React.RefSetter<React.ElementRef<typeof Animated.View>> =
+    // $FlowFixMe[prop-missing] - Instance is mutated to have `setNextHeaderY`.
     useMergeRefs<Instance>(callbackRef, forwardedRef);
 
   const offset = useMemo(

packages/react-native/Libraries/Utilities/__tests__/useMergeRefs-test.js

@@ -74,6 +74,36 @@ test('accepts a ref callback', () => {
   ]);
 });
 
+test('accepts a ref callback that returns a cleanup function', () => {
+  const screen = new Screen();
+  const ledger: Array<{[string]: string | null}> = [];
+
+  // TODO: Remove `| null` after Flow supports ref cleanup functions.
+  const ref = (current: HostInstance | null) => {
+    ledger.push({ref: id(current)});
+    return () => {
+      ledger.push({ref: null});
+    };
+  };
+
+  screen.render(() => <View id="foo" key="foo" ref={useMergeRefs(ref)} />);
+
+  expect(ledger).toEqual([{ref: 'foo'}]);
+
+  screen.render(() => <View id="bar" key="bar" ref={useMergeRefs(ref)} />);
+
+  expect(ledger).toEqual([{ref: 'foo'}, {ref: null}, {ref: 'bar'}]);
+
+  screen.unmount();
+
+  expect(ledger).toEqual([
+    {ref: 'foo'},
+    {ref: null},
+    {ref: 'bar'},
+    {ref: null},
+  ]);
+});
+
 test('accepts a ref object', () => {
   const screen = new Screen();
   const ledger: Array<{[string]: string | null}> = [];

packages/react-native/Libraries/Utilities/useMergeRefs.js

@@ -8,6 +8,7 @@
  * @format
  */
 
+import useRefEffect from './useRefEffect';
 import * as React from 'react';
 import {useCallback} from 'react';
 
@@ -22,19 +23,37 @@ import {useCallback} from 'react';
  */
 export default function useMergeRefs<Instance>(
   ...refs: $ReadOnlyArray<?React.RefSetter<Instance>>
-): (Instance | null) => void {
-  return useCallback(
-    (current: Instance | null) => {
-      for (const ref of refs) {
-        if (ref != null) {
+): React.RefSetter<Instance> {
+  const refEffect = useCallback(
+    (current: Instance) => {
+      const cleanups: $ReadOnlyArray<void | (() => void)> = refs.map(ref => {
+        if (ref == null) {
+          return undefined;
+        } else {
           if (typeof ref === 'function') {
-            ref(current);
+            // $FlowIssue[incompatible-type] - Flow does not understand ref cleanup.
+            const cleanup: void | (() => void) = ref(current);
+            return typeof cleanup === 'function'
+              ? cleanup
+              : () => {
+                  ref(null);
+                };
           } else {
             ref.current = current;
+            return () => {
+              ref.current = null;
+            };
           }
         }
-      }
+      });
+
+      return () => {
+        for (const cleanup of cleanups) {
+          cleanup?.();
+        }
+      };
     },
     [...refs], // eslint-disable-line react-hooks/exhaustive-deps
   );
+  return useRefEffect(refEffect);
 }

packages/react-native/Libraries/__tests__/__snapshots__/public-api-test.js.snap

@@ -9423,7 +9423,7 @@ exports[`public API should not change unintentionally Libraries/Utilities/useCol
 exports[`public API should not change unintentionally Libraries/Utilities/useMergeRefs.js 1`] = `
 "declare export default function useMergeRefs<Instance>(
   ...refs: $ReadOnlyArray<?React.RefSetter<Instance>>
-): (Instance | null) => void;
+): React.RefSetter<Instance>;
 "
 `;