checkpoint 2025-08-22

src/graph.rs

@@ -76,6 +76,8 @@ pub trait FlowGraph {
 
     fn reset_flow(&mut self);
 
+    fn prune_zero(&self) -> StableGraph<(f32, f32), (u64, u64)>;
+
     fn valid_flow(&self, source: NodeIndex, sink: NodeIndex) -> bool;
     
     fn saturated_cut(&self, source: NodeIndex, sink: NodeIndex) -> bool;
@@ -153,6 +155,19 @@ impl FlowGraph for StableGraph<(f32, f32), (u64, u64)> {
         }
     }
 
+    // removes all edges where the current flow is zero & returns a pruned graph
+    fn prune_zero(&self) -> StableGraph<(f32, f32), (u64, u64)> {
+        let mut zero_graph = self.clone();
+        for edge in zero_graph.edge_indices().collect::<Vec<_>>() {
+            let weight = zero_graph.edge_weight(edge).expect("edge not found");
+            if weight.0 == 0 {
+                zero_graph.remove_edge(edge);
+            }
+        }
+
+        zero_graph
+    }
+
     // check if flow is valid:
     // - for every edge, the flow doesn't exceed the capacity
     // - for every node, the amount of incoming flows is equal to the amount of outgoing flows
@@ -199,7 +214,6 @@ impl FlowGraph for StableGraph<(f32, f32), (u64, u64)> {
         return true;
     }
 
-    // TODO: assert that incoming flows at sink are equal to outgoing flows at source (otherwise return 0)
     fn total_flow(&self, source: NodeIndex, sink: NodeIndex) -> u64 {
         let outgoing_flows_source: u64 = self.edges_directed(source, Direction::Outgoing).map(|e| e.weight().0).sum();
         let incoming_flows_source: u64 = self.edges_directed(source, Direction::Incoming).map(|e| e.weight().0).sum();
@@ -207,7 +221,9 @@ impl FlowGraph for StableGraph<(f32, f32), (u64, u64)> {
         let incoming_flows_sink: u64 = self.edges_directed(sink, Direction::Incoming).map(|e| e.weight().0).sum();
         let total_flow_source = outgoing_flows_source - incoming_flows_source;
         let total_flow_sink = incoming_flows_sink - outgoing_flows_sink;
-        assert!(total_flow_sink == total_flow_source);
+        if total_flow_sink != total_flow_source {
+            return 0;
+        }
         return total_flow_source;
     }
 
@@ -233,22 +249,38 @@ impl FlowGraph for StableGraph<(f32, f32), (u64, u64)> {
         self.edge_weight(edge).expect("edge not found").1 - self.edge_weight(edge).expect("edge not found").0
     }
     
-    fn add_flow(&mut self, edge: EdgeIndex, flow: u64) {
-        // search for reverse edge and subtract existing flow if possible
-        let (source, target) = self.edge_endpoints(edge).expect("edge not connected");
-        let residual_edge = self.find_edge(target, source).expect("residual edge not found");
-        let residual_weight: &mut (u64, u64) = self.edge_weight_mut(residual_edge).expect("residual edge not found");
-        let forward_flow = i128::from(flow) - i128::from((*residual_weight).0);
-        // subtract minimum of additional_flow and residual_weight
-        if residual_weight.0 > 0 {
-            (*residual_weight).0 -= min(residual_weight.0, flow);
+    // adds the specified flow to the edge of a residual graph
+    fn add_flow(&mut self, forward_edge: EdgeIndex, flow: u64) {
+        // search for reverse edge
+        let (source, target) = self.edge_endpoints(forward_edge).expect("edge not connected");
+        let reverse_edge = self.find_edge(target, source).expect("reverse edge not found");
+        
+        // calculate how much flow needs to be added on the reverse edge
+        let current_reverse_flow = self.edge_weight(reverse_edge).expect("reverse erge not found").0;
+        // maximum flow we can subtract from reverse edge
+        let reverse_flow = min(current_reverse_flow, flow);
+        let forward_flow = flow - reverse_flow;
+        if reverse_flow > 0 {
+            // decrease reverse flow
+            let reverse_weight: &mut (u64, u64) = self.edge_weight_mut(reverse_edge).expect("reverse edge not found");
+            (*reverse_weight).0 -= reverse_flow;
+            assert!(reverse_weight.0 <= reverse_weight.1);
+            // decrease forward capacity 
+            let forward_weight: &mut (u64, u64) = self.edge_weight_mut(forward_edge).expect("edge not found");
+            (*forward_weight).1 -= reverse_flow;
+            assert!(forward_weight.0 <= forward_weight.1);
         }
-        assert!((*residual_weight).0 <= (*residual_weight).1);
+
         // add remaining flow to forward edge
-        let forward_weight: &mut (u64, u64) = self.edge_weight_mut(edge).expect("edge not found");
         if forward_flow > 0 {
-            (*forward_weight).0 += u64::try_from(forward_flow).expect("error casting flow to u64");
+            // increase forward flow
+            let forward_weight: &mut (u64, u64) = self.edge_weight_mut(forward_edge).expect("edge not found");
+            (*forward_weight).0 += forward_flow;
+            assert!(forward_weight.0 <= forward_weight.1);
+            // increase reverse capacity
+            let reverse_weight: &mut (u64, u64) = self.edge_weight_mut(reverse_edge).expect("reverse edge not found");
+            (*reverse_weight).1 += forward_flow;
+            assert!(reverse_weight.0 <= reverse_weight.1);
         }
-        assert!((*forward_weight).0 <= (*forward_weight).1);
     }
 }