fix(async-component): forward refs on async component wrapper

fix #2671

packages/runtime-core/__tests__/apiAsyncComponent.spec.ts

@@ -652,4 +652,49 @@ describe('api: defineAsyncComponent', () => {
     expect(loaderCallCount).toBe(2)
     expect(serializeInner(root)).toBe('<!---->')
   })
+
+  test('template ref forwarding', async () => {
+    let resolve: (comp: Component) => void
+    const Foo = defineAsyncComponent(
+      () =>
+        new Promise(r => {
+          resolve = r as any
+        })
+    )
+
+    const fooRef = ref()
+    const toggle = ref(true)
+    const root = nodeOps.createElement('div')
+    createApp({
+      render: () => (toggle.value ? h(Foo, { ref: fooRef }) : null)
+    }).mount(root)
+
+    expect(serializeInner(root)).toBe('<!---->')
+    expect(fooRef.value).toBe(null)
+
+    resolve!({
+      data() {
+        return {
+          id: 'foo'
+        }
+      },
+      render: () => 'resolved'
+    })
+    // first time resolve, wait for macro task since there are multiple
+    // microtasks / .then() calls
+    await timeout()
+    expect(serializeInner(root)).toBe('resolved')
+    expect(fooRef.value.id).toBe('foo')
+
+    toggle.value = false
+    await nextTick()
+    expect(serializeInner(root)).toBe('<!---->')
+    expect(fooRef.value).toBe(null)
+
+    // already resolved component should update on nextTick
+    toggle.value = true
+    await nextTick()
+    expect(serializeInner(root)).toBe('resolved')
+    expect(fooRef.value.id).toBe('foo')
+  })
 })

packages/runtime-core/src/apiAsyncComponent.ts

@@ -3,11 +3,12 @@ import {
   ConcreteComponent,
   currentInstance,
   ComponentInternalInstance,
-  isInSSRComponentSetup
+  isInSSRComponentSetup,
+  ComponentOptions
 } from './component'
 import { isFunction, isObject } from '@vue/shared'
 import { ComponentPublicInstance } from './componentPublicInstance'
-import { createVNode } from './vnode'
+import { createVNode, VNode } from './vnode'
 import { defineComponent } from './apiDefineComponent'
 import { warn } from './warning'
 import { ref } from '@vue/reactivity'
@@ -34,6 +35,9 @@ export interface AsyncComponentOptions<T = any> {
   ) => any
 }
 
+export const isAsyncWrapper = (i: ComponentInternalInstance | VNode): boolean =>
+  !!(i.type as ComponentOptions).__asyncLoader
+
 export function defineAsyncComponent<
   T extends Component = { new (): ComponentPublicInstance }
 >(source: AsyncComponentLoader<T> | AsyncComponentOptions<T>): T {
@@ -193,7 +197,10 @@ export function defineAsyncComponent<
 
 function createInnerComp(
   comp: ConcreteComponent,
-  { vnode: { props, children } }: ComponentInternalInstance
+  { vnode: { ref, props, children } }: ComponentInternalInstance
 ) {
-  return createVNode(comp, props, children)
+  const vnode = createVNode(comp, props, children)
+  // ensure inner component inherits the async wrapper's ref owner
+  vnode.ref = ref
+  return vnode
 }

packages/runtime-core/src/renderer.ts

@@ -74,6 +74,7 @@ import { startMeasure, endMeasure } from './profiling'
 import { ComponentPublicInstance } from './componentPublicInstance'
 import { devtoolsComponentRemoved, devtoolsComponentUpdated } from './devtools'
 import { initFeatureFlags } from './featureFlags'
+import { isAsyncWrapper } from './apiAsyncComponent'
 
 export interface Renderer<HostElement = RendererElement> {
   render: RootRenderFunction<HostElement>
@@ -289,7 +290,6 @@ export const queuePostRenderEffect = __FEATURE_SUSPENSE__
 export const setRef = (
   rawRef: VNodeNormalizedRef,
   oldRawRef: VNodeNormalizedRef | null,
-  parentComponent: ComponentInternalInstance,
   parentSuspense: SuspenseBoundary | null,
   vnode: VNode | null
 ) => {
@@ -298,7 +298,6 @@ export const setRef = (
       setRef(
         r,
         oldRawRef && (isArray(oldRawRef) ? oldRawRef[i] : oldRawRef),
-        parentComponent,
         parentSuspense,
         vnode
       )
@@ -307,7 +306,7 @@ export const setRef = (
   }
 
   let value: ComponentPublicInstance | RendererNode | Record<string, any> | null
-  if (!vnode) {
+  if (!vnode || isAsyncWrapper(vnode)) {
     value = null
   } else {
     if (vnode.shapeFlag & ShapeFlags.STATEFUL_COMPONENT) {
@@ -368,10 +367,7 @@ export const setRef = (
       doSet()
     }
   } else if (isFunction(ref)) {
-    callWithErrorHandling(ref, parentComponent, ErrorCodes.FUNCTION_REF, [
-      value,
-      refs
-    ])
+    callWithErrorHandling(ref, owner, ErrorCodes.FUNCTION_REF, [value, refs])
   } else if (__DEV__) {
     warn('Invalid template ref type:', value, `(${typeof value})`)
   }
@@ -552,7 +548,7 @@ function baseCreateRenderer(
 
     // set ref
     if (ref != null && parentComponent) {
-      setRef(ref, n1 && n1.ref, parentComponent, parentSuspense, n2)
+      setRef(ref, n1 && n1.ref, parentSuspense, n2)
     }
   }
 
@@ -1983,8 +1979,8 @@ function baseCreateRenderer(
       dirs
     } = vnode
     // unset ref
-    if (ref != null && parentComponent) {
-      setRef(ref, null, parentComponent, parentSuspense, null)
+    if (ref != null) {
+      setRef(ref, null, parentSuspense, null)
     }
 
     if (shapeFlag & ShapeFlags.COMPONENT_SHOULD_KEEP_ALIVE) {

packages/runtime-core/src/vnode.ts

@@ -21,7 +21,7 @@ import {
   isClassComponent
 } from './component'
 import { RawSlots } from './componentSlots'
-import { isProxy, Ref, toRaw, ReactiveFlags } from '@vue/reactivity'
+import { isProxy, Ref, toRaw, ReactiveFlags, isRef } from '@vue/reactivity'
 import { AppContext } from './apiCreateApp'
 import {
   SuspenseImpl,
@@ -304,9 +304,9 @@ const normalizeKey = ({ key }: VNodeProps): VNode['key'] =>
 
 const normalizeRef = ({ ref }: VNodeProps): VNodeNormalizedRefAtom | null => {
   return (ref != null
-    ? isArray(ref)
-      ? ref
-      : { i: currentRenderingInstance, r: ref }
+    ? isString(ref) || isRef(ref) || isFunction(ref)
+      ? { i: currentRenderingInstance, r: ref }
+      : ref
     : null) as any
 }