feat(compiler): support templateRef

Author: zhiyuanzmj

src/core/compiler/compile.ts

@@ -18,6 +18,7 @@ import {
 
 import { transformElement } from './transforms/transformElement'
 import { transformChildren } from './transforms/transformChildren'
+import { transformTemplateRef } from './transforms/transformTemplateRef'
 import {
   type HackOptions,
   IRNodeTypes,
@@ -132,7 +133,13 @@ export function getBaseTransformPreset(
   prefixIdentifiers?: boolean,
 ): TransformPreset {
   return [
-    [transformText, transformElement, transformVSlot, transformChildren],
+    [
+      transformTemplateRef,
+      transformText,
+      transformElement,
+      transformVSlot,
+      transformChildren,
+    ],
     {
       bind: transformVBind,
       on: transformVOn,

src/core/compiler/index.ts

@@ -10,6 +10,7 @@ export * from './ir'
 export { transformText } from './transforms/transformText'
 export { transformElement } from './transforms/transformElement'
 export { transformChildren } from './transforms/transformChildren'
+export { transformTemplateRef } from './transforms/transformTemplateRef'
 export { transformVBind } from './transforms/vBind'
 export { transformVOn } from './transforms/vOn'
 export { transformVSlot } from './transforms/vSlot'

src/core/compiler/transforms/transformTemplateRef.ts

@@ -0,0 +1,29 @@
+import { IRNodeTypes } from '../ir'
+import { findProp, isConstantExpression, resolveExpression } from '../utils'
+import type { NodeTransform } from '../transform'
+
+export const transformTemplateRef: NodeTransform = (node, context) => {
+  if (node.type !== 'JSXElement') return
+
+  const dir = findProp(node, 'ref')
+  if (!dir?.value) return
+
+  const value = resolveExpression(dir.value, context)
+
+  return () => {
+    const id = context.reference()
+    const effect = !isConstantExpression(value)
+    effect &&
+      context.registerOperation({
+        type: IRNodeTypes.DECLARE_OLD_REF,
+        id,
+      })
+    context.registerEffect([value], {
+      type: IRNodeTypes.SET_TEMPLATE_REF,
+      element: id,
+      value,
+      refFor: !!context.inVFor,
+      effect,
+    })
+  }
+}

src/core/compiler/transforms/vFor.ts

@@ -29,7 +29,8 @@ export function processMapCallExpression(
     argument.params[2] && resolveExpression(argument.params[2], context)
 
   const returnExpression = getReturnExpression(argument)
-  const keyProperty = findProp(returnExpression, context)
+  const keyProp = findProp(returnExpression, 'key')
+  const keyProperty = keyProp && resolveExpression(keyProp.value, context)
   return () => {
     exitBlock()
     context.registerOperation({

src/core/compiler/utils.ts

@@ -266,14 +266,11 @@ export function isMapCallExpression(
   )
 }
 
-export function findProp(
-  expression: Expression | undefined,
-  context: TransformContext,
-) {
+export function findProp(expression: Expression | undefined, key: string) {
   if (expression?.type === 'JSXElement') {
     for (const attr of expression.openingElement.attributes) {
-      if (attr.type === 'JSXAttribute' && attr.name.name === 'key') {
-        return resolveExpression(attr.value, context)
+      if (attr.type === 'JSXAttribute' && attr.name.name === key) {
+        return attr
       }
     }
   }

test/transforms/__snapshots__/transformTemplateRef.spec.ts.snap

@@ -0,0 +1,24 @@
+// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
+
+exports[`compiler: template ref transform > dynamic ref 1`] = `
+"import { renderEffect as _renderEffect, setRef as _setRef, template as _template } from 'vue/vapor';
+const t0 = _template("<div></div>")
+
+export function render(_ctx) {
+  const n0 = t0()
+  let r0
+  _renderEffect(() => r0 = _setRef(n0, foo, r0))
+  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
+}"
+`;

test/transforms/transformTemplateRef.spec.ts

@@ -0,0 +1,120 @@
+import { describe, expect, test } from 'vitest'
+import {
+  DynamicFlag,
+  IRNodeTypes,
+  transformChildren,
+  transformElement,
+  transformTemplateRef,
+} from '../../src/core/compiler/index'
+
+import { makeCompile } from './_utils'
+
+const compileWithTransformRef = makeCompile({
+  nodeTransforms: [transformTemplateRef, transformElement, transformChildren],
+  prefixIdentifiers: false,
+})
+
+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_TEMPLATE_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).toMatchObject([
+      {
+        type: IRNodeTypes.DECLARE_OLD_REF,
+        id: 0,
+      },
+    ])
+    expect(ir.block.effect).toMatchObject([
+      {
+        operations: [
+          {
+            type: IRNodeTypes.SET_TEMPLATE_REF,
+            element: 0,
+            value: {
+              content: 'foo',
+              isStatic: false,
+            },
+          },
+        ],
+      },
+    ])
+    expect(code).matchSnapshot()
+    expect(code).contains('_setRef(n0, foo, r0)')
+  })
+
+  // 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).toMatchObject([
+  //     {
+  //       type: IRNodeTypes.SET_TEMPLATE_REF,
+  //       element: 2,
+  //       value: {
+  //         content: 'foo',
+  //         isStatic: true,
+  //       },
+  //       effect: false,
+  //     },
+  //   ])
+  //   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).toMatchObject([
+  //     {
+  //       type: IRNodeTypes.SET_TEMPLATE_REF,
+  //       element: 2,
+  //       value: {
+  //         content: 'foo',
+  //         isStatic: true,
+  //       },
+  //       refFor: true,
+  //       effect: false,
+  //     },
+  //   ])
+  //   expect(code).matchSnapshot()
+  //   expect(code).contains('_setRef(n2, "foo", void 0, true)')
+  // })
+})