add traits for stable graph: gui graph conversion, residual graph, empty flow filter

src/graph.rs

@@ -0,0 +1,80 @@
+use egui_graphs::{DefaultNodeShape, Graph, default_edge_transform, default_node_transform, to_graph_custom};
+use petgraph::{stable_graph::{StableGraph, NodeIndex}, Directed};
+use egui::Pos2;
+use crate::layout::CustomEdgeShape;
+use egui::Color32;
+
+pub trait ResidualGraph {
+    fn residual(&self) -> StableGraph<(f32, f32), (u64, u64)>;
+}
+
+impl ResidualGraph for StableGraph<(f32, f32), (u64, u64)> {
+    fn residual(&self) -> StableGraph<(f32, f32), (u64, u64)> {
+        let mut res_graph = self.clone();
+
+        for edge in res_graph.edge_indices().map(|e| e).collect::<Vec<_>>() {
+            let (source, target) = res_graph.edge_endpoints(edge).expect("edge not connected");
+            if let Some(reverse_edge) = res_graph.find_edge(target, source) {
+                let (flow, capacity) = *res_graph.edge_weight(edge).expect("forward edge not found");
+                let (reverse_flow, reverse_capacity) = *res_graph.edge_weight(reverse_edge).expect("reverse edge not found");
+                // update the forward edge
+                // TODO: this seems to overflow in certain cases
+                res_graph.update_edge(source, target, (flow, capacity + reverse_flow - flow));
+                // update the reverse edge
+                res_graph.update_edge(target, source, (reverse_flow, reverse_capacity + flow - reverse_flow));
+            } else {
+                // add a residual edge with a flow of 0 if the reverse edge doesn't exist
+                res_graph.add_edge(target, source, (0, 0));
+            }
+        }
+        res_graph
+    }
+}
+
+pub trait GuiGraph {
+    fn to_gui_graph(&self, s: NodeIndex, t: NodeIndex, active_flows_only: bool) -> Graph<(f32, f32), (u64, u64), Directed, u32, DefaultNodeShape, CustomEdgeShape>;
+}
+
+impl GuiGraph for StableGraph<(f32, f32), (u64, u64)> {
+    fn to_gui_graph(&self, s: NodeIndex, t: NodeIndex, active_flows_only: bool) -> Graph<(f32, f32), (u64, u64), Directed, u32, DefaultNodeShape, CustomEdgeShape> {
+        let mut graph = to_graph_custom(&self,
+            |n| {
+                let (x, y) = *n.payload();
+                default_node_transform(n);
+                n.set_location(Pos2::new(x, y));
+                n.set_label(String::from(""));
+            },
+            |e| {
+                let (flow, capacity): (u64, u64) = *e.payload();
+                default_edge_transform(e);
+                e.set_label(format!("{flow}:{capacity}"));
+                if flow > 0 {
+                    e.set_selected(true);
+                }
+            }
+        );
+        graph.node_mut(s).expect("node index not found").set_label(String::from("s"));
+        graph.node_mut(t).expect("node index not found").set_label(String::from("t"));
+        graph.node_mut(s).expect("node index not found").set_color(Color32::RED);
+        graph.node_mut(t).expect("node index not found").set_color(Color32::GREEN);
+        graph
+    }
+}
+
+pub trait FlowGraph {
+    fn filter_empty_flows(&self) -> StableGraph<(f32, f32), (u64, u64)>;
+}
+
+impl FlowGraph for StableGraph<(f32, f32), (u64, u64)> {
+    fn filter_empty_flows(&self) -> StableGraph<(f32, f32), (u64, u64)> {
+        self.filter_map( |_, n| {
+            Some(*n)
+        },
+        |_, &e| {
+            match e.0 {
+                0 => None,
+                _ => Some(e),
+            }
+        })
+    }
+}