From c4acb0d0f67b7a094d55d824de930e73e07d5ea3 Mon Sep 17 00:00:00 2001
From: braks <78412429+bcakmakoglu@users.noreply.github.com>
Date: Mon, 24 Jul 2023 22:10:50 +0200
Subject: [PATCH] feat(core): add flag to `removeNodes` to remove children of a
 node

---
 packages/core/src/store/actions.ts | 42 ++++++++++++++++++++++++------
 packages/core/src/types/store.ts   |  3 ++-
 2 files changed, 36 insertions(+), 9 deletions(-)

diff --git a/packages/core/src/store/actions.ts b/packages/core/src/store/actions.ts
index cac417ff..74815101 100644
--- a/packages/core/src/store/actions.ts
+++ b/packages/core/src/store/actions.ts
@@ -506,13 +506,42 @@ export function useActions(
     }
   }
 
-  const removeNodes: Actions['removeNodes'] = (nodes, removeConnectedEdges = true) => {
+  const removeNodes: Actions['removeNodes'] = (nodes, removeConnectedEdges = true, removeChildren = false) => {
     let nodesToRemove = nodes instanceof Function ? nodes(state.nodes) : nodes
     nodesToRemove = Array.isArray(nodesToRemove) ? nodesToRemove : [nodesToRemove]
 
     const nodeChanges: NodeRemoveChange[] = []
     const edgeChanges: EdgeRemoveChange[] = []
 
+    function createEdgeRemovalChanges(nodes: Node[]) {
+      const connections = getConnectedEdges(nodes, state.edges).filter((edge) => {
+        if (isDef(edge.deletable)) {
+          return edge.deletable
+        }
+        return true
+      })
+
+      edgeChanges.push(...connections.map((connection) => createRemoveChange(connection.id)))
+    }
+
+    // recursively get all children and if the child is a parent, get those children as well until all nodes have been removed that are children of the current node
+    function createChildrenRemovalChanges(id: string) {
+      const children = state.nodes.filter((n) => n.parentNode === id)
+
+      if (children.length) {
+        const childIds = children.map((n) => n.id)
+        nodeChanges.push(...childIds.map((id) => createRemoveChange(id)))
+
+        if (removeConnectedEdges) {
+          createEdgeRemovalChanges(children)
+        }
+
+        children.forEach((child) => {
+          createChildrenRemovalChanges(child.id)
+        })
+      }
+    }
+
     nodesToRemove.forEach((item) => {
       const currNode = typeof item === 'string' ? findNode(item) : item
 
@@ -527,14 +556,11 @@ export function useActions(
       nodeChanges.push(createRemoveChange(currNode.id))
 
       if (removeConnectedEdges) {
-        const connections = getConnectedEdges([currNode], state.edges).filter((edge) => {
-          if (isDef(edge.deletable)) {
-            return edge.deletable
-          }
-          return true
-        })
+        createEdgeRemovalChanges([currNode])
+      }
 
-        edgeChanges.push(...connections.map((connection) => createRemoveChange(connection.id)))
+      if (removeChildren) {
+        createChildrenRemovalChanges(currNode.id)
       }
     })
 
diff --git a/packages/core/src/types/store.ts b/packages/core/src/types/store.ts
index 81f25184..5be8f548 100644
--- a/packages/core/src/types/store.ts
+++ b/packages/core/src/types/store.ts
@@ -155,6 +155,7 @@ export type AddNodes = (nodes: Node | Node[] | ((nodes: GraphNode[]) => Node | N
 export type RemoveNodes = (
   nodes: (string | Node) | (Node | string)[] | ((nodes: GraphNode[]) => (string | Node) | (Node | string)[]),
   removeConnectedEdges?: boolean,
+  removeChildren?: boolean,
 ) => void
 
 export type RemoveEdges = (
@@ -209,7 +210,7 @@ export interface Actions extends ViewportFunctions {
   addNodes: AddNodes
   /** parses edges and adds to state */
   addEdges: AddEdges
-  /** remove nodes (and possibly connected edges) from state */
+  /** remove nodes (and possibly connected edges and children) from state */
   removeNodes: RemoveNodes
   /** remove edges from state */
   removeEdges: RemoveEdges