From d6f3a240517fdc787665f38fd1148a0b1aa0b18a Mon Sep 17 00:00:00 2001
From: moklick <info@moritzklack.com>
Date: Tue, 22 Mar 2022 12:07:24 +0100
Subject: [PATCH] fix(sub-flows): delete child nodes when parent gets deleted

---
 src/hooks/useGlobalKeyHandler.ts | 25 ++++++++++++++++---------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/src/hooks/useGlobalKeyHandler.ts b/src/hooks/useGlobalKeyHandler.ts
index 1dff6db1..2db51868 100644
--- a/src/hooks/useGlobalKeyHandler.ts
+++ b/src/hooks/useGlobalKeyHandler.ts
@@ -4,7 +4,7 @@ import shallow from 'zustand/shallow';
 import { useStore, useStoreApi } from '../store';
 import useKeyPress from './useKeyPress';
 import { getConnectedEdges } from '../utils/graph';
-import { EdgeChange, KeyCode, NodeChange, ReactFlowState } from '../types';
+import { EdgeChange, KeyCode, NodeChange, Node, ReactFlowState } from '../types';
 
 interface HookParams {
   deleteKeyCode: KeyCode | null;
@@ -26,18 +26,25 @@ export default ({ deleteKeyCode, multiSelectionKeyCode }: HookParams): void => {
 
   useEffect(() => {
     const { nodeInternals, edges, hasDefaultNodes, hasDefaultEdges, onNodesDelete, onEdgesDelete } = store.getState();
-    // @TODO: work with nodeInternals instead of converting it to an array
-    const nodes = Array.from(nodeInternals).map(([_, node]) => node);
-    const selectedNodes = nodes.filter((n) => n.selected);
+    const nodes = Array.from(nodeInternals.values());
+    const nodesToRemove = nodes.reduce<Node[]>((res, node) => {
+      if (!node.selected && node.parentNode && res.find((n) => n.id === node.parentNode)) {
+        res.push(node);
+      } else if (node.selected) {
+        res.push(node);
+      }
+
+      return res;
+    }, []);
     const selectedEdges = edges.filter((e) => e.selected);
 
-    if (deleteKeyPressed && (selectedNodes || selectedEdges)) {
-      const connectedEdges = getConnectedEdges(selectedNodes, edges);
+    if (deleteKeyPressed && (nodesToRemove || selectedEdges)) {
+      const connectedEdges = getConnectedEdges(nodesToRemove, edges);
       const edgesToRemove = [...selectedEdges, ...connectedEdges];
       const edgeIdsToRemove = edgesToRemove.map((e) => e.id);
 
       if (hasDefaultNodes) {
-        selectedNodes.forEach((node) => {
+        nodesToRemove.forEach((node) => {
           nodeInternals.delete(node.id);
         });
       }
@@ -49,11 +56,11 @@ export default ({ deleteKeyCode, multiSelectionKeyCode }: HookParams): void => {
         });
       }
 
-      onNodesDelete?.(selectedNodes);
+      onNodesDelete?.(nodesToRemove);
       onEdgesDelete?.(edgesToRemove);
 
       if (onNodesChange) {
-        const nodeChanges: NodeChange[] = selectedNodes.map((n) => ({ id: n.id, type: 'remove' }));
+        const nodeChanges: NodeChange[] = nodesToRemove.map((n) => ({ id: n.id, type: 'remove' }));
         onNodesChange(nodeChanges);
       }