[WIP] Support for ref callback (#4807)

* ✨ Tests for ref callback

* ✨ Support for ref callback

* Add test of inline ref callback

* adjust ref implementation strategy

* fix patch ref registration

* fix tests

* fix flow

* fix test for phantomjs

flow/component.js

@@ -2,6 +2,10 @@ import type { Config } from '../src/core/config'
 import type VNode from '../src/core/vdom/vnode'
 import type Watcher from '../src/core/observer/watcher'
 
+declare type Refs = {
+  [key: string]: Component | Element | Array<Component | Element> | void;
+};
+
 declare interface Component {
   // constructor information
   static cid: number;
@@ -23,7 +27,7 @@ declare interface Component {
   $parent: Component | void;
   $root: Component;
   $children: Array<Component>;
-  $refs: { [key: string]: Component | Element | Array<Component | Element> | void };
+  $refs: Refs;
   $slots: { [key: string]: Array<VNode> };
   $scopedSlots: { [key: string]: () => VNodeChildren };
   $vnode: VNode;

src/core/instance/lifecycle.js

@@ -1,6 +1,7 @@
 /* @flow */
 
 import Watcher from '../observer/watcher'
+import { resetRefs } from '../vdom/modules/ref'
 import { createEmptyVNode } from '../vdom/vnode'
 import { observerState } from '../observer/index'
 import { updateComponentListeners } from './events'
@@ -83,6 +84,7 @@ export function lifecycleMixin (Vue: Class<Component>) {
     const prevVnode = vm._vnode
     const prevActiveInstance = activeInstance
     activeInstance = vm
+    vm.$refs = resetRefs(vm.$refs)
     vm._vnode = vnode
     // Vue.prototype.__patch__ is injected in entry points
     // based on the rendering backend used.

src/core/vdom/modules/ref.js

@@ -1,44 +1,40 @@
 /* @flow */
 
-import { remove } from 'shared/util'
-
 export default {
-  create (_: any, vnode: VNodeWithData) {
-    registerRef(vnode)
-  },
-  update (oldVnode: VNodeWithData, vnode: VNodeWithData) {
-    if (oldVnode.data.ref !== vnode.data.ref) {
-      registerRef(oldVnode, true)
-      registerRef(vnode)
-    }
-  },
-  destroy (vnode: VNodeWithData) {
-    registerRef(vnode, true)
-  }
+  create: registerRef,
+  update: registerRef
 }
 
-export function registerRef (vnode: VNodeWithData, isRemoval: ?boolean) {
+export function registerRef (_: any, vnode: VNodeWithData) {
   const key = vnode.data.ref
   if (!key) return
 
-  const vm = vnode.context
   const ref = vnode.componentInstance || vnode.elm
-  const refs = vm.$refs
-  if (isRemoval) {
-    if (Array.isArray(refs[key])) {
-      remove(refs[key], ref)
-    } else if (refs[key] === ref) {
-      refs[key] = undefined
-    }
-  } else {
-    if (vnode.data.refInFor) {
-      if (Array.isArray(refs[key]) && refs[key].indexOf(ref) < 0) {
-        refs[key].push(ref)
-      } else {
-        refs[key] = [ref]
+  const refs = vnode.context.$refs
+
+  if (typeof key === 'function') {
+    key(ref)
+  } else if (vnode.data.refInFor) {
+    const refArray = refs[key]
+    if (Array.isArray(refArray)) {
+      if (refArray.indexOf(ref) < 0) {
+        refArray.push(ref)
       }
     } else {
-      refs[key] = ref
+      refs[key] = [ref]
+    }
+  } else {
+    refs[key] = ref
+  }
+}
+
+export function resetRefs (refs: Refs): Refs {
+  const res = {}
+  // keep existing v-for ref arrays even if empty
+  for (const key in refs) {
+    if (Array.isArray(refs[key])) {
+      res[key] = []
     }
   }
+  return res
 }

src/core/vdom/patch.js

@@ -189,7 +189,7 @@ export function createPatchFunction (backend) {
     } else {
       // empty component root.
       // skip all element-related modules except for ref (#3455)
-      registerRef(vnode)
+      registerRef(null, vnode)
       // make sure to invoke the insert hook
       insertedVnodeQueue.push(vnode)
     }

test/unit/features/ref.spec.js

@@ -157,4 +157,72 @@ describe('ref', () => {
     }).$mount()
     expect(vm.$refs.test).toBe(vm.$children[0])
   })
+
+  it('should should call callback method (v-for)', done => {
+    const vm = new Vue({
+      data: {
+        items: [1, 2, 3]
+      },
+      template: `
+        <div>
+          <test v-for="n in items" :key="n" :ref="onRef" :n="n"></test>
+        </div>
+      `,
+      components: {
+        test: {
+          props: ['n'],
+          template: '<div>{{ n }}</div>'
+        }
+      },
+      methods: {
+        onRef (ref) {
+          (this.$refs.list || (this.$refs.list = [])).push(ref)
+        }
+      }
+    }).$mount()
+    assertRefs()
+    // updating
+    vm.items.push(4)
+    waitForUpdate(assertRefs)
+      .then(() => { vm.items = [] })
+      .then(assertRefs)
+      .then(done)
+
+    function assertRefs () {
+      expect(Array.isArray(vm.$refs.list)).toBe(true)
+      expect(vm.$refs.list.length).toBe(vm.items.length)
+      expect(vm.$refs.list.every((comp, i) => comp.$el.textContent === String(i + 1))).toBe(true)
+    }
+  })
+
+  it('should should call inline callback (v-for)', done => {
+    const vm = new Vue({
+      data: {
+        items: [1, 2, 3]
+      },
+      template: `
+        <div>
+          <test v-for="n in items" :key="n" :ref="function (ref) { $refs[n] = ref }" :n="n"></test>
+        </div>
+      `,
+      components: {
+        test: {
+          props: ['n'],
+          template: '<div>{{ n }}</div>'
+        }
+      }
+    }).$mount()
+    assertRefs()
+    // updating
+    vm.items.push(4)
+    waitForUpdate(assertRefs)
+      .then(() => { vm.items = [] })
+      .then(assertRefs)
+      .then(done)
+
+    function assertRefs () {
+      expect(Object.keys(vm.$refs).length).toBe(vm.items.length)
+      expect(Object.keys(vm.$refs).every(i => vm.$refs[i].$el.textContent === String(i))).toBe(true)
+    }
+  })
 })