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
>,
lazy_adjacencies: HashMap>,
capacity_adjustments: HashMap>,
@@ -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 = 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 {