diff --git a/src/flow/adjacencies.rs b/src/flow/adjacencies.rs index 4f914f3..e0c6f43 100644 --- a/src/flow/adjacencies.rs +++ b/src/flow/adjacencies.rs @@ -1,3 +1,95 @@ -struct Adjacencies { +use crate::flow::Node; +use crate::types::{Address, Edge, U256}; +use std::cmp::Reverse; +use std::collections::HashMap; +use std::collections::HashSet; -} \ No newline at end of file +struct Adjacencies<'a> { + edges: &'a HashMap>, + lazy_adjacencies: HashMap>, + capacity_adjustments: HashMap>, +} + +fn pseudo_node(edge: Edge) -> Node { + Node::TokenEdge(edge.from, edge.token) +} + +fn source_address_of(node: &Node) -> &Address { + match node { + Node::Node(addr) => addr, + Node::TokenEdge(from, _) => from, + } +} + +impl<'a> Adjacencies<'a> { + pub fn new(edges: &'a HashMap>) -> Self { + Adjacencies { + edges, + lazy_adjacencies: HashMap::new(), + capacity_adjustments: HashMap::new(), + } + } + + pub fn outgoing_edges_sorted_by_capacity(&mut self, from: &Node) -> Vec<(Node, U256)> { + let mut adjacencies = self.adjacencies_from(from); + if let Some(adjustments) = self.capacity_adjustments.get(from) { + for (node, c) in adjustments { + *adjacencies.entry(node.clone()).or_default() += *c; + } + } + let mut result = adjacencies.into_iter().collect::>(); + result.sort_unstable_by_key(|(_, capacity)| Reverse(*capacity)); + result + } + + pub fn adjust_capacity(&mut self, from: Node, to: Node, adjustment: U256) { + *self + .capacity_adjustments + .entry(from) + .or_default() + .entry(to) + .or_default() += adjustment; + } + + pub fn is_adjacent(&mut self, from: &Node, to: &Node) -> bool { + // TODO More efficiently? + if let Some(capacity) = self.adjacencies_from(from).get(to) { + *capacity > U256::from(0) + } else { + false + } + } + + fn adjacencies_from(&mut self, from: &Node) -> HashMap { + self.lazy_adjacencies + .entry(from.clone()) + .or_insert_with(|| { + let mut result: HashMap = HashMap::new(); + for edge in &self.edges[source_address_of(from)] { + match from { + Node::Node(_) => { + // One edge from "from" to "from x token" with a capacity + // as the max over all contributing edges (the balance of the sender) + result + .entry(pseudo_node(*edge)) + .and_modify(|c| { + if edge.capacity > *c { + *c = edge.capacity; + } + }) + .or_insert(edge.capacity); + } + Node::TokenEdge(_, _) => { + // Another edge from "from x token" to "to" with its + // own capacity (based on the trust) + if pseudo_node(*edge) == *from { + result.insert(Node::Node(edge.to), edge.capacity); + } + } + } + } + result + }) + .clone() + } +} diff --git a/src/flow/mod.rs b/src/flow/mod.rs index c76a6ca..e1e1903 100644 --- a/src/flow/mod.rs +++ b/src/flow/mod.rs @@ -1,12 +1,9 @@ -use crate::address::Address; -use crate::edge::Edge; -use crate::u256::U256; +mod adjacencies; +use crate::types::{Address, Edge, U256}; + +#[derive(Eq, PartialEq, Hash, Clone)] enum Node { Node(Address), TokenEdge(Address, Address), } - -fn pseudoNode(edge: Edge) -> Node { - Node::TokenEdge(edge.from, edge.to) -} \ No newline at end of file diff --git a/src/io.rs b/src/io.rs index 639d18c..a76c720 100644 --- a/src/io.rs +++ b/src/io.rs @@ -6,7 +6,7 @@ use std::io::Read; use crate::types::{Address, Edge, U256}; -pub fn read_edges_binary(path: &String) -> Result, io::Error> { +pub fn read_edges_binary(path: &String) -> Result>, io::Error> { let mut f = File::open(path)?; let address_index = read_address_index(&mut f)?; read_edges(&mut f, &address_index) @@ -55,15 +55,15 @@ fn read_u256(file: &mut File) -> Result { fn read_edges( file: &mut File, address_index: &HashMap, -) -> Result, io::Error> { +) -> Result>, io::Error> { let edge_count = read_u32(file)?; - let mut edges = HashSet::new(); + let mut edges: HashMap> = HashMap::new(); for _i in 0..edge_count { let from = read_address(file, address_index)?; let to = read_address(file, address_index)?; let token = read_address(file, address_index)?; let capacity = read_u256(file)?; - edges.insert(Edge { + edges.entry(from).or_insert(vec![]).push(Edge { from, to, token, diff --git a/src/types/address.rs b/src/types/address.rs index 2a13a65..8d9afbd 100644 --- a/src/types/address.rs +++ b/src/types/address.rs @@ -1,4 +1,4 @@ -#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, PartialOrd)] +#[derive(Clone, Copy, Debug, Default, Hash, Eq, PartialEq, PartialOrd)] pub struct Address([u8; 20]); impl From<[u8; 20]> for Address { diff --git a/src/types/u256.rs b/src/types/u256.rs index 576a16f..d8b1af7 100644 --- a/src/types/u256.rs +++ b/src/types/u256.rs @@ -2,7 +2,7 @@ use std::fmt::Display; use std::fmt::Formatter; use std::ops::{Add, AddAssign}; -#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, PartialOrd)] +#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq, PartialOrd, Ord)] pub struct U256([u128; 2]); impl U256 { @@ -56,11 +56,9 @@ impl Add for U256 { type Output = Self; fn add(self, rhs: Self) -> Self { let (low, carry) = self.0[1].overflowing_add(rhs.0[1]); - let (mut high, mut overflow) = self.0[0].overflowing_add(rhs.0[0]); - assert!(!overflow); + let mut high = self.0[0].wrapping_add(rhs.0[0]); if carry { - (high, overflow) = high.overflowing_add(1); - assert!(!overflow); + high = high.wrapping_add(1); } Self([high, low]) }