feat: cleanup selector during vFor item unmount

Author: edison1105

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

@@ -120,18 +120,19 @@ exports[`compiler: v-for > selector pattern 1`] = `
 const t0 = _template("<tr> </tr>", true)
 
 export function render(_ctx) {
-  const _selector0 = _useSelectorPattern(() => _ctx.selected, (..._args) => {
+  const _selector0_0 = _useSelectorPattern(() => _ctx.selected, (..._args) => {
   _setClass(_args[0], _args[1])
   }, 'danger', '')
   const n0 = _createFor(() => (_ctx.rows), (_for_item0) => {
-    _selector0.register(_for_item0.value.id, _for_item0)
+    const _cleanup0 = _selector0_0.register(_for_item0.value.id, _for_item0)
     const n2 = t0()
     const x2 = _child(n2)
-    _renderEffect(() => {
+    const _e = _renderEffect(() => {
       const _row = _for_item0.value
       const _row_id = _row.id
       _setText(x2, _toDisplayString(_row_id))
     })
+    _e.cleanup = () => _cleanup0()
     return n2
   }, (row) => (row.id))
   return n0

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

@@ -7,7 +7,7 @@ import {
 import { genBlockContent } from './block'
 import { genExpression } from './expression'
 import type { CodegenContext } from '../generate'
-import type { BlockIRNode, ForIRNode, IREffect } from '../ir'
+import type { BlockIRNode, ForIRNode, IREffect, SetPropIRNode } from '../ir'
 import {
   type CodeFragment,
   INDENT_END,
@@ -16,7 +16,7 @@ import {
   genCall,
   genMulti,
 } from './utils'
-import { isNodesEquivalent, type Identifier } from '@babel/types'
+import { type Identifier, isNodesEquivalent } from '@babel/types'
 import { parseExpression } from '@babel/parser'
 import { VaporVForFlags } from '../../../shared/src/vaporFlags'
 import { getRuntimeHelper } from './prop'
@@ -92,13 +92,13 @@ export function genFor(
 
   for (let i = 0; i < selectorPatterns.length; i++) {
     const { selector, trueValue, falseValue, effect } = selectorPatterns[i]
-    const selectorName = `_selector${i}`
+    const selectorName = `_selector${id}_${i}`
     const args: CodeFragment[] = [
       `() => `,
       ...genExpression(selector, context),
       `, (..._args) => {`,
     ]
-    for (const oper of effect.operations) {
+    for (const oper of effect.operations as SetPropIRNode[]) {
       const resolvedHelper = getRuntimeHelper(
         oper.tag,
         oper.prop.key.content,
@@ -123,13 +123,17 @@ export function genFor(
     for (let i = 0; i < selectorPatterns.length; i++) {
       frag.push(
         NEWLINE,
-        `_selector${i}.register(`,
+        `const _cleanup${id} = _selector${id}_${i}.register(`,
         ...genExpression(keyProp!, context),
         `, `,
         itemVar,
         `)`,
       )
     }
+    if (selectorPatterns.length) {
+      // save cleanup function to the first effect
+      render.effect[0].cleanup = [`_cleanup${id}()`] as CodeFragment[]
+    }
     frag.push(...genBlockContent(render, context))
     frag.push(INDENT_END, NEWLINE, '}')
     return frag
@@ -333,7 +337,8 @@ function matchSelectorPattern(
     if (
       keyProp !== undefined &&
       typeof ast === 'object' &&
-      ast?.type === 'ConditionalExpression' &&
+      ast &&
+      ast.type === 'ConditionalExpression' &&
       ast.test.type === 'BinaryExpression' &&
       ast.test.operator === '===' &&
       ast.test.left.type !== 'PrivateName' &&

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

@@ -128,8 +128,18 @@ export function genEffects(
   }
 
   if (effects.length) {
-    unshift(NEWLINE, `${helper('renderEffect')}(() => `)
+    // cleanup is only used for `vFor` + `useSelectorPattern`
+    // the cleanup code fragment can be found in the first effect
+    const cleanup = effects.length === 1 && effects[0].cleanup
+    unshift(
+      NEWLINE,
+      cleanup ? `const _e = ` : '',
+      `${helper('renderEffect')}(() => `,
+    )
     push(`)`)
+    if (cleanup) {
+      push(NEWLINE, `_e.cleanup = () => `, ...cleanup)
+    }
   }
 
   return frag

packages/compiler-vapor/src/ir/index.ts

@@ -8,6 +8,7 @@ import type {
 import type { Prettify } from '@vue/shared'
 import type { DirectiveTransform, NodeTransform } from '../transform'
 import type { IRProp, IRProps, IRSlots } from './component'
+import type { CodeFragment } from '..'
 
 export * from './component'
 
@@ -266,6 +267,7 @@ export interface IRDynamicInfo {
 export interface IREffect {
   expressions: SimpleExpressionNode[]
   operations: OperationNode[]
+  cleanup?: CodeFragment[]
 }
 
 type Overwrite<T, U> = Pick<T, Exclude<keyof T, keyof U>> &

packages/runtime-vapor/src/apiCreateFor.ts

@@ -3,11 +3,13 @@ import {
   type ShallowRef,
   isReactive,
   isShallow,
+  onScopeDispose,
   pauseTracking,
   resetTracking,
   shallowReadArray,
   shallowRef,
   toReactive,
+  watch,
 } from '@vue/reactivity'
 import { getSequence, isArray, isObject, isString } from '@vue/shared'
 import { createComment, createTextNode } from './dom/node'
@@ -451,3 +453,50 @@ export function getRestElement(val: any, keys: string[]): any {
 export function getDefaultValue(val: any, defaultVal: any): any {
   return val === undefined ? defaultVal : val
 }
+
+export function useSelectorPattern(
+  getActiveKey: () => any,
+  setter: (element: any, value: any) => void,
+  trueValue: any,
+  falseValue: any,
+): {
+  register: (key: any, element: any) => () => void
+} {
+  const nodeMap = new Map()
+  let trueElements: any[] | undefined
+
+  onScopeDispose(() => nodeMap.clear(), true)
+
+  watch(getActiveKey, newValue => {
+    if (trueElements !== undefined) {
+      for (const n of trueElements) {
+        setter(n, falseValue)
+      }
+    }
+    trueElements = nodeMap.get(newValue)
+    if (trueElements !== undefined) {
+      for (const n of trueElements) {
+        setter(n, trueValue)
+      }
+    }
+  })
+
+  return { register }
+
+  function register(key: any, element: any): () => void {
+    const isTrue = getActiveKey() === key
+    setter(element, isTrue ? trueValue : falseValue)
+    let elements = nodeMap.get(key)
+    if (elements !== undefined) {
+      elements.push(element)
+    } else {
+      elements = [element]
+      nodeMap.set(key, elements)
+      if (isTrue) {
+        trueElements = elements
+      }
+    }
+
+    return () => nodeMap.delete(key)
+  }
+}

packages/runtime-vapor/src/index.ts

@@ -37,6 +37,7 @@ export {
   createForSlots,
   getRestElement,
   getDefaultValue,
+  useSelectorPattern,
 } from './apiCreateFor'
 export { createTemplateRefSetter } from './apiTemplateRef'
 export { createDynamicComponent } from './apiCreateDynamicComponent'
@@ -54,48 +55,3 @@ export { withVaporDirectives } from './directives/custom'
  * @internal
  */
 export { simpleSetCurrentInstance } from '@vue/runtime-core'
-
-import { watch } from '@vue/reactivity'
-
-export function useSelectorPattern(
-  getActiveKey: () => any,
-  setter: (element: any, value: any) => void,
-  trueValue: any,
-  falseValue: any,
-): {
-  register: (key: any, element: any) => void
-} {
-  const nodeMap = new Map()
-  let trueElements: any[] | undefined
-
-  watch(getActiveKey, newValue => {
-    if (trueElements !== undefined) {
-      for (const n of trueElements) {
-        setter(n, falseValue)
-      }
-    }
-    trueElements = nodeMap.get(newValue)
-    if (trueElements !== undefined) {
-      for (const n of trueElements) {
-        setter(n, trueValue)
-      }
-    }
-  })
-
-  return { register }
-
-  function register(key: any, element: any): void {
-    const isTrue = getActiveKey() === key
-    setter(element, isTrue ? trueValue : falseValue)
-    let elements = nodeMap.get(key)
-    if (elements !== undefined) {
-      elements.push(element)
-    } else {
-      elements = [element]
-      nodeMap.set(key, elements)
-      if (isTrue) {
-        trueElements = elements
-      }
-    }
-  }
-}

packages/runtime-vapor/src/renderEffect.ts

@@ -58,7 +58,7 @@ class RenderEffect extends ReactiveEffect {
     // TODO recurse handling
   }
 
-  callback() {
+  callback(): void {
     const instance = this.i!
     const scope = this.scope
     // renderEffect is always called after user has registered all hooks
@@ -84,7 +84,7 @@ class RenderEffect extends ReactiveEffect {
     }
   }
 
-  scheduler() {
+  scheduler(): void {
     queueJob(this.baseJob)
   }
 }
@@ -99,10 +99,16 @@ class RenderEffect_NoLifecycle extends RenderEffect {
   }
 }
 
-export function renderEffect(fn: () => void, noLifecycle = false): void {
+export function renderEffect(
+  fn: () => void,
+  noLifecycle = false,
+): RenderEffect {
+  let effect
   if (noLifecycle) {
-    new RenderEffect_NoLifecycle(fn).run()
+    effect = new RenderEffect_NoLifecycle(fn)
   } else {
-    new RenderEffect(fn).run()
+    effect = new RenderEffect(fn)
   }
+  effect.run()
+  return effect
 }