From d55e5fdb68794e8587313b594b0f02ef86d0d71f Mon Sep 17 00:00:00 2001
From: chriseth <chris@ethereum.org>
Date: Wed, 31 Aug 2022 19:20:39 +0200
Subject: [PATCH] Flow

---
 src/flow/adjacencies.rs |  8 ++++----
 src/flow/mod.rs         |  7 +++++--
 src/main.rs             | 15 ++++++++++++++-
 src/types/address.rs    | 12 ++++++++++++
 src/types/u256.rs       | 23 ++++++++++++++++++++++-
 5 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/src/flow/adjacencies.rs b/src/flow/adjacencies.rs
index e0c6f43..682f744 100644
--- a/src/flow/adjacencies.rs
+++ b/src/flow/adjacencies.rs
@@ -4,7 +4,7 @@ use std::cmp::Reverse;
 use std::collections::HashMap;
 use std::collections::HashSet;
 
-struct Adjacencies<'a> {
+pub struct Adjacencies<'a> {
     edges: &'a HashMap<Address, Vec<Edge>>,
     lazy_adjacencies: HashMap<Node, HashMap<Node, U256>>,
     capacity_adjustments: HashMap<Node, HashMap<Node, U256>>,
@@ -42,12 +42,12 @@ impl<'a> Adjacencies<'a> {
         result
     }
 
-    pub fn adjust_capacity(&mut self, from: Node, to: Node, adjustment: U256) {
+    pub fn adjust_capacity(&mut self, from: &Node, to: &Node, adjustment: U256) {
         *self
             .capacity_adjustments
-            .entry(from)
+            .entry(from.clone())
             .or_default()
-            .entry(to)
+            .entry(to.clone())
             .or_default() += adjustment;
     }
 
diff --git a/src/flow/mod.rs b/src/flow/mod.rs
index e1e1903..7072613 100644
--- a/src/flow/mod.rs
+++ b/src/flow/mod.rs
@@ -1,9 +1,12 @@
-mod adjacencies;
+use crate::types::{Address};
 
-use crate::types::{Address, Edge, U256};
+mod adjacencies;
+mod flow;
 
 #[derive(Eq, PartialEq, Hash, Clone)]
 enum Node {
     Node(Address),
     TokenEdge(Address, Address),
 }
+
+pub use crate::flow::flow::compute_flow;
diff --git a/src/main.rs b/src/main.rs
index 8841ca2..0aa39cd 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,8 +1,21 @@
+use std::env;
+
+
 mod flow;
 mod io;
 mod types;
 
+use types::Address;
+
 fn main() {
-    let edges = io::read_edges_binary(&String::from("./edges.dat")).expect("Error loading edges.");
+    let args: Vec<String> = env::args().collect();
+    if args.len() != 4 {
+        panic!("Expected three arguments");
+    }
+    let (from_str, to_str, edges_file) = (&args[1], &args[2], &args[3]);
+
+    println!("Computing flow {from_str} -> {to_str} using {edges_file}");
+    let edges = io::read_edges_binary(edges_file).expect("Error loading edges.");
     print!("Read {} edges", edges.len());
+    flow::compute_flow(&Address::from(from_str.as_str()), &Address::from(to_str.as_str()), &edges);
 }
diff --git a/src/types/address.rs b/src/types/address.rs
index 8d9afbd..d68c42f 100644
--- a/src/types/address.rs
+++ b/src/types/address.rs
@@ -6,3 +6,15 @@ impl From<[u8; 20]> for Address {
         Address(item)
     }
 }
+
+impl From<&str> for Address {
+    fn from(item: &str) -> Self {
+        assert!(item.starts_with("0x"));
+        assert!(item.len() == 2 + 20 * 2);
+        let mut data = [0u8; 20];
+        for i in (2..item.len()).step_by(2) {
+            data[i / 2 - 1] = u8::from_str_radix(&item[i..i + 2], 16).unwrap();
+        }
+        Address(data)
+    }
+}
diff --git a/src/types/u256.rs b/src/types/u256.rs
index d8b1af7..662ea86 100644
--- a/src/types/u256.rs
+++ b/src/types/u256.rs
@@ -1,6 +1,6 @@
 use std::fmt::Display;
 use std::fmt::Formatter;
-use std::ops::{Add, AddAssign};
+use std::ops::{Add, AddAssign, Neg, Sub, SubAssign};
 
 #[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq, PartialOrd, Ord)]
 pub struct U256([u128; 2]);
@@ -70,6 +70,27 @@ impl AddAssign for U256 {
     }
 }
 
+impl Neg for U256 {
+    type Output = Self;
+    fn neg(self) -> Self {
+        let result = U256([!self.0[0], !self.0[1]]);
+        result + U256::from(1)
+    }
+}
+
+impl Sub for U256 {
+    type Output = Self;
+    fn sub(self, rhs: Self) -> Self {
+        self + (-rhs)
+    }
+}
+
+impl SubAssign for U256 {
+    fn sub_assign(&mut self, rhs: Self) {
+        *self = *self - rhs;
+    }
+}
+
 impl Display for U256 {
     fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
         if self.0[0] == 0 {