feat(compiler-vapor): ref for `v-for` (#167)

Co-authored-by: Rizumu Ayaka <rizumu@ayaka.moe>
Co-authored-by: 三咲智子 Kevin Deng <sxzz@sxzz.moe>

packages/compiler-vapor/__tests__/transforms/__snapshots__/transformRef.spec.ts.snap

@@ -0,0 +1,51 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`compiler: template ref transform > dynamic ref 1`] = `
+"import { setRef as _setRef, template as _template } from 'vue/vapor';
+const t0 = _template("<div></div>")
+
+export function render(_ctx) {
+  const n0 = t0()
+  _setRef(n0, _ctx.foo)
+  return n0
+}"
+`;
+
+exports[`compiler: template ref transform > ref + v-for 1`] = `
+"import { setRef as _setRef, createFor as _createFor, template as _template } from 'vue/vapor';
+const t0 = _template("<div></div>")
+
+export function render(_ctx) {
+  const n0 = _createFor(() => ([1,2,3]), (_block) => {
+    const n2 = t0()
+    _setRef(n2, "foo", true)
+    return [n2, () => {}]
+  })
+  return n0
+}"
+`;
+
+exports[`compiler: template ref transform > ref + v-if 1`] = `
+"import { setRef as _setRef, createIf as _createIf, template as _template } from 'vue/vapor';
+const t0 = _template("<div></div>")
+
+export function render(_ctx) {
+  const n0 = _createIf(() => (true), () => {
+    const n2 = t0()
+    _setRef(n2, "foo")
+    return n2
+  })
+  return n0
+}"
+`;
+
+exports[`compiler: template ref transform > static ref 1`] = `
+"import { setRef as _setRef, template as _template } from 'vue/vapor';
+const t0 = _template("<div></div>")
+
+export function render(_ctx) {
+  const n0 = t0()
+  _setRef(n0, "foo")
+  return n0
+}"
+`;

packages/compiler-vapor/__tests__/transforms/transformRef.spec.ts

@@ -1,4 +1,127 @@
-// TODO: add tests for this transform
-describe('compiler: template ref transform', () => {
-  test.todo('basic')
+import {
+  DynamicFlag,
+  type ForIRNode,
+  IRNodeTypes,
+  type IfIRNode,
+  transformChildren,
+  transformElement,
+  transformRef,
+  transformVFor,
+  transformVIf,
+} from '../../src'
+import { makeCompile } from './_utils'
+
+const compileWithTransformRef = makeCompile({
+  nodeTransforms: [
+    transformVIf,
+    transformVFor,
+    transformRef,
+    transformElement,
+    transformChildren,
+  ],
+})
+
+describe('compiler: template ref transform', () => {
+  test('static ref', () => {
+    const { ir, code } = compileWithTransformRef(`<div ref="foo" />`)
+
+    expect(ir.block.dynamic.children[0]).toMatchObject({
+      id: 0,
+      flags: DynamicFlag.REFERENCED,
+    })
+    expect(ir.template).toEqual(['<div></div>'])
+    expect(ir.block.operation).lengthOf(1)
+    expect(ir.block.operation[0]).toMatchObject({
+      type: IRNodeTypes.SET_REF,
+      element: 0,
+      value: {
+        content: 'foo',
+        isStatic: true,
+        loc: {
+          start: { line: 1, column: 10, offset: 9 },
+          end: { line: 1, column: 15, offset: 14 },
+        },
+      },
+    })
+
+    expect(code).matchSnapshot()
+    expect(code).contains('_setRef(n0, "foo")')
+  })
+
+  test('dynamic ref', () => {
+    const { ir, code } = compileWithTransformRef(`<div :ref="foo" />`)
+
+    expect(ir.block.dynamic.children[0]).toMatchObject({
+      id: 0,
+      flags: DynamicFlag.REFERENCED,
+    })
+    expect(ir.template).toEqual(['<div></div>'])
+    expect(ir.block.operation).lengthOf(1)
+    expect(ir.block.operation[0]).toMatchObject({
+      type: IRNodeTypes.SET_REF,
+      element: 0,
+      value: {
+        content: 'foo',
+        isStatic: false,
+        loc: {
+          start: { line: 1, column: 12, offset: 11 },
+          end: { line: 1, column: 15, offset: 14 },
+        },
+      },
+    })
+    expect(code).matchSnapshot()
+    expect(code).contains('_setRef(n0, _ctx.foo)')
+  })
+
+  test('ref + v-if', () => {
+    const { ir, code } = compileWithTransformRef(
+      `<div ref="foo" v-if="true" />`,
+    )
+
+    expect(ir.block.operation).lengthOf(1)
+    expect(ir.block.operation[0].type).toBe(IRNodeTypes.IF)
+
+    const { positive } = ir.block.operation[0] as IfIRNode
+
+    expect(positive.operation).lengthOf(1)
+    expect(positive.operation[0]).toMatchObject({
+      type: IRNodeTypes.SET_REF,
+      element: 2,
+      value: {
+        content: 'foo',
+        isStatic: true,
+        loc: {
+          start: { line: 1, column: 10, offset: 9 },
+          end: { line: 1, column: 15, offset: 14 },
+        },
+      },
+    })
+
+    expect(code).matchSnapshot()
+    expect(code).contains('_setRef(n2, "foo")')
+  })
+
+  test('ref + v-for', () => {
+    const { ir, code } = compileWithTransformRef(
+      `<div ref="foo" v-for="item in [1,2,3]" />`,
+    )
+
+    const { render } = ir.block.operation[0] as ForIRNode
+    expect(render.operation).lengthOf(1)
+    expect(render.operation[0]).toMatchObject({
+      type: IRNodeTypes.SET_REF,
+      element: 2,
+      value: {
+        content: 'foo',
+        isStatic: true,
+        loc: {
+          start: { line: 1, column: 10, offset: 9 },
+          end: { line: 1, column: 15, offset: 14 },
+        },
+      },
+      refFor: true,
+    })
+    expect(code).matchSnapshot()
+    expect(code).contains('_setRef(n2, "foo", true)')
+  })
 })

packages/compiler-vapor/src/compile.ts

@@ -102,9 +102,9 @@ export function getBaseTransformPreset(
   return [
     [
       transformOnce,
-      transformRef,
       transformVIf,
       transformVFor,
+      transformRef,
       transformText,
       transformElement,
       transformComment,

packages/compiler-vapor/src/generators/ref.ts

@@ -14,6 +14,7 @@ export function genSetRef(
       vaporHelper('setRef'),
       [`n${oper.element}`],
       genExpression(oper.value, context),
+      oper.refFor && 'true',
     ),
   ]
 }

packages/compiler-vapor/src/index.ts

@@ -36,6 +36,7 @@ export {
 
 export { transformElement } from './transforms/transformElement'
 export { transformChildren } from './transforms/transformChildren'
+export { transformRef } from './transforms/transformRef'
 export { transformText } from './transforms/transformText'
 export { transformVBind } from './transforms/vBind'
 export { transformVHtml } from './transforms/vHtml'

packages/compiler-vapor/src/ir.ts

@@ -138,6 +138,7 @@ export interface SetRefIRNode extends BaseIRNode {
   type: IRNodeTypes.SET_REF
   element: number
   value: SimpleExpressionNode
+  refFor: boolean
 }
 
 export interface SetModelValueIRNode extends BaseIRNode {

packages/compiler-vapor/src/transform.ts

@@ -72,8 +72,9 @@ export interface TransformContext<T extends AllNode = AllNode> {
   comment: CommentNode[]
 
   inVOnce: boolean
+  inVFor: number
 
-  enterBlock(ir: TransformContext['block']): () => void
+  enterBlock(ir: TransformContext['block'], isVFor?: boolean): () => void
   reference(): number
   increaseId(): number
   registerTemplate(): number
@@ -122,23 +123,26 @@ function createRootContext(
     index: 0,
     root: null!, // set later
     block: root.block,
-    enterBlock(ir) {
+    enterBlock(ir, inVFor = false) {
       const { block, template, dynamic, childrenTemplate } = this
       this.block = ir
       this.dynamic = ir.dynamic
       this.template = ''
       this.childrenTemplate = []
+      inVFor && this.inVFor++
       return () => {
         // exit
         this.block = block
         this.template = template
         this.dynamic = dynamic
         this.childrenTemplate = childrenTemplate
+        inVFor && this.inVFor--
       }
     },
     options: extend({}, defaultOptions, options),
     dynamic: root.block.dynamic,
     inVOnce: false,
+    inVFor: 0,
     comment: [],
 
     increaseId: () => globalId++,

packages/compiler-vapor/src/transforms/transformRef.ts

@@ -11,6 +11,7 @@ import { EMPTY_EXPRESSION } from './utils'
 
 export const transformRef: NodeTransform = (node, context) => {
   if (node.type !== NodeTypes.ELEMENT) return
+
   const dir = findProp(node, 'ref', false, true)
   if (!dir) return
 
@@ -22,10 +23,12 @@ export const transformRef: NodeTransform = (node, context) => {
       ? createSimpleExpression(dir.value.content, true, dir.value.loc)
       : EMPTY_EXPRESSION
   }
+
   return () =>
     context.registerOperation({
       type: IRNodeTypes.SET_REF,
       element: context.reference(),
       value,
+      refFor: !!context.inVFor,
     })
 }

packages/compiler-vapor/src/transforms/vFor.ts

@@ -60,7 +60,7 @@ export function processFor(
     operation: [],
     returns: [],
   }
-  const exitBlock = context.enterBlock(render)
+  const exitBlock = context.enterBlock(render, true)
   context.reference()
 
   return () => {

packages/runtime-vapor/src/dom/templateRef.ts

@@ -1,7 +1,19 @@
-import { type Ref, type SchedulerJob, isRef } from '@vue/reactivity'
+import {
+  type Ref,
+  type SchedulerJob,
+  isRef,
+  onScopeDispose,
+} from '@vue/reactivity'
 import { currentInstance } from '../component'
 import { VaporErrorCodes, callWithErrorHandling } from '../errorHandling'
-import { EMPTY_OBJ, hasOwn, isFunction, isString } from '@vue/shared'
+import {
+  EMPTY_OBJ,
+  hasOwn,
+  isArray,
+  isFunction,
+  isString,
+  remove,
+} from '@vue/shared'
 import { warn } from '../warning'
 import { queuePostRenderEffect } from '../scheduler'
 
@@ -10,48 +22,90 @@ export type NodeRef = string | Ref | ((ref: Element) => void)
 /**
  * Function for handling a template ref
  */
-export function setRef(el: Element, ref: NodeRef) {
+export function setRef(el: Element, ref: NodeRef, refFor = false) {
   if (!currentInstance) return
   const { setupState, isUnmounted } = currentInstance
 
-  const value = isUnmounted ? null : el
+  if (isUnmounted) {
+    return
+  }
+
   const refs =
     currentInstance.refs === EMPTY_OBJ
       ? (currentInstance.refs = {})
       : currentInstance.refs
 
   if (isFunction(ref)) {
-    callWithErrorHandling(ref, currentInstance, VaporErrorCodes.FUNCTION_REF, [
-      value,
-      refs,
-    ])
+    const invokeRefSetter = (value: Element | null) => {
+      callWithErrorHandling(
+        ref,
+        currentInstance,
+        VaporErrorCodes.FUNCTION_REF,
+        [value, refs],
+      )
+    }
+
+    invokeRefSetter(el)
+    onScopeDispose(() => invokeRefSetter(null))
   } else {
     const _isString = isString(ref)
     const _isRef = isRef(ref)
+    let existing: unknown
 
     if (_isString || _isRef) {
-      const doSet = () => {
-        if (_isString) {
-          refs[ref] = value
+      const doSet: SchedulerJob = () => {
+        if (refFor) {
+          existing = _isString
+            ? hasOwn(setupState, ref)
+              ? setupState[ref]
+              : refs[ref]
+            : ref.value
+
+          if (!isArray(existing)) {
+            existing = [el]
+            if (_isString) {
+              refs[ref] = existing
+              if (hasOwn(setupState, ref)) {
+                setupState[ref] = refs[ref]
+                // if setupState[ref] is a reactivity ref,
+                // the existing will also become reactivity too
+                // need to get the Proxy object by resetting
+                existing = setupState[ref]
+              }
+            } else {
+              ref.value = existing
+            }
+          } else if (!existing.includes(el)) {
+            existing.push(el)
+          }
+        } else if (_isString) {
+          refs[ref] = el
           if (hasOwn(setupState, ref)) {
-            setupState[ref] = value
+            setupState[ref] = el
           }
         } else if (_isRef) {
-          ref.value = value
+          ref.value = el
         } else if (__DEV__) {
           warn('Invalid template ref type:', ref, `(${typeof ref})`)
         }
       }
-      // #9908 ref on v-for mutates the same array for both mount and unmount
-      // and should be done together
-      if (isUnmounted /* || isVFor */) {
-        doSet()
-      } else {
-        // #1789: set new refs in a post job so that they don't get overwritten
-        // by unmounting ones.
-        ;(doSet as SchedulerJob).id = -1
-        queuePostRenderEffect(doSet)
-      }
+      doSet.id = -1
+      queuePostRenderEffect(doSet)
+
+      onScopeDispose(() => {
+        queuePostRenderEffect(() => {
+          if (isArray(existing)) {
+            remove(existing, el)
+          } else if (_isString) {
+            refs[ref] = null
+            if (hasOwn(setupState, ref)) {
+              setupState[ref] = null
+            }
+          } else if (_isRef) {
+            ref.value = null
+          }
+        })
+      })
     } else if (__DEV__) {
       warn('Invalid template ref type:', ref, `(${typeof ref})`)
     }