diff --git a/src/flow/flow.rs b/src/flow/flow.rs index 1a77dd2..cd71773 100644 --- a/src/flow/flow.rs +++ b/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); - assert!(!next.is_empty()); - transfers.extend(next.into_iter()); + println!( + "Finding next transfers. Number of non-zero-balance accounts: {}", + 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 -/// we cannot transfer the full balance and start over. -fn extract_next_transfers( - used_edges: &mut HashMap>, - account_balances: &mut HashMap, -) -> Vec { - 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( +fn next_full_capacity_edge( used_edges: &HashMap>, account_balances: &HashMap, - first_edge: bool, -) -> 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 { - // We do not have enough balance yet, there will be another transfer along this edge. - if first_edge { - continue; - } else { - return None; - } - } else { - return Some(Edge { + if *balance >= *capacity { + println!("Found an edge: {from} -> {to} [{token}] {capacity}"); + return Edge { from: *from, to: *to, token: *token, capacity: *capacity, - }); + }; } } } } - None + panic!(); }