fix(templateRef): prevent duplicate onScopeDispose registrations for dynamic template refs (#14161)

Author: edison1105

packages/runtime-vapor/__tests__/dom/templateRef.spec.ts

@@ -182,7 +182,7 @@ describe('api: template ref', () => {
     expect(fn.mock.calls[0][0]).toBe(host.children[0])
     toggle.value = false
     await nextTick()
-    expect(fn.mock.calls[1][0]).toBe(undefined)
+    expect(fn.mock.calls[1][0]).toBe(null)
   })
 
   test('useTemplateRef mount', () => {
@@ -756,6 +756,90 @@ describe('api: template ref', () => {
     }
   })
 
+  it('should not register duplicate onScopeDispose callbacks for dynamic function refs', async () => {
+    const fn1 = vi.fn()
+    const fn2 = vi.fn()
+    const toggle = ref(true)
+    const t0 = template('<div></div>')
+
+    const { app } = define({
+      render() {
+        const n0 = t0()
+        let r0: any
+        renderEffect(() => {
+          r0 = createTemplateRefSetter()(
+            n0 as Element,
+            toggle.value ? fn1 : fn2,
+            r0,
+          )
+        })
+        return n0
+      },
+    }).render()
+
+    expect(fn1).toHaveBeenCalledTimes(1)
+    expect(fn2).toHaveBeenCalledTimes(0)
+    expect(app._instance!.scope.cleanups.length).toBe(1)
+
+    toggle.value = false
+    await nextTick()
+    expect(fn1).toHaveBeenCalledTimes(1)
+    expect(fn2).toHaveBeenCalledTimes(1)
+    expect(app._instance!.scope.cleanups.length).toBe(1)
+
+    toggle.value = true
+    await nextTick()
+    expect(fn1).toHaveBeenCalledTimes(2)
+    expect(fn2).toHaveBeenCalledTimes(1)
+    expect(app._instance!.scope.cleanups.length).toBe(1)
+
+    app.unmount()
+    await nextTick()
+    // expected fn1 to be called again during scope dispose
+    expect(fn1).toHaveBeenCalledTimes(3)
+    expect(fn2).toHaveBeenCalledTimes(1)
+  })
+
+  it('should not register duplicate onScopeDispose callbacks for dynamic string refs', async () => {
+    const el1 = ref(null)
+    const el2 = ref(null)
+    const toggle = ref(true)
+    const t0 = template('<div></div>')
+
+    const { app, host } = define({
+      setup() {
+        return { ref1: el1, ref2: el2 }
+      },
+      render() {
+        const n0 = t0()
+        let r0: any
+        renderEffect(() => {
+          r0 = createTemplateRefSetter()(
+            n0 as Element,
+            toggle.value ? 'ref1' : 'ref2',
+            r0,
+          )
+        })
+        return n0
+      },
+    }).render()
+
+    expect(el1.value).toBe(host.children[0])
+    expect(el2.value).toBe(null)
+    expect(app._instance!.scope.cleanups.length).toBe(1)
+
+    toggle.value = false
+    await nextTick()
+    expect(el1.value).toBe(null)
+    expect(el2.value).toBe(host.children[0])
+    expect(app._instance!.scope.cleanups.length).toBe(1)
+
+    app.unmount()
+    await nextTick()
+    expect(el1.value).toBe(null)
+    expect(el2.value).toBe(null)
+  })
+
   // TODO: can not reproduce in Vapor
   // // #2078
   // test('handling multiple merged refs', async () => {

packages/runtime-vapor/src/apiTemplateRef.ts

@@ -16,7 +16,8 @@ import {
 } from '@vue/runtime-dom'
 import {
   EMPTY_OBJ,
-  hasOwn,
+  NO,
+  NOOP,
   isArray,
   isFunction,
   isString,
@@ -38,6 +39,20 @@ export type setRefFn = (
   refKey?: string,
 ) => NodeRef | undefined
 
+const refCleanups = new WeakMap<RefEl, { fn: () => void }>()
+
+function ensureCleanup(el: RefEl): { fn: () => void } {
+  let cleanupRef = refCleanups.get(el)
+  if (!cleanupRef) {
+    refCleanups.set(el, (cleanupRef = { fn: NOOP }))
+    onScopeDispose(() => {
+      cleanupRef!.fn()
+      refCleanups.delete(el)
+    })
+  }
+  return cleanupRef
+}
+
 export function createTemplateRefSetter(): setRefFn {
   const instance = currentInstance as VaporComponentInstance
   return (...args) => setRef(instance, ...args)
@@ -80,12 +95,12 @@ export function setRef(
   const refs =
     instance.refs === EMPTY_OBJ ? (instance.refs = {}) : instance.refs
 
-  const canSetSetupRef = createCanSetSetupRefChecker(setupState)
+  const canSetSetupRef = __DEV__ ? createCanSetSetupRefChecker(setupState) : NO
   // dynamic ref changed. unset old ref
   if (oldRef != null && oldRef !== ref) {
     if (isString(oldRef)) {
       refs[oldRef] = null
-      if (__DEV__ && hasOwn(setupState, oldRef)) {
+      if (__DEV__ && canSetSetupRef(oldRef)) {
         setupState[oldRef] = null
       }
     } else if (isRef(oldRef)) {
@@ -94,16 +109,15 @@ export function setRef(
   }
 
   if (isFunction(ref)) {
-    const invokeRefSetter = (value?: Element | Record<string, any>) => {
+    const invokeRefSetter = (value?: Element | Record<string, any> | null) => {
       callWithErrorHandling(ref, currentInstance, ErrorCodes.FUNCTION_REF, [
         value,
         refs,
       ])
     }
 
     invokeRefSetter(refValue)
-    // TODO this gets called repeatedly in renderEffect when it's dynamic ref?
-    onScopeDispose(() => invokeRefSetter())
+    ensureCleanup(el).fn = () => invokeRefSetter(null)
   } else {
     const _isString = isString(ref)
     const _isRef = isRef(ref)
@@ -150,8 +164,7 @@ export function setRef(
       }
       queuePostFlushCb(doSet, -1)
 
-      // TODO this gets called repeatedly in renderEffect when it's dynamic ref?
-      onScopeDispose(() => {
+      ensureCleanup(el).fn = () => {
         queuePostFlushCb(() => {
           if (isArray(existing)) {
             remove(existing, refValue)
@@ -165,7 +178,7 @@ export function setRef(
             if (refKey) refs[refKey] = null
           }
         })
-      })
+      }
     } else if (__DEV__) {
       warn('Invalid template ref type:', ref, `(${typeof ref})`)
     }