refactor(runtime-vapor): improve template ref handling for DynamicFragment and async components (#14472)

packages/runtime-vapor/__tests__/apiDefineAsyncComponent.spec.ts

@@ -690,6 +690,120 @@ describe('api: defineAsyncComponent', () => {
     expect(fooRef.value.id).toBe('foo')
   })
 
+  test('template ref forwarding should not keep stale ref callbacks before resolve', async () => {
+    let resolve: (comp: VaporComponent) => void
+    const Foo = defineVaporAsyncComponent(
+      () =>
+        new Promise(r => {
+          resolve = r as any
+        }),
+    )
+
+    const refA = ref<any>(null)
+    const refB = ref<any>(null)
+    const useA = ref(true)
+    const root = document.createElement('div')
+
+    const { mount } = define({
+      setup() {
+        return { refA, refB, useA }
+      },
+      render() {
+        const setTemplateRef = createTemplateRefSetter()
+        const n0 = createComponent(Foo, null, null, true)
+        renderEffect(() => {
+          setTemplateRef(n0, useA.value ? 'refA' : 'refB')
+        })
+        return n0
+      },
+    }).create()
+
+    mount(root)
+    expect(root.innerHTML).toBe('<!--async component-->')
+    expect(refA.value).toBe(null)
+    expect(refB.value).toBe(null)
+
+    useA.value = false
+    await nextTick()
+    useA.value = true
+    await nextTick()
+
+    resolve!({
+      setup: (props, { expose }) => {
+        expose({
+          id: 'foo',
+        })
+        return template('resolved')()
+      },
+    })
+    await timeout()
+
+    expect(root.innerHTML).toBe('resolved<!--async component-->')
+    expect(refA.value.id).toBe('foo')
+    expect(refB.value).toBe(null)
+  })
+
+  test('template ref forwarding should not keep stale ref callbacks after resolve', async () => {
+    let resolve: (comp: VaporComponent) => void
+    const Foo = defineVaporAsyncComponent(
+      () =>
+        new Promise(r => {
+          resolve = r as any
+        }),
+    )
+
+    const refA = ref<any>(null)
+    const refB = ref<any>(null)
+    const useA = ref(true)
+    const root = document.createElement('div')
+    let asyncWrapper: any
+
+    const { mount } = define({
+      setup() {
+        return { refA, refB, useA }
+      },
+      render() {
+        const setTemplateRef = createTemplateRefSetter()
+        const n0 = (asyncWrapper = createComponent(Foo, null, null, true))
+        renderEffect(() => {
+          setTemplateRef(n0, useA.value ? 'refA' : 'refB')
+        })
+        return n0
+      },
+    }).create()
+
+    mount(root)
+    expect(root.innerHTML).toBe('<!--async component-->')
+    expect(refA.value).toBe(null)
+    expect(refB.value).toBe(null)
+
+    resolve!({
+      setup: (props, { expose }) => {
+        expose({
+          id: 'foo',
+        })
+        return template('resolved')()
+      },
+    })
+    await timeout()
+
+    expect(root.innerHTML).toBe('resolved<!--async component-->')
+    expect(refA.value.id).toBe('foo')
+    expect(refB.value).toBe(null)
+
+    useA.value = false
+    await nextTick()
+    expect(refA.value).toBe(null)
+    expect(refB.value.id).toBe('foo')
+
+    const onUpdated = asyncWrapper.block.onUpdated
+    if (onUpdated) onUpdated.forEach((hook: any) => hook())
+    await nextTick()
+
+    expect(refA.value).toBe(null)
+    expect(refB.value.id).toBe('foo')
+  })
+
   test('the forwarded template ref should always exist when doing multi patching', async () => {
     let resolve: (comp: VaporComponent) => void
     const Foo = defineVaporAsyncComponent(

packages/runtime-vapor/__tests__/apiExpose.spec.ts

@@ -1,6 +1,6 @@
 import { ref, shallowRef } from '@vue/reactivity'
 import { type VaporComponentInstance, createComponent } from '../src/component'
-import { setRef } from '../src/apiTemplateRef'
+import { createTemplateRefSetter } from '../src/apiTemplateRef'
 import { makeRender } from './_utils'
 import { currentInstance } from '@vue/runtime-dom'
 import { defineVaporComponent } from '../src/apiDefineComponent'
@@ -23,7 +23,8 @@ describe('api: expose', () => {
     define({
       setup: () => {
         const n0 = (i = createComponent(Child))
-        setRef(currentInstance as VaporComponentInstance, n0, childRef)
+        const setRef = createTemplateRefSetter()
+        setRef(n0, childRef)
         return n0
       },
     }).render()
@@ -46,7 +47,8 @@ describe('api: expose', () => {
     define({
       setup: () => {
         const n0 = createComponent(Child)
-        setRef(currentInstance as VaporComponentInstance, n0, childRef)
+        const setRef = createTemplateRefSetter()
+        setRef(n0, childRef)
         return n0
       },
     }).render()

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

@@ -737,6 +737,230 @@ describe('api: template ref', () => {
     expect(msg.value).toBe('two')
   })
 
+  test('useTemplateRef should update when switching dynamic components', async () => {
+    const One = defineVaporComponent({
+      setup(_, { expose }) {
+        expose({ name: 'one' })
+        return template('<div>one</div>')()
+      },
+    })
+
+    const Two = defineVaporComponent({
+      setup(_, { expose }) {
+        expose({ name: 'two' })
+        return template('<div>two</div>')()
+      },
+    })
+
+    const views: VaporComponent[] = [One, Two]
+    const view = ref(0)
+    const refNames: string[] = []
+    let compRef: ShallowRef
+
+    const { html } = define({
+      setup() {
+        compRef = useTemplateRef('compRef')
+        watchEffect(() => {
+          const value = compRef.value as { name?: string } | null
+          if (value?.name) {
+            refNames.push(value.name)
+          }
+        })
+
+        const setRef = createTemplateRefSetter()
+        const n0 = createDynamicComponent(() => views[view.value]) as any
+        setRef(n0, compRef, false, 'compRef')
+        return n0
+      },
+    }).render()
+
+    await nextTick()
+    const one = compRef!.value
+    expect(one).toMatchObject({ name: 'one' })
+    expect(html()).toBe('<div>one</div><!--dynamic-component-->')
+
+    view.value = 1
+    await nextTick()
+    const two = compRef!.value
+    expect(two).toMatchObject({ name: 'two' })
+    expect(two).not.toBe(one)
+    expect(html()).toBe('<div>two</div><!--dynamic-component-->')
+
+    view.value = 0
+    await nextTick()
+    expect(compRef!.value).toMatchObject({ name: 'one' })
+    expect(compRef!.value).not.toBe(two)
+    expect(refNames).toContain('two')
+  })
+
+  test('dynamic component should not register duplicate onUpdated handlers for refs', async () => {
+    const One = defineVaporComponent({
+      setup() {
+        return template('<div>one</div>')()
+      },
+    })
+
+    const Two = defineVaporComponent({
+      setup() {
+        return template('<div>two</div>')()
+      },
+    })
+
+    const views: VaporComponent[] = [One, Two]
+    const view = ref(0)
+    const useA = ref(true)
+    const refA = ref<any>(null)
+    const refB = ref<any>(null)
+    let frag: any
+
+    define({
+      setup() {
+        const setRef = createTemplateRefSetter()
+        frag = createDynamicComponent(() => views[view.value]) as any
+        renderEffect(() => {
+          setRef(frag, useA.value ? refA : refB)
+        })
+        return frag
+      },
+    }).render()
+
+    expect(frag.onUpdated.length).toBe(1)
+
+    useA.value = false
+    await nextTick()
+    expect(frag.onUpdated.length).toBe(1)
+
+    view.value = 1
+    await nextTick()
+    expect(frag.onUpdated.length).toBe(1)
+  })
+
+  test('dynamic component function ref should cleanup old branch with null', async () => {
+    const One = defineVaporComponent({
+      setup(_, { expose }) {
+        expose({ name: 'one' })
+        return template('<div>one</div>')()
+      },
+    })
+
+    const Two = defineVaporComponent({
+      setup(_, { expose }) {
+        expose({ name: 'two' })
+        return template('<div>two</div>')()
+      },
+    })
+
+    const views: VaporComponent[] = [One, Two]
+    const view = ref(0)
+    const fnRef = vi.fn()
+
+    define({
+      setup() {
+        const setRef = createTemplateRefSetter()
+        const n0 = createDynamicComponent(() => views[view.value]) as any
+        setRef(n0, fnRef as any)
+        return n0
+      },
+    }).render()
+
+    expect(fnRef).toHaveBeenCalledTimes(1)
+    expect(fnRef.mock.calls[0][0]).toMatchObject({ name: 'one' })
+
+    view.value = 1
+    await nextTick()
+
+    const callArgs = fnRef.mock.calls.map(args => args[0])
+    expect(callArgs).toContain(null)
+    expect(fnRef.mock.calls[fnRef.mock.calls.length - 1][0]).toMatchObject({
+      name: 'two',
+    })
+  })
+
+  test('dynamic component function ref should cleanup previous callback when ref changes', async () => {
+    const One = defineVaporComponent({
+      setup(_, { expose }) {
+        expose({ name: 'one' })
+        return template('<div>one</div>')()
+      },
+    })
+
+    const useA = ref(true)
+    const fnA = vi.fn()
+    const fnB = vi.fn()
+
+    define({
+      setup() {
+        const setRef = createTemplateRefSetter()
+        const n0 = createDynamicComponent(() => One) as any
+        renderEffect(() => {
+          setRef(n0, useA.value ? (fnA as any) : (fnB as any))
+        })
+        return n0
+      },
+    }).render()
+
+    expect(fnA.mock.calls[fnA.mock.calls.length - 1][0]).toMatchObject({
+      name: 'one',
+    })
+    expect(fnB).toHaveBeenCalledTimes(0)
+
+    useA.value = false
+    await nextTick()
+
+    const fnAArgs = fnA.mock.calls.map(args => args[0])
+    expect(fnAArgs).toContain(null)
+    expect(fnB.mock.calls[fnB.mock.calls.length - 1][0]).toMatchObject({
+      name: 'one',
+    })
+  })
+
+  test('dynamic component ref_for should keep sibling refs when one branch updates', async () => {
+    const One = defineVaporComponent({
+      setup(_, { expose }) {
+        expose({ name: 'one' })
+        return template('<div>one</div>')()
+      },
+    })
+
+    const Two = defineVaporComponent({
+      setup(_, { expose }) {
+        expose({ name: 'two' })
+        return template('<div>two</div>')()
+      },
+    })
+
+    const views: VaporComponent[] = [One, Two]
+    const view = ref(0)
+    const listRef = ref<any[]>([])
+
+    define({
+      setup() {
+        return { listRef }
+      },
+      render() {
+        const n0 = template('<div></div>')() as Element
+        const setRef = createTemplateRefSetter()
+        const n1 = createDynamicComponent(() => views[view.value]) as any
+        const n2 = createDynamicComponent(() => One) as any
+        setRef(n1, listRef as any, true)
+        setRef(n2, listRef as any, true)
+        insert(n1, n0 as ParentNode)
+        insert(n2, n0 as ParentNode)
+        return n0
+      },
+    }).render()
+
+    await nextTick()
+    expect(listRef.value).toHaveLength(2)
+    expect(listRef.value.filter(i => i?.name === 'one')).toHaveLength(2)
+
+    view.value = 1
+    await nextTick()
+    expect(listRef.value).toHaveLength(2)
+    expect(listRef.value.filter(i => i?.name === 'one')).toHaveLength(1)
+    expect(listRef.value.some(i => i?.name === 'two')).toBe(true)
+  })
+
   test('should not attempt to set when variable name is same as key', () => {
     let tRef: ShallowRef
     const key = 'refKey'

packages/runtime-vapor/src/apiDefineAsyncComponent.ts

@@ -225,8 +225,5 @@ function createInnerComp(
     appContext,
   )
 
-  // set ref
-  frag && frag.setAsyncRef && frag.setAsyncRef(instance)
-
   return instance
 }

packages/runtime-vapor/src/apiTemplateRef.ts

@@ -25,13 +25,23 @@ import {
   isString,
   remove,
 } from '@vue/shared'
-import { DynamicFragment, isDynamicFragment, isFragment } from './fragment'
+import {
+  type DynamicFragment,
+  type VaporFragment,
+  isDynamicFragment,
+  isFragment,
+} from './fragment'
+import { isInteropEnabled } from './vdomInteropState'
 
 export type NodeRef =
   | string
   | Ref
   | ((ref: Element | VaporComponentInstance, refs: Record<string, any>) => void)
-export type RefEl = Element | VaporComponentInstance
+export type RefEl =
+  | Element
+  | VaporComponentInstance
+  | DynamicFragment
+  | VaporFragment
 
 export type setRefFn = (
   el: RefEl,
@@ -57,12 +67,25 @@ function ensureCleanup(el: RefEl): { fn: () => void } {
 export function createTemplateRefSetter(): setRefFn {
   const instance = currentInstance as VaporComponentInstance
   const oldRefMap = new WeakMap<RefEl, NodeRef | undefined>()
+  const setRefMap = new WeakMap<DynamicFragment, () => void>()
+
   return (el, ref, refFor, refKey) => {
-    if (isDynamicFragment(el)) {
-      ;(el.onUpdated || (el.onUpdated = [])).push(() => {
-        setRef(instance, el, ref, oldRefMap.get(el), refFor, refKey)
-      })
+    // Re-apply refs after DynamicFragment updates.
+    if (isDynamicFragment(el) || (isVaporComponent(el) && isAsyncWrapper(el))) {
+      const frag = isDynamicFragment(el)
+        ? (el as DynamicFragment)
+        : ((el as VaporComponentInstance).block as DynamicFragment)
+      const doSet = () =>
+        oldRefMap.set(
+          el,
+          setRef(instance, el, ref, oldRefMap.get(el), refFor, refKey),
+        )
+      const prevSet = setRefMap.get(frag)
+      if (prevSet && frag.onUpdated) remove(frag.onUpdated, prevSet)
+      ;(frag.onUpdated || (frag.onUpdated = [])).push(doSet)
+      setRefMap.set(frag, doSet)
     }
+
     const oldRef = setRef(instance, el, ref, oldRefMap.get(el), refFor, refKey)
     oldRefMap.set(el, oldRef)
     return oldRef
@@ -72,7 +95,7 @@ export function createTemplateRefSetter(): setRefFn {
 /**
  * Function for handling a template ref
  */
-export function setRef(
+function setRef(
   instance: VaporComponentInstance,
   el: RefEl,
   ref: NodeRef,
@@ -83,22 +106,17 @@ export function setRef(
   if (!instance || instance.isUnmounted) return
 
   // vdom interop
-  if (isFragment(el) && el.setRef) {
+  if (isInteropEnabled && isFragment(el) && el.setRef) {
     el.setRef(instance, ref, refFor, refKey)
     return
   }
 
   if (isVaporComponent(el) && isAsyncWrapper(el)) {
-    const frag = el.block as DynamicFragment
-    // async component not resolved yet, register ref setter
-    // it will be called when the async component is resolved
-    if (!el.type.__asyncResolved) {
-      frag.setAsyncRef = i => setRef(instance, i, ref, oldRef, refFor)
-      return
-    }
+    // unresolved: handled in DynamicFragment's updated hook
+    if (!el.type.__asyncResolved) return
 
-    // set ref to the inner component instead
-    el = frag.nodes as VaporComponentInstance
+    // resolved: set ref to the inner component
+    el = (el.block as DynamicFragment).nodes as VaporComponentInstance
   }
 
   const setupState: any = __DEV__ ? instance.setupState || {} : null
@@ -121,7 +139,7 @@ export function setRef(
   }
 
   // dynamic ref changed. unset old ref
-  if (oldRef != null && (oldRef !== ref || isDynamicFragment(el))) {
+  if (oldRef != null && oldRef !== ref) {
     if (isString(oldRef)) {
       refs[oldRef] = null
       if (__DEV__ && canSetSetupRef(oldRef)) {
@@ -129,12 +147,28 @@ export function setRef(
       }
     } else if (isRef(oldRef)) {
       if (canSetRef(oldRef)) oldRef.value = null
+    } else if (isFunction(oldRef) && isDynamicFragment(el)) {
+      callWithErrorHandling(oldRef, instance, ErrorCodes.FUNCTION_REF, [
+        null,
+        refs,
+      ])
+    }
+  } else if (oldRef != null && isDynamicFragment(el)) {
+    if (isFunction(oldRef)) {
+      callWithErrorHandling(oldRef, instance, ErrorCodes.FUNCTION_REF, [
+        null,
+        refs,
+      ])
+    } else if (refFor) {
+      // For dynamic ref-for branches, remove only this branch's previous value.
+      const cleanup = refCleanups.get(el)
+      if (cleanup) cleanup.fn()
     }
   }
 
   if (isFunction(ref)) {
     const invokeRefSetter = (value?: Element | Record<string, any> | null) => {
-      callWithErrorHandling(ref, currentInstance, ErrorCodes.FUNCTION_REF, [
+      callWithErrorHandling(ref, instance, ErrorCodes.FUNCTION_REF, [
         value,
         refs,
       ])
@@ -215,7 +249,8 @@ export function setRef(
 const getRefValue = (el: RefEl) => {
   if (isVaporComponent(el)) {
     return getExposed(el) || el
-  } else if (el instanceof DynamicFragment) {
+  } else if (isDynamicFragment(el)) {
+    if (isArray(el.nodes)) return null
     return getRefValue(el.nodes as RefEl)
   }
   return el

packages/runtime-vapor/src/fragment.ts

@@ -92,9 +92,6 @@ export class DynamicFragment extends VaporFragment {
   // fallthrough attrs
   attrs?: Record<string, any>
 
-  // set ref for async wrapper
-  setAsyncRef?: (instance: VaporComponentInstance) => void
-
   keepAliveCtx: VaporKeepAliveContext | null
 
   slotOwner: VaporComponentInstance | null