feat: Support Ref Cleanup (#7758)

* feat: Support Ref Cleanup

* cleanup

* add some preliminary mergeRef tests

---------

Co-authored-by: Robert Snow <[email protected]>
Co-authored-by: GitHub <[email protected]>

Author: chirokas

packages/@react-aria/utils/src/mergeRefs.ts

@@ -10,23 +10,43 @@
  * governing permissions and limitations under the License.
  */
 
-import {ForwardedRef, MutableRefObject} from 'react';
+import {MutableRefObject, Ref} from 'react';
 
 /**
  * Merges multiple refs into one. Works with either callback or object refs.
  */
-export function mergeRefs<T>(...refs: Array<ForwardedRef<T> | MutableRefObject<T> | null | undefined>): ForwardedRef<T> {
+export function mergeRefs<T>(...refs: Array<Ref<T> | MutableRefObject<T> | null | undefined>): Ref<T> {
   if (refs.length === 1 && refs[0]) {
     return refs[0];
   }
 
   return (value: T | null) => {
-    for (let ref of refs) {
-      if (typeof ref === 'function') {
-        ref(value);
-      } else if (ref != null) {
-        ref.current = value;
-      }
+    let hasCleanup = false;
+
+    const cleanups = refs.map(ref => {
+      const cleanup = setRef(ref, value);
+      hasCleanup ||= typeof cleanup == 'function';
+      return cleanup;
+    });
+
+    if (hasCleanup) {
+      return () => {
+        cleanups.forEach((cleanup, i) => {
+          if (typeof cleanup === 'function') {
+            cleanup();
+          } else {
+            setRef(refs[i], null);
+          }
+        });
+      };
     }
   };
 }
+
+function setRef<T>(ref: Ref<T> | MutableRefObject<T> | null | undefined, value: T) {
+  if (typeof ref === 'function') {
+    return ref(value);
+  } else if (ref != null) {
+    ref.current = value;
+  }
+}

packages/@react-aria/utils/src/useObjectRef.ts

@@ -10,30 +10,60 @@
  * governing permissions and limitations under the License.
  */
 
-import {MutableRefObject, useMemo, useRef} from 'react';
+import {MutableRefObject, useCallback, useMemo, useRef} from 'react';
 
 /**
  * Offers an object ref for a given callback ref or an object ref. Especially
  * helfpul when passing forwarded refs (created using `React.forwardRef`) to
  * React Aria hooks.
  *
- * @param forwardedRef The original ref intended to be used.
+ * @param ref The original ref intended to be used.
  * @returns An object ref that updates the given ref.
- * @see https://reactjs.org/docs/forwarding-refs.html
+ * @see https://react.dev/reference/react/forwardRef
  */
-export function useObjectRef<T>(forwardedRef?: ((instance: T | null) => void) | MutableRefObject<T | null> | null): MutableRefObject<T | null> {
+export function useObjectRef<T>(ref?: ((instance: T | null) => (() => void) | void) | MutableRefObject<T | null> | null): MutableRefObject<T | null> {
   const objRef: MutableRefObject<T | null> = useRef<T>(null);
-  return useMemo(() => ({
-    get current() {
-      return objRef.current;
+  const cleanupRef: MutableRefObject<(() => void) | void> = useRef(undefined);
+
+  const refEffect = useCallback(
+    (instance: T | null) => {
+      if (typeof ref === 'function') {
+        const refCallback = ref;
+        const refCleanup = refCallback(instance);
+        return () => {
+          if (typeof refCleanup === 'function') {
+            refCleanup();
+          } else {
+            refCallback(null);
+          }
+        };
+      } else if (ref) {
+        ref.current = instance;
+        return () => {
+          ref.current = null;
+        };
+      }
     },
-    set current(value) {
-      objRef.current = value;
-      if (typeof forwardedRef === 'function') {
-        forwardedRef(value);
-      } else if (forwardedRef) {
-        forwardedRef.current = value;
+    [ref]
+  );
+
+  return useMemo(
+    () => ({
+      get current() {
+        return objRef.current;
+      },
+      set current(value) {
+        objRef.current = value;
+        if (cleanupRef.current) {
+          cleanupRef.current();
+          cleanupRef.current = undefined;
+        }
+
+        if (value != null) {
+          cleanupRef.current = refEffect(value);
+        }
       }
-    }
-  }), [forwardedRef]);
+    }),
+    [refEffect]
+  );
 }

packages/@react-aria/utils/test/mergeRefs.test.tsx

@@ -0,0 +1,65 @@
+/*
+ * Copyright 2021 Adobe. All rights reserved.
+ * This file is licensed to you under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License. You may obtain a copy
+ * of the License at http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under
+ * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+ * OF ANY KIND, either express or implied. See the License for the specific language
+ * governing permissions and limitations under the License.
+ */
+
+import {mergeRefs} from '../';
+import React, {useCallback, useRef} from 'react';
+import {render} from '@react-spectrum/test-utils-internal';
+
+describe('mergeRefs', () => {
+  it('merge Refs', () => {
+    let ref1;
+    let ref2;
+
+    const TextField = (props) => {
+      ref1 = useRef(null);
+      ref2 = useRef(null);
+
+      const ref = mergeRefs(ref1, ref2);
+      return <input {...props} ref={ref} />;
+    };
+
+    render(<TextField foo="foo" />);
+
+    expect(ref1.current).toBe(ref2.current);
+  });
+
+  if (parseInt(React.version.split('.')[0], 10) >= 19) {
+    it('merge Ref Cleanup', () => {
+      const cleanUp = jest.fn();
+      let ref1;
+      let ref2;
+      let target = null;
+
+      const TextField = (props) => {
+        ref1 = useRef(null);
+        ref2 = useRef(null);
+        let ref3 = useCallback((node) => {
+          target = node;
+          return cleanUp;
+        }, []);
+
+        const ref = mergeRefs(ref1, ref2, ref3);
+        return <input {...props} ref={ref} />;
+      };
+
+      const {unmount} = render(<TextField foo="foo" />);
+
+      expect(cleanUp).toHaveBeenCalledTimes(0);
+      expect(ref1.current).toBe(target);
+      expect(ref2.current).toBe(target);
+      unmount();
+
+      // Now cleanup has been called
+      expect(cleanUp).toHaveBeenCalledTimes(1);
+    });
+  }
+});

packages/@react-aria/utils/test/useObjectRef.test.js

@@ -146,4 +146,38 @@ describe('useObjectRef', () => {
       expect(screen.getAllByPlaceholderText(/foo/i)).toHaveLength(2);
     });
   });
+
+  it('calls cleanup function on unmount', () => {
+    const cleanUp = jest.fn();
+    const setup = jest.fn();
+    const nullHandler = jest.fn();
+
+    function ref(_ref) {
+      if (_ref) {
+        setup();
+      } else {
+        nullHandler();
+      }
+      return cleanUp;
+    }
+
+    const TextField = React.forwardRef((props, forwardedRef) => {
+      const ref = useObjectRef(forwardedRef);
+      return <input {...props} ref={ref} />;
+    });
+
+    const {unmount} = render(<TextField ref={ref} />);
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    expect(cleanUp).toHaveBeenCalledTimes(0);
+    expect(nullHandler).toHaveBeenCalledTimes(0);
+
+    unmount();
+
+    expect(setup).toHaveBeenCalledTimes(1);
+    // Now cleanup has been called
+    expect(cleanUp).toHaveBeenCalledTimes(1);
+    // Ref callback never called with null when cleanup is returned
+    expect(nullHandler).toHaveBeenCalledTimes(0);
+  });
 });