[utils] Support cleanup callbacks in useForkRef (@DiegoAndai) (#45733)

Co-authored-by: Diego Andai <[email protected]>
Co-authored-by: Aarón García Hervás <[email protected]>

Author: 3 people

Diff for: packages/mui-utils/src/useForkRef/useForkRef.test.js

@@ -1,5 +1,6 @@
 import * as React from 'react';
 import { expect } from 'chai';
+import { spy } from 'sinon';
 import { createRenderer, screen } from '@mui/internal-test-utils';
 import useForkRef from './useForkRef';
 import getReactElementRef from '../getReactElementRef';
@@ -108,4 +109,52 @@ describe('useForkRef', () => {
       expect(secondRightRef.current.id).to.equal('test');
     });
   });
+
+  it('calls clean up function if it exists', () => {
+    const cleanUp = spy();
+    const setup = spy();
+    const setup2 = spy();
+    const nullHandler = spy();
+
+    function onRefChangeWithCleanup(ref) {
+      if (ref) {
+        setup(ref.id);
+      } else {
+        nullHandler();
+      }
+      return cleanUp;
+    }
+
+    function onRefChangeWithoutCleanup(ref) {
+      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);
+  });
 });

Diff for: packages/mui-utils/src/useForkRef/useForkRef.ts

@@ -1,10 +1,8 @@
 'use client';
 import * as React from 'react';
-import setRef from '../setRef';
 
 /**
- * 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.
+ * Merges refs into a single memoized callback ref or `null`.
  *
  * ```tsx
  * const rootRef = React.useRef<Instance>(null);
@@ -21,20 +19,50 @@ import setRef from '../setRef';
 export default 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.
-   */
+  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