[useForkRef] Support ref cleanup functions

Author: atomiks

packages/react/src/utils/useForkRef.test.tsx

@@ -1,6 +1,7 @@
 import * as React from 'react';
 import { expect } from 'chai';
 import { createRenderer, MuiRenderResult, screen } from '@mui/internal-test-utils';
+import { spy } from 'sinon';
 import { useForkRef } from './useForkRef';
 import { getReactElementRef } from './getReactElementRef';
 
@@ -123,4 +124,52 @@ describe('useForkRef', () => {
       expect(secondRightRef.current!.id).to.equal('test');
     });
   });
+
+  test('calls clean up function if it exists', () => {
+    const cleanUp = spy();
+    const setup = spy();
+    const setup2 = spy();
+    const nullHandler = spy();
+
+    function onRefChangeWithCleanup(ref: HTMLDivElement | null) {
+      if (ref) {
+        setup(ref.id);
+      } else {
+        nullHandler();
+      }
+      return cleanUp;
+    }
+
+    function onRefChangeWithoutCleanup(ref: HTMLDivElement | null) {
+      if (ref) {
+        setup2(ref.id);
+      } else {
+        nullHandler();
+      }
+    }
+
+    function App() {
+      const ref = useForkRef(onRefChangeWithCleanup, onRefChangeWithoutCleanup);
+      return <div id="test" ref={ref} />;
+    }
+
+    const { unmount } = render(<App />);
+
+    expect(setup.args[0][0]).to.equal('test');
+    expect(setup.callCount).to.equal(1);
+    expect(cleanUp.callCount).to.equal(0);
+
+    expect(setup2.args[0][0]).to.equal('test');
+    expect(setup2.callCount).to.equal(1);
+
+    unmount();
+
+    expect(setup.callCount).to.equal(1);
+    expect(cleanUp.callCount).to.equal(1);
+
+    // Setup was not called again
+    expect(setup2.callCount).to.equal(1);
+    // Null handler hit because no cleanup is returned
+    expect(nullHandler.callCount).to.equal(1);
+  });
 });

packages/react/src/utils/useForkRef.ts

@@ -1,52 +1,56 @@
-'use client';
 import * as React from 'react';
 
 /**
- * Takes an array of refs and returns a new ref which will apply any modification to all of the refs.
- * This is useful when you want to have the ref used in multiple places.
- *
- * ```tsx
- * const rootRef = React.useRef<Instance>(null);
- * const refFork = useForkRef(rootRef, props.ref);
- *
- * return (
- *   <Root {...props} ref={refFork} />
- * );
- * ```
- *
- * @param {Array<React.Ref<Instance> | undefined>} refs The ref array.
- * @returns {React.RefCallback<Instance> | null} The new ref callback.
+ * Merges refs into a single memoized callback ref or `null`.
  */
 export function useForkRef<Instance>(
   ...refs: Array<React.Ref<Instance> | undefined>
-): React.RefCallback<Instance> | null {
-  /**
-   * This will create a new function if the refs passed to this hook change and are all defined.
-   * This means react will call the old forkRef with `null` and the new forkRef
-   * with the ref. Cleanup naturally emerges from this behavior.
-   */
+): null | React.RefCallback<Instance> {
+  const cleanupRef = React.useRef<void | (() => void)>(undefined);
+
+  const refEffect = React.useCallback((instance: Instance | null) => {
+    const cleanups = refs.map((ref) => {
+      if (ref == null) {
+        return null;
+      }
+
+      if (typeof ref === 'function') {
+        const refCallback = ref;
+        const refCleanup: void | (() => void) = refCallback(instance);
+        return typeof refCleanup === 'function'
+          ? refCleanup
+          : () => {
+              refCallback(null);
+            };
+      }
+
+      (ref as React.RefObject<Instance | null>).current = instance;
+      return () => {
+        (ref as React.RefObject<Instance | null>).current = null;
+      };
+    });
+
+    return () => {
+      cleanups.forEach((refCleanup) => refCleanup?.());
+    };
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, refs);
+
   return React.useMemo(() => {
     if (refs.every((ref) => ref == null)) {
       return null;
     }
 
-    return (instance) => {
-      refs.forEach((ref) => {
-        setRef(ref, instance);
-      });
+    return (value) => {
+      if (cleanupRef.current) {
+        cleanupRef.current();
+        (cleanupRef as React.RefObject<void | (() => void)>).current = undefined;
+      }
+
+      if (value != null) {
+        (cleanupRef as React.RefObject<void | (() => void)>).current = refEffect(value);
+      }
     };
-    // TODO: uncomment once we enable eslint-plugin-react-compiler // eslint-disable-next-line react-compiler/react-compiler -- intentionally ignoring that the dependency array must be an array literal
     // eslint-disable-next-line react-hooks/exhaustive-deps
   }, refs);
 }
-
-function setRef<T>(
-  ref: React.RefObject<T | null> | ((instance: T | null) => void) | null | undefined,
-  value: T | null,
-): void {
-  if (typeof ref === 'function') {
-    ref(value);
-  } else if (ref) {
-    ref.current = value;
-  }
-}