From d55e5fdb68794e8587313b594b0f02ef86d0d71f Mon Sep 17 00:00:00 2001 From: chriseth 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>, 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 {