fix: merged ref behavior in react 19 (#102)

* fix: fix merged ref behavior in react 19

* Update src/useMergedRefs.ts

Co-authored-by: Kyle Tsang <[email protected]>

---------

Co-authored-by: Kyle Tsang <[email protected]>

Author: jquense

src/useMergedRefs.ts

@@ -1,22 +1,42 @@
-import { useMemo } from 'react'
+import { RefCallback, useMemo, version } from 'react'
+
+const isReact19 = parseInt(version.split('.')[0]!, 10) >= 19
 
 type CallbackRef<T> = (ref: T | null) => void
 type Ref<T> = React.MutableRefObject<T> | CallbackRef<T>
 
 function toFnRef<T>(ref?: Ref<T> | null) {
-  return !ref || typeof ref === 'function'
-    ? ref
-    : (value: T | null) => {
-        ref.current = value as T
-      }
+  if (!ref || typeof ref === 'function') {
+    return ref
+  }
+
+  return (value: T | null) => {
+    ref.current = value as T
+  }
+}
+
+function cleanUp<T>(cleanup: unknown, ref: RefCallback<T> | null | undefined) {
+  if (typeof cleanup === 'function') {
+    cleanup()
+  } else if (ref) {
+    ref(null)
+  }
 }
 
 export function mergeRefs<T>(refA?: Ref<T> | null, refB?: Ref<T> | null) {
-  const a = toFnRef(refA)
-  const b = toFnRef(refB)
+  const refASetter = toFnRef(refA)
+  const refBSetter = toFnRef(refB)
+
   return (value: T | null) => {
-    if (a) a(value)
-    if (b) b(value)
+    const cleanupA = refASetter?.(value)
+    const cleanupB = refBSetter?.(value)
+
+    if (isReact19) {
+      return () => {
+        cleanUp(cleanupA, refASetter)
+        cleanUp(cleanupB, refBSetter)
+      }
+    }
   }
 }
 
@@ -36,7 +56,10 @@ export function mergeRefs<T>(refA?: Ref<T> | null, refB?: Ref<T> | null) {
  * @param refB A Callback or mutable Ref
  * @category refs
  */
-function useMergedRefs<T>(refA?: Ref<T> | null, refB?: Ref<T> | null) {
+function useMergedRefs<T>(
+  refA?: Ref<T | null> | null,
+  refB?: Ref<T | null> | null,
+) {
   return useMemo(() => mergeRefs(refA, refB), [refA, refB])
 }
 

test/useMergedRefs.test.tsx

@@ -5,7 +5,7 @@ import useMergedRefs from '../src/useMergedRefs.js'
 import { render } from '@testing-library/react'
 
 describe('useMergedRefs', () => {
-  it('should return a function that returns mount state', () => {
+  it('should work with forwardRef', () => {
     let innerRef: HTMLButtonElement
     const outerRef = React.createRef<HTMLButtonElement>()
 
@@ -18,9 +18,80 @@ describe('useMergedRefs', () => {
       return <button ref={mergedRef} {...props} />
     })
 
+    render(<Button ref={outerRef} />)
+
+    expect(innerRef!.tagName).toEqual('BUTTON')
+    expect(outerRef.current!.tagName).toEqual('BUTTON')
+  })
+
+  it('should work with plain function component', () => {
+    let innerRef: HTMLButtonElement
+    const outerRef = React.createRef<HTMLButtonElement>()
+
+    const Button = ({ ref }: { ref: React.Ref<HTMLButtonElement> }) => {
+      const [buttonEl, attachRef] = useCallbackRef<HTMLButtonElement>()
+      innerRef = buttonEl!
+
+      const mergedRef = useMergedRefs(ref, attachRef)
+
+      return <button ref={mergedRef} />
+    }
+
+    render(<Button ref={outerRef} />)
+
+    expect(innerRef!.tagName).toEqual('BUTTON')
+    expect(outerRef.current!.tagName).toEqual('BUTTON')
+  })
+
+  it('should call refs with null when unmounting', () => {
+    let innerRef: HTMLButtonElement | null = null
+
+    const outerRef = React.createRef<HTMLButtonElement>()
+
+    function Button({ ref }: { ref: React.Ref<HTMLButtonElement> }) {
+      const mergedRef = useMergedRefs(ref, (value) => {
+        innerRef = value
+      })
+
+      return <button ref={mergedRef} />
+    }
+
+    const result = render(<Button ref={outerRef} />)
+
+    expect(innerRef!.tagName).toEqual('BUTTON')
+    expect(outerRef.current!.tagName).toEqual('BUTTON')
+
+    result.unmount()
+
+    expect(innerRef).toBeNull()
+    expect(outerRef.current).toBeNull()
+  })
+
+  it('should call refs cleanup functions', () => {
+    let innerRef: HTMLButtonElement | null = null
+
+    const outerRef = React.createRef<HTMLButtonElement>()
+
+    function Button({ ref }: { ref: React.Ref<HTMLButtonElement> }) {
+      const mergedRef = useMergedRefs(ref, (value) => {
+        innerRef = value
+
+        return () => {
+          innerRef = 'hi' as any
+        }
+      })
+
+      return <button ref={mergedRef} />
+    }
+
     const result = render(<Button ref={outerRef} />)
 
     expect(innerRef!.tagName).toEqual('BUTTON')
     expect(outerRef.current!.tagName).toEqual('BUTTON')
+
+    result.unmount()
+
+    expect(innerRef).toEqual('hi')
+    expect(outerRef.current).toBeNull()
   })
 })