fix: remove in-between stage not working

app/home/graph.ts

@@ -48,5 +48,28 @@ const DEFAULT_NODE = {
 	},
 } as const
 
-export { $Graph, DEFAULT_NODE }
+function disconnectNode(srcNode: Node, removeNode: Node) {
+	const conn = srcNode.outs[removeNode.key]
+	if (conn) {
+		if (conn.weight === 1) {
+			delete srcNode.outs[removeNode.key]
+		} else {
+			conn.weight--
+		}
+	}
+}
+
+function connectNode(srcNode: Node, node: Node) {
+	const conn = srcNode.outs[node.key]
+	if (conn) {
+		conn.weight++
+	} else {
+		srcNode.outs[node.key] = {
+			nodeKey: node.key,
+			weight: 1,
+		}
+	}
+}
+
+export { $Graph, DEFAULT_NODE, connectNode, disconnectNode }
 export type { Graph, Node, Entry, Connection }

app/home/store.ts

@@ -1,7 +1,14 @@
 import { is } from "superstruct"
 import { create } from "zustand/index"
 import { immer } from "zustand/middleware/immer"
-import { $Graph, DEFAULT_NODE, type Entry, type Node } from "~/home/graph"
+import {
+	$Graph,
+	DEFAULT_NODE,
+	type Entry,
+	type Node,
+	connectNode,
+	disconnectNode,
+} from "~/home/graph"
 
 interface RootStore {
 	nodes: Record<string, Node>
@@ -128,20 +135,33 @@ const useRootStore = create<RootStore>()(
 				const entry = state.entries[entryName]
 				const node = state.nodes[stage]
 				if (entry && node) {
-					const lastStageNodeKey = entry.stages.at(-2)
+					const i = entry.stages.indexOf(stage)
-					if (!lastStageNodeKey) {
+					if (i < 0) {
 						return
 					}
 
-					const lastStageNode = state.nodes[lastStageNodeKey]
+					const lastStageNodeKey = entry.stages.at(i - 1)
-					const conn = lastStageNode.outs[node.key]
+					let lastStageNode: Node | null = null
-					if (conn) {
+					if (lastStageNodeKey) {
-						if (conn.weight === 1) {
+						lastStageNode = state.nodes[lastStageNodeKey]
-							delete lastStageNode.outs[node.key]
+						if (lastStageNode) {
-						} else {
+							disconnectNode(lastStageNode, node)
-							conn.weight--
 						}
 					}
+
+					const nextStageNodeKey = entry.stages.at(i + 1)
+					let nextStageNode: Node | null = null
+					if (nextStageNodeKey) {
+						nextStageNode = state.nodes[nextStageNodeKey]
+						if (nextStageNode) {
+							disconnectNode(node, nextStageNode)
+						}
+					}
+
+					if (lastStageNode && nextStageNode) {
+						connectNode(lastStageNode, nextStageNode)
+					}
+
 					entry.stages = entry.stages.filter((step) => step !== stage)
 				}
 			}),
@@ -152,14 +172,7 @@ const useRootStore = create<RootStore>()(
 				for (let i = 1; i < entry.stages.length; ++i) {
 					const lastStageNode = state.nodes[entry.stages[i - 1]]
 					const node = state.nodes[entry.stages[i]]
-					const conn = lastStageNode.outs[node.key]
+					disconnectNode(lastStageNode, node)
-					if (conn) {
-						if (conn.weight === 1) {
-							delete lastStageNode.outs[node.key]
-						} else {
-							conn.weight--
-						}
-					}
 				}
 				delete state.entries[entryName]
 			}),