From 35b78920c43a2e4f7f41d43da466a830012a255e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E4=B8=89=E5=92=B2=E6=99=BA=E5=AD=90=20Kevin=20Deng?=
 <sxzz@sxzz.moe>
Date: Thu, 8 Feb 2024 20:28:10 +0800
Subject: [PATCH] feat(compiler-vapor): add getKey function for v-for

---
 .../__snapshots__/vFor.spec.ts.snap           |  2 +-
 .../__tests__/transforms/vFor.spec.ts         |  6 ++++-
 packages/compiler-vapor/src/generators/for.ts | 23 ++++++++++++++++--
 packages/compiler-vapor/src/ir.ts             |  1 +
 .../src/transforms/transformRef.ts            | 10 +++-----
 .../compiler-vapor/src/transforms/vBind.ts    | 21 ++++++++--------
 .../compiler-vapor/src/transforms/vFor.ts     |  4 ++++
 packages/compiler-vapor/src/utils.ts          | 24 +++++++++++++++++++
 8 files changed, 69 insertions(+), 22 deletions(-)
 create mode 100644 packages/compiler-vapor/src/utils.ts

diff --git a/packages/compiler-vapor/__tests__/transforms/__snapshots__/vFor.spec.ts.snap b/packages/compiler-vapor/__tests__/transforms/__snapshots__/vFor.spec.ts.snap
index 99c4a223c..af4572b36 100644
--- a/packages/compiler-vapor/__tests__/transforms/__snapshots__/vFor.spec.ts.snap
+++ b/packages/compiler-vapor/__tests__/transforms/__snapshots__/vFor.spec.ts.snap
@@ -14,7 +14,7 @@ export function render(_ctx) {
       _setText(n3, item)
     })
     return n2
-  })
+  }, (item) => (item.id))
   return [n1]
 }"
 `;
diff --git a/packages/compiler-vapor/__tests__/transforms/vFor.spec.ts b/packages/compiler-vapor/__tests__/transforms/vFor.spec.ts
index 66c0ab241..aa9a670c1 100644
--- a/packages/compiler-vapor/__tests__/transforms/vFor.spec.ts
+++ b/packages/compiler-vapor/__tests__/transforms/vFor.spec.ts
@@ -21,7 +21,7 @@ const compileWithVFor = makeCompile({
 describe('compiler: v-for', () => {
   test('basic v-for', () => {
     const { code, ir, vaporHelpers, helpers } = compileWithVFor(
-      `<div v-for="item of items" @click="remove(item)">{{ item }}</div>`,
+      `<div v-for="item of items" :key="item.id" @click="remove(item)">{{ item }}</div>`,
     )
 
     expect(code).matchSnapshot()
@@ -52,6 +52,10 @@ describe('compiler: v-for', () => {
           type: IRNodeTypes.BLOCK_FUNCTION,
           templateIndex: 0,
         },
+        keyProperty: {
+          type: NodeTypes.SIMPLE_EXPRESSION,
+          content: 'item.id',
+        },
       },
     ])
     expect(ir.returns).toEqual([1])
diff --git a/packages/compiler-vapor/src/generators/for.ts b/packages/compiler-vapor/src/generators/for.ts
index c1446e7d9..8b9bf85da 100644
--- a/packages/compiler-vapor/src/generators/for.ts
+++ b/packages/compiler-vapor/src/generators/for.ts
@@ -17,7 +17,7 @@ export function genFor(
   context: CodegenContext,
 ): CodeFragment[] {
   const { call, vaporHelper } = context
-  const { source, value, key, render } = oper
+  const { source, value, key, render, keyProperty } = oper
 
   const rawValue = value && value.content
   const rawKey = key && key.content
@@ -34,12 +34,31 @@ export function genFor(
     idMap,
   )
 
+  let getKeyFn: CodeFragment[] | false = false
+  if (keyProperty) {
+    const idMap: Record<string, null> = {}
+    if (rawValue) idMap[rawValue] = null
+    if (rawKey) idMap[rawKey] = null
+    const expr = context.withId(
+      () => genExpression(keyProperty, context),
+      idMap,
+    )
+    getKeyFn = [
+      '(',
+      rawValue ? rawValue : rawKey ? '_' : '',
+      rawKey && `, ${rawKey}`,
+      ') => (',
+      ...expr,
+      ')',
+    ]
+  }
+
   context.genEffect = undefined
 
   return [
     NEWLINE,
     `const n${oper.id} = `,
-    ...call(vaporHelper('createFor'), sourceExpr, blockFn),
+    ...call(vaporHelper('createFor'), sourceExpr, blockFn, getKeyFn),
   ]
 
   function genEffectInFor(effects: IREffect[]): CodeFragment[] {
diff --git a/packages/compiler-vapor/src/ir.ts b/packages/compiler-vapor/src/ir.ts
index d020613d8..fc2e63bfd 100644
--- a/packages/compiler-vapor/src/ir.ts
+++ b/packages/compiler-vapor/src/ir.ts
@@ -76,6 +76,7 @@ export interface ForIRNode extends BaseIRNode {
   value?: SimpleExpressionNode
   key?: SimpleExpressionNode
   index?: SimpleExpressionNode
+  keyProperty?: SimpleExpressionNode
   render: BlockFunctionIRNode
 }
 
diff --git a/packages/compiler-vapor/src/transforms/transformRef.ts b/packages/compiler-vapor/src/transforms/transformRef.ts
index f27a7b683..f09c2951d 100644
--- a/packages/compiler-vapor/src/transforms/transformRef.ts
+++ b/packages/compiler-vapor/src/transforms/transformRef.ts
@@ -1,20 +1,16 @@
 import {
-  type AttributeNode,
   NodeTypes,
   type SimpleExpressionNode,
   createSimpleExpression,
-  findProp,
 } from '@vue/compiler-dom'
 import { EMPTY_EXPRESSION, type NodeTransform } from '../transform'
-import { IRNodeTypes, type VaporDirectiveNode } from '../ir'
+import { IRNodeTypes } from '../ir'
 import { normalizeBindShorthand } from './vBind'
+import { findProp } from '../utils'
 
 export const transformRef: NodeTransform = (node, context) => {
   if (node.type !== NodeTypes.ELEMENT) return
-  const dir = findProp(node, 'ref', false, true) as
-    | VaporDirectiveNode
-    | AttributeNode
-
+  const dir = findProp(node, 'ref', false, true)
   if (!dir) return
 
   let value: SimpleExpressionNode
diff --git a/packages/compiler-vapor/src/transforms/vBind.ts b/packages/compiler-vapor/src/transforms/vBind.ts
index 1b7bae137..6be54ca79 100644
--- a/packages/compiler-vapor/src/transforms/vBind.ts
+++ b/packages/compiler-vapor/src/transforms/vBind.ts
@@ -35,19 +35,8 @@ export const transformVBind: DirectiveTransform = (dir, node, context) => {
   let { exp } = dir
   const arg = dir.arg!
 
-  if (arg.isStatic && isReservedProp(arg.content)) return
-
   if (!exp) exp = normalizeBindShorthand(arg, context)
 
-  let camel = false
-  if (modifiers.includes('camel')) {
-    if (arg.isStatic) {
-      arg.content = camelize(arg.content)
-    } else {
-      camel = true
-    }
-  }
-
   if (!exp.content.trim()) {
     if (!__BROWSER__) {
       // #10280 only error against empty expression in non-browser build
@@ -60,6 +49,16 @@ export const transformVBind: DirectiveTransform = (dir, node, context) => {
     exp = createSimpleExpression('', true, loc)
   }
 
+  if (arg.isStatic && isReservedProp(arg.content)) return
+  let camel = false
+  if (modifiers.includes('camel')) {
+    if (arg.isStatic) {
+      arg.content = camelize(arg.content)
+    } else {
+      camel = true
+    }
+  }
+
   return {
     key: arg,
     value: exp,
diff --git a/packages/compiler-vapor/src/transforms/vFor.ts b/packages/compiler-vapor/src/transforms/vFor.ts
index 01eec91cb..6980cd78f 100644
--- a/packages/compiler-vapor/src/transforms/vFor.ts
+++ b/packages/compiler-vapor/src/transforms/vFor.ts
@@ -18,6 +18,7 @@ import {
   type VaporDirectiveNode,
 } from '../ir'
 import { extend } from '@vue/shared'
+import { findProp, propToExpression } from '../utils'
 
 export const transformVFor = createStructuralDirectiveTransform(
   'for',
@@ -45,6 +46,8 @@ export function processFor(
 
   const { source, value, key, index } = parseResult
 
+  const keyProp = findProp(node, 'key')
+  const keyProperty = keyProp && propToExpression(keyProp)
   context.node = node = wrapTemplate(node, ['for'])
   context.dynamic.flags |= DynamicFlag.NON_TEMPLATE | DynamicFlag.INSERT
   const id = context.reference()
@@ -71,6 +74,7 @@ export function processFor(
       value: value as SimpleExpressionNode | undefined,
       key: key as SimpleExpressionNode | undefined,
       index: index as SimpleExpressionNode | undefined,
+      keyProperty,
       render,
     })
   }
diff --git a/packages/compiler-vapor/src/utils.ts b/packages/compiler-vapor/src/utils.ts
new file mode 100644
index 000000000..bbc8e69bf
--- /dev/null
+++ b/packages/compiler-vapor/src/utils.ts
@@ -0,0 +1,24 @@
+import {
+  type AttributeNode,
+  type ElementNode,
+  NodeTypes,
+  findProp as _findProp,
+  createSimpleExpression,
+} from '@vue/compiler-dom'
+import type { VaporDirectiveNode } from './ir'
+import { EMPTY_EXPRESSION } from './transform'
+
+export const findProp = _findProp as (
+  node: ElementNode,
+  name: string,
+  dynamicOnly?: boolean,
+  allowEmpty?: boolean,
+) => AttributeNode | VaporDirectiveNode | undefined
+
+export function propToExpression(prop: AttributeNode | VaporDirectiveNode) {
+  return prop.type === NodeTypes.ATTRIBUTE
+    ? prop.value
+      ? createSimpleExpression(prop.value.content, true, prop.value.loc)
+      : EMPTY_EXPRESSION
+    : prop.exp
+}