Try new way to do it.

src/flow/flow.rs

@@ -46,7 +46,11 @@ pub fn compute_flow(
     // TODO prune
 
     println!("Max flow: {flow}");
-    let transfers = extract_transfers(source, sink, &flow, used_edges);
+    let transfers = if flow == U256::from(0) {
+        vec![]
+    } else {
+        extract_transfers(source, sink, &flow, used_edges)
+    };
     println!("Num transfers: {}", transfers.len());
     flow.to_string()
 }
@@ -106,74 +110,57 @@ fn extract_transfers(
     while !account_balances.is_empty()
         && (account_balances.len() > 1 || *account_balances.iter().nth(0).unwrap().0 != *sink)
     {
-        let next = extract_next_transfers(&mut used_edges, &mut account_balances);
+        println!(
-        assert!(!next.is_empty());
+            "Finding next transfers. Number of non-zero-balance accounts: {}",
-        transfers.extend(next.into_iter());
+            account_balances.len()
+        );
+        let edge = next_full_capacity_edge(&mut used_edges, &mut account_balances);
+        account_balances
+            .entry(edge.from)
+            .and_modify(|balance| *balance -= edge.capacity);
+        account_balances
+            .entry(edge.to)
+            .and_modify(|balance| *balance += edge.capacity);
+        account_balances.retain(|_account, balance| balance > &mut U256::from(0));
+        used_edges
+            .entry(Node::Node(edge.from))
+            .and_modify(|outgoing| {
+                outgoing.remove(&Node::TokenEdge(edge.from, edge.token));
+            });
+        transfers.push(edge);
     }
 
     transfers
 }
 
-/// Extract the next list of transfers until we get to a situation where
+fn next_full_capacity_edge(
-/// we cannot transfer the full balance and start over.
-fn extract_next_transfers(
-    used_edges: &mut HashMap<Node, HashMap<Node, U256>>,
-    account_balances: &mut HashMap<Address, U256>,
-) -> Vec<Edge> {
-    let mut transfers = Vec::new();
-
-    loop {
-        let first_edge = transfers.is_empty();
-        if let Some(edge) = next_nonzero_edge(used_edges, account_balances, first_edge) {
-            account_balances
-                .entry(edge.from)
-                .and_modify(|balance| *balance -= edge.capacity);
-            account_balances
-                .entry(edge.to)
-                .and_modify(|balance| *balance += edge.capacity);
-            account_balances.retain(|_account, balance| balance > &mut U256::from(0));
-            used_edges
-                .entry(Node::Node(edge.from))
-                .and_modify(|outgoing| {
-                    outgoing.remove(&Node::TokenEdge(edge.from, edge.token));
-                });
-            transfers.push(edge);
-        } else {
-            return transfers;
-        }
-    }
-}
-
-fn next_nonzero_edge(
     used_edges: &HashMap<Node, HashMap<Node, U256>>,
     account_balances: &HashMap<Address, U256>,
-    first_edge: bool,
+) -> Edge {
-) -> Option<Edge> {
     for (account, balance) in account_balances {
-        for (intermediate, _) in &used_edges[&Node::Node(*account)] {
+        println!("Account: {account} - balance: {balance}");
+        for (intermediate, _) in used_edges
+            .get(&Node::Node(*account))
+            .unwrap_or(&HashMap::default())
+        {
             let (from, token) = node_as_token_edge(intermediate);
             for (to_node, capacity) in &used_edges[intermediate] {
+                println!(" - used edge to {to_node} with capacity {capacity}");
                 let to = node_as_address(to_node);
                 if *capacity == U256::from(0) {
                     continue;
                 }
-                if *balance < *capacity {
+                if *balance >= *capacity {
-                    // We do not have enough balance yet, there will be another transfer along this edge.
+                    println!("Found an edge: {from} -> {to} [{token}] {capacity}");
-                    if first_edge {
+                    return Edge {
-                        continue;
-                    } else {
-                        return None;
-                    }
-                } else {
-                    return Some(Edge {
                         from: *from,
                         to: *to,
                         token: *token,
                         capacity: *capacity,
-                    });
+                    };
                 }
             }
         }
     }
-    None
+    panic!();
 }