Adjacencies

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;
 
+struct Adjacencies<'a> {
+    edges: &'a HashMap<Address, Vec<Edge>>,
+    lazy_adjacencies: HashMap<Node, HashMap<Node, U256>>,
+    capacity_adjustments: HashMap<Node, HashMap<Node, U256>>,
+}
+
+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<Address, Vec<Edge>>) -> 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::<Vec<(Node, U256)>>();
+        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<Node, U256> {
+        self.lazy_adjacencies
+            .entry(from.clone())
+            .or_insert_with(|| {
+                let mut result: HashMap<Node, U256> = 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()
+    }
 }

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

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<HashSet<Edge>, io::Error> {
+pub fn read_edges_binary(path: &String) -> Result<HashMap<Address, Vec<Edge>>, 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<U256, io::Error> {
 fn read_edges(
     file: &mut File,
     address_index: &HashMap<u32, Address>,
-) -> Result<HashSet<Edge>, io::Error> {
+) -> Result<HashMap<Address, Vec<Edge>>, io::Error> {
     let edge_count = read_u32(file)?;
-    let mut edges = HashSet::new();
+    let mut edges: HashMap<Address, Vec<Edge>> = 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,

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 {

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