fix(ssr): should set ref on hydration

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

@@ -124,6 +124,15 @@ describe('SSR hydration', () => {
     expect(vnode.el.innerHTML).toBe(`<span>bar</span><span class="bar"></span>`)
   })
 
+  test('element with ref', () => {
+    const el = ref()
+    const { vnode, container } = mountWithHydration('<div></div>', () =>
+      h('div', { ref: el })
+    )
+    expect(vnode.el).toBe(container.firstChild)
+    expect(el.value).toBe(vnode.el)
+  })
+
   test('Fragment', async () => {
     const msg = ref('foo')
     const fn = jest.fn()

packages/runtime-core/src/hydration.ts

@@ -12,7 +12,7 @@ import { ComponentOptions, ComponentInternalInstance } from './component'
 import { invokeDirectiveHook } from './directives'
 import { warn } from './warning'
 import { PatchFlags, ShapeFlags, isReservedProp, isOn } from '@vue/shared'
-import { RendererInternals, invokeVNodeHook } from './renderer'
+import { RendererInternals, invokeVNodeHook, setRef } from './renderer'
 import {
   SuspenseImpl,
   SuspenseBoundary,
@@ -88,15 +88,16 @@ export function createHydrationFunctions(
         isFragmentStart
       )
 
-    const { type, shapeFlag } = vnode
+    const { type, ref, shapeFlag } = vnode
     const domType = node.nodeType
     vnode.el = node
 
+    let nextNode: Node | null = null
     switch (type) {
       case Text:
         if (domType !== DOMNodeTypes.TEXT) {
-          return onMismatch()
+          nextNode = onMismatch()
-        }
+        } else {
           if ((node as Text).data !== vnode.children) {
             hasMismatch = true
             __DEV__ &&
@@ -107,55 +108,65 @@ export function createHydrationFunctions(
               )
             ;(node as Text).data = vnode.children as string
           }
-        return nextSibling(node)
+          nextNode = nextSibling(node)
+        }
+        break
       case Comment:
         if (domType !== DOMNodeTypes.COMMENT || isFragmentStart) {
-          return onMismatch()
+          nextNode = onMismatch()
+        } else {
+          nextNode = nextSibling(node)
         }
-        return nextSibling(node)
+        break
       case Static:
         if (domType !== DOMNodeTypes.ELEMENT) {
-          return onMismatch()
+          nextNode = onMismatch()
-        }
+        } else {
           // determine anchor, adopt content
-        let cur = node
+          nextNode = node
           // if the static vnode has its content stripped during build,
           // adopt it from the server-rendered HTML.
           const needToAdoptContent = !(vnode.children as string).length
           for (let i = 0; i < vnode.staticCount; i++) {
-          if (needToAdoptContent) vnode.children += (cur as Element).outerHTML
+            if (needToAdoptContent)
+              vnode.children += (nextNode as Element).outerHTML
             if (i === vnode.staticCount - 1) {
-            vnode.anchor = cur
+              vnode.anchor = nextNode
             }
-          cur = nextSibling(cur)!
+            nextNode = nextSibling(nextNode)!
           }
-        return cur
+          return nextNode
+        }
+        break
       case Fragment:
         if (!isFragmentStart) {
-          return onMismatch()
+          nextNode = onMismatch()
-        }
+        } else {
-        return hydrateFragment(
+          nextNode = hydrateFragment(
             node as Comment,
             vnode,
             parentComponent,
             parentSuspense,
             optimized
           )
+        }
+        break
       default:
         if (shapeFlag & ShapeFlags.ELEMENT) {
           if (
             domType !== DOMNodeTypes.ELEMENT ||
             vnode.type !== (node as Element).tagName.toLowerCase()
           ) {
-            return onMismatch()
+            nextNode = onMismatch()
-          }
+          } else {
-          return hydrateElement(
+            nextNode = hydrateElement(
               node as Element,
               vnode,
               parentComponent,
               parentSuspense,
               optimized
             )
+          }
         } else if (shapeFlag & ShapeFlags.COMPONENT) {
           // when setting up the render effect, if the initial vnode already
           // has .el set, the component will perform hydration instead of mount
@@ -182,14 +193,14 @@ export function createHydrationFunctions(
           // component may be async, so in the case of fragments we cannot rely
           // on component's rendered output to determine the end of the fragment
           // instead, we do a lookahead to find the end anchor node.
-          return isFragmentStart
+          nextNode = isFragmentStart
             ? locateClosingAsyncAnchor(node)
             : nextSibling(node)
         } else if (shapeFlag & ShapeFlags.TELEPORT) {
           if (domType !== DOMNodeTypes.COMMENT) {
-            return onMismatch()
+            nextNode = onMismatch()
-          }
+          } else {
-          return (vnode.type as typeof TeleportImpl).hydrate(
+            nextNode = (vnode.type as typeof TeleportImpl).hydrate(
               node,
               vnode,
               parentComponent,
@@ -198,8 +209,9 @@ export function createHydrationFunctions(
               rendererInternals,
               hydrateChildren
             )
+          }
         } else if (__FEATURE_SUSPENSE__ && shapeFlag & ShapeFlags.SUSPENSE) {
-          return (vnode.type as typeof SuspenseImpl).hydrate(
+          nextNode = (vnode.type as typeof SuspenseImpl).hydrate(
             node,
             vnode,
             parentComponent,
@@ -212,8 +224,13 @@ export function createHydrationFunctions(
         } else if (__DEV__) {
           warn('Invalid HostVNode type:', type, `(${typeof type})`)
         }
-        return null
     }
+
+    if (ref != null && parentComponent) {
+      setRef(ref, null, parentComponent, vnode)
+    }
+
+    return nextNode
   }
 
   const hydrateElement = (
@@ -386,7 +403,7 @@ export function createHydrationFunctions(
     parentComponent: ComponentInternalInstance | null,
     parentSuspense: SuspenseBoundary | null,
     isFragment: boolean
-  ) => {
+  ): Node | null => {
     hasMismatch = true
     __DEV__ &&
       warn(

packages/runtime-core/src/renderer.ts

@@ -47,7 +47,6 @@ import { effect, stop, ReactiveEffectOptions, isRef } from '@vue/reactivity'
 import { updateProps } from './componentProps'
 import { updateSlots } from './componentSlots'
 import { pushWarningContext, popWarningContext, warn } from './warning'
-import { ComponentPublicInstance } from './componentProxy'
 import { createAppAPI, CreateAppFunction } from './apiCreateApp'
 import {
   SuspenseBoundary,
@@ -271,6 +270,55 @@ export const queuePostRenderEffect = __FEATURE_SUSPENSE__
   ? queueEffectWithSuspense
   : queuePostFlushCb
 
+export const setRef = (
+  rawRef: VNodeNormalizedRef,
+  oldRawRef: VNodeNormalizedRef | null,
+  parent: ComponentInternalInstance,
+  vnode: VNode | null
+) => {
+  const value = vnode
+    ? vnode.shapeFlag & ShapeFlags.STATEFUL_COMPONENT
+      ? vnode.component!.proxy
+      : vnode.el
+    : null
+  const [owner, ref] = rawRef
+  if (__DEV__ && !owner) {
+    warn(
+      `Missing ref owner context. ref cannot be used on hoisted vnodes. ` +
+        `A vnode with ref must be created inside the render function.`
+    )
+    return
+  }
+  const oldRef = oldRawRef && oldRawRef[1]
+  const refs = owner.refs === EMPTY_OBJ ? (owner.refs = {}) : owner.refs
+  const setupState = owner.setupState
+
+  // unset old ref
+  if (oldRef != null && oldRef !== ref) {
+    if (isString(oldRef)) {
+      refs[oldRef] = null
+      if (hasOwn(setupState, oldRef)) {
+        setupState[oldRef] = null
+      }
+    } else if (isRef(oldRef)) {
+      oldRef.value = null
+    }
+  }
+
+  if (isString(ref)) {
+    refs[ref] = value
+    if (hasOwn(setupState, ref)) {
+      setupState[ref] = value
+    }
+  } else if (isRef(ref)) {
+    ref.value = value
+  } else if (isFunction(ref)) {
+    callWithErrorHandling(ref, parent, ErrorCodes.FUNCTION_REF, [value, refs])
+  } else if (__DEV__) {
+    warn('Invalid template ref type:', value, `(${typeof value})`)
+  }
+}
+
 /**
  * The createRenderer function accepts two generic arguments:
  * HostNode and HostElement, corresponding to Node and Element types in the
@@ -440,9 +488,7 @@ function baseCreateRenderer(
 
     // set ref
     if (ref != null && parentComponent) {
-      const refValue =
+      setRef(ref, n1 && n1.ref, parentComponent, n2)
-        shapeFlag & ShapeFlags.STATEFUL_COMPONENT ? n2.component!.proxy : n2.el
-      setRef(ref, n1 && n1.ref, parentComponent, refValue)
     }
   }
 
@@ -1984,50 +2030,6 @@ function baseCreateRenderer(
     return hostNextSibling((vnode.anchor || vnode.el)!)
   }
 
-  const setRef = (
-    rawRef: VNodeNormalizedRef,
-    oldRawRef: VNodeNormalizedRef | null,
-    parent: ComponentInternalInstance,
-    value: RendererNode | ComponentPublicInstance | null
-  ) => {
-    const [owner, ref] = rawRef
-    if (__DEV__ && !owner) {
-      warn(
-        `Missing ref owner context. ref cannot be used on hoisted vnodes. ` +
-          `A vnode with ref must be created inside the render function.`
-      )
-      return
-    }
-    const oldRef = oldRawRef && oldRawRef[1]
-    const refs = owner.refs === EMPTY_OBJ ? (owner.refs = {}) : owner.refs
-    const setupState = owner.setupState
-
-    // unset old ref
-    if (oldRef != null && oldRef !== ref) {
-      if (isString(oldRef)) {
-        refs[oldRef] = null
-        if (hasOwn(setupState, oldRef)) {
-          setupState[oldRef] = null
-        }
-      } else if (isRef(oldRef)) {
-        oldRef.value = null
-      }
-    }
-
-    if (isString(ref)) {
-      refs[ref] = value
-      if (hasOwn(setupState, ref)) {
-        setupState[ref] = value
-      }
-    } else if (isRef(ref)) {
-      ref.value = value
-    } else if (isFunction(ref)) {
-      callWithErrorHandling(ref, parent, ErrorCodes.FUNCTION_REF, [value, refs])
-    } else if (__DEV__) {
-      warn('Invalid template ref type:', value, `(${typeof value})`)
-    }
-  }
-
   /**
    * #1156
    * When a component is HMR-enabled, we need to make sure that all static nodes