From 7204ac79c5c6e05fd0e121c96c459d59ecd8743d Mon Sep 17 00:00:00 2001
From: Setzer22 <jsanchezfsms@gmail.com>
Date: Tue, 7 Jun 2022 21:44:51 +0200
Subject: [PATCH] Address review comments

---
 egui_node_graph/src/editor_ui.rs   | 10 ++++++--
 egui_node_graph/src/graph_impls.rs | 38 +++++++++++++++++++++---------
 2 files changed, 35 insertions(+), 13 deletions(-)

diff --git a/egui_node_graph/src/editor_ui.rs b/egui_node_graph/src/editor_ui.rs
index 4896b04..03bb537 100644
--- a/egui_node_graph/src/editor_ui.rs
+++ b/egui_node_graph/src/editor_ui.rs
@@ -203,8 +203,14 @@ where
                     self.selected_node = Some(node_id);
                 }
                 NodeResponse::DeleteNode(node_id) => {
-                    let (ins, _) = self.graph.remove_node(node_id);
-                    extra_responses.extend(ins.into_iter().map(NodeResponse::DisconnectEvent));
+                    let removed = self.graph.remove_node(node_id);
+                    extra_responses.extend(
+                        removed
+                            .into_iter()
+                            .map(|x| x.0)
+                            .into_iter()
+                            .map(NodeResponse::DisconnectEvent),
+                    );
                     self.node_positions.remove(node_id);
                     // Make sure to not leave references to old nodes hanging
                     if self.selected_node.map(|x| x == node_id).unwrap_or(false) {
diff --git a/egui_node_graph/src/graph_impls.rs b/egui_node_graph/src/graph_impls.rs
index 42cd96e..9c26881 100644
--- a/egui_node_graph/src/graph_impls.rs
+++ b/egui_node_graph/src/graph_impls.rs
@@ -63,20 +63,36 @@ impl<NodeData, DataType, ValueType> Graph<NodeData, DataType, ValueType> {
         output_id
     }
 
-    /// Returns the list of input and output ids that were disconnected.
-    pub fn remove_node(&mut self, node_id: NodeId) -> (SVec<InputId>, SVec<OutputId>) {
-        self.connections
-            .retain(|i, o| !(self.outputs[*o].node == node_id || self.inputs[i].node == node_id));
-        let inputs: SVec<_> = self[node_id].input_ids().collect();
-        for input in &inputs {
-            self.inputs.remove(*input);
+    /// Removes a node from the graph with given `node_id`. This also removes
+    /// any incoming or outgoing connections from that node
+    ///
+    /// This function returns the list of connections that has been removed
+    /// after deleting this node as input-output pairs. Note that one of the two
+    /// ids in the pair (the one on `node_id`'s end) will be invalid after
+    /// calling this function.
+    pub fn remove_node(&mut self, node_id: NodeId) -> Vec<(InputId, OutputId)> {
+        let mut disconnect_events = vec![];
+
+        self.connections.retain(|i, o| {
+            if self.outputs[*o].node == node_id || self.inputs[i].node == node_id {
+                disconnect_events.push((i, *o));
+                false
+            } else {
+                true
+            }
+        });
+
+        // NOTE: Collect is needed because we can't borrow the input ids while
+        // we remove them inside the loop.
+        for input in self[node_id].input_ids().collect::<SVec<_>>() {
+            self.inputs.remove(input);
         }
-        let outputs: SVec<_> = self[node_id].output_ids().collect();
-        for output in &outputs {
-            self.outputs.remove(*output);
+        for output in self[node_id].output_ids().collect::<SVec<_>>() {
+            self.outputs.remove(output);
         }
         self.nodes.remove(node_id);
-        (inputs, outputs)
+
+        disconnect_events
     }
 
     pub fn remove_connection(&mut self, input_id: InputId) -> Option<OutputId> {