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])
}