From bcb9209c4c19e656ca30cb5adbdd34c4532f07c8 Mon Sep 17 00:00:00 2001
From: Evan You <evan@vuejs.org>
Date: Sat, 14 Dec 2024 16:28:05 +0800
Subject: [PATCH] wip(vapor): optimize unmounted children removal

---
 packages/runtime-vapor/src/block.ts     | 18 ++++++++++
 packages/runtime-vapor/src/component.ts | 48 +++++++++++++++++++++----
 2 files changed, 59 insertions(+), 7 deletions(-)

diff --git a/packages/runtime-vapor/src/block.ts b/packages/runtime-vapor/src/block.ts
index 3951be239..65b7c70de 100644
--- a/packages/runtime-vapor/src/block.ts
+++ b/packages/runtime-vapor/src/block.ts
@@ -126,7 +126,19 @@ export function prepend(parent: ParentNode, ...blocks: Block[]): void {
   while (i--) insert(blocks[i], parent, 0)
 }
 
+/**
+ * Optimized children removal: record all parents with unmounted children
+ * during each root remove call, and update their children list by filtering
+ * unmounted children
+ */
+export let parentsWithUnmountedChildren: Set<VaporComponentInstance> | null =
+  null
+
 export function remove(block: Block, parent: ParentNode): void {
+  const isRoot = !parentsWithUnmountedChildren
+  if (isRoot) {
+    parentsWithUnmountedChildren = new Set()
+  }
   if (block instanceof Node) {
     parent.removeChild(block)
   } else if (isVaporComponent(block)) {
@@ -143,4 +155,10 @@ export function remove(block: Block, parent: ParentNode): void {
       ;(block as DynamicFragment).scope!.stop()
     }
   }
+  if (isRoot) {
+    for (const i of parentsWithUnmountedChildren!) {
+      i.children = i.children.filter(n => !n.isUnmounted)
+    }
+    parentsWithUnmountedChildren = null
+  }
 }
diff --git a/packages/runtime-vapor/src/component.ts b/packages/runtime-vapor/src/component.ts
index 2a7ffb009..61476a098 100644
--- a/packages/runtime-vapor/src/component.ts
+++ b/packages/runtime-vapor/src/component.ts
@@ -25,7 +25,13 @@ import {
   unregisterHMR,
   warn,
 } from '@vue/runtime-dom'
-import { type Block, insert, isBlock, remove } from './block'
+import {
+  type Block,
+  insert,
+  isBlock,
+  parentsWithUnmountedChildren,
+  remove,
+} from './block'
 import {
   markRaw,
   pauseTracking,
@@ -33,7 +39,14 @@ import {
   resetTracking,
   unref,
 } from '@vue/reactivity'
-import { EMPTY_OBJ, invokeArrayFns, isFunction, isString } from '@vue/shared'
+import {
+  EMPTY_ARR,
+  EMPTY_OBJ,
+  invokeArrayFns,
+  isFunction,
+  isString,
+  remove as removeItem,
+} from '@vue/shared'
 import {
   type DynamicPropsSource,
   type RawProps,
@@ -465,24 +478,45 @@ export function mountComponent(
 
 export function unmountComponent(
   instance: VaporComponentInstance,
-  parent?: ParentNode,
+  parentNode?: ParentNode,
 ): void {
   if (instance.isMounted && !instance.isUnmounted) {
     if (__DEV__ && instance.type.__hmrId) {
       unregisterHMR(instance)
     }
-    if (instance.bum) invokeArrayFns(instance.bum)
+    if (instance.bum) {
+      invokeArrayFns(instance.bum)
+    }
+
     instance.scope.stop()
+
     for (const c of instance.children) {
       unmountComponent(c)
     }
-    if (parent) remove(instance.block, parent)
+    instance.children = EMPTY_ARR as any
+
+    if (parentNode) {
+      // root remove: need to both remove this instance's DOM nodes
+      // and also remove it from the parent's children list.
+      remove(instance.block, parentNode)
+      const parentInstance = instance.parent
+      if (isVaporComponent(parentInstance)) {
+        if (parentsWithUnmountedChildren) {
+          // for optimize children removal
+          parentsWithUnmountedChildren.add(parentInstance)
+        } else {
+          removeItem(parentInstance.children, instance)
+        }
+        instance.parent = null
+      }
+    }
+
     if (instance.um) {
       queuePostFlushCb(() => invokeArrayFns(instance.um!))
     }
     instance.isUnmounted = true
-  } else if (parent) {
-    remove(instance.block, parent)
+  } else if (parentNode) {
+    remove(instance.block, parentNode)
   }
 }