From a20455ef276a8427c5b41f3c1503e3d1baf9ea99 Mon Sep 17 00:00:00 2001
From: eps1lon <sebastian.silbermann@vercel.com>
Date: Mon, 2 Sep 2024 11:13:56 +0200
Subject: [PATCH] Ensure  merging refs does not trigger warnings in React 18

test plan:
- link-ref-app
- link-ref-pages
---
 packages/next/src/client/use-merged-ref.ts | 41 ++++++++++++----------
 1 file changed, 23 insertions(+), 18 deletions(-)

diff --git a/packages/next/src/client/use-merged-ref.ts b/packages/next/src/client/use-merged-ref.ts
index 7fce7fb1f0ee12..65bb8dd4dcb428 100644
--- a/packages/next/src/client/use-merged-ref.ts
+++ b/packages/next/src/client/use-merged-ref.ts
@@ -1,29 +1,34 @@
-import { useMemo, type Ref } from 'react'
+import { useMemo, useRef, type Ref } from 'react'
 
+// This is a compatibility hook to support React 18 and 19 refs.
+// In 19, a cleanup function from refs may be returned.
+// In 18, returning a cleanup function creates a warning.
+// Since we take userspace refs, we don't know ahead of time if a cleanup function will be returned.
+// This implements cleanup functions with the old behavior in 18.
+// We know refs are always called alternating with `null` and then `T`.
+// So a call with `null` means we need to call the previous cleanup functions.
 export function useMergedRef<TElement>(
   refA: Ref<TElement>,
   refB: Ref<TElement>
 ): Ref<TElement> {
-  return useMemo(() => mergeRefs(refA, refB), [refA, refB])
-}
+  const cleanupA = useRef<() => void>(() => {})
+  const cleanupB = useRef<() => void>(() => {})
 
-export function mergeRefs<TElement>(
-  refA: Ref<TElement>,
-  refB: Ref<TElement>
-): Ref<TElement> {
-  if (!refA || !refB) {
-    return refA || refB
-  }
-
-  return (current: TElement) => {
-    const cleanupA = applyRef(refA, current)
-    const cleanupB = applyRef(refB, current)
+  return useMemo(() => {
+    if (!refA || !refB) {
+      return refA || refB
+    }
 
-    return () => {
-      cleanupA()
-      cleanupB()
+    return (current: TElement | null): void => {
+      if (current === null) {
+        cleanupA.current()
+        cleanupB.current()
+      } else {
+        cleanupA.current = applyRef(refA, current)
+        cleanupB.current = applyRef(refB, current)
+      }
     }
-  }
+  }, [refA, refB])
 }
 
 function applyRef<TElement>(