From 12bb20d58da77466c7bb53c9d66c9fde86490dea Mon Sep 17 00:00:00 2001
From: chriseth <chris@ethereum.org>
Date: Tue, 30 Aug 2022 18:53:34 +0200
Subject: [PATCH] Loading edge database.

---
 src/address.rs |  8 ++++++
 src/edge.rs    | 15 ++++++++++
 src/io.rs      | 76 ++++++++++++++++++++++++++++++++++++++++++++++++++
 src/main.rs    |  6 +++-
 src/u256.rs    | 21 ++++++++++----
 5 files changed, 120 insertions(+), 6 deletions(-)
 create mode 100644 src/address.rs
 create mode 100644 src/edge.rs
 create mode 100644 src/io.rs

diff --git a/src/address.rs b/src/address.rs
new file mode 100644
index 0000000..2a13a65
--- /dev/null
+++ b/src/address.rs
@@ -0,0 +1,8 @@
+#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, PartialOrd)]
+pub struct Address([u8; 20]);
+
+impl From<[u8; 20]> for Address {
+    fn from(item: [u8; 20]) -> Self {
+        Address(item)
+    }
+}
diff --git a/src/edge.rs b/src/edge.rs
new file mode 100644
index 0000000..498e0c7
--- /dev/null
+++ b/src/edge.rs
@@ -0,0 +1,15 @@
+use crate::address::Address;
+use crate::u256::U256;
+
+#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, PartialOrd)]
+pub struct Edge {
+    pub from: Address,
+    pub to: Address,
+    pub token: Address,
+    pub capacity: U256,
+}
+
+// TODO comparison, hash, etc. can ignore the capacity field.
+
+// TODO can we derive it?
+impl Eq for Edge {}
diff --git a/src/io.rs b/src/io.rs
new file mode 100644
index 0000000..5cec6f2
--- /dev/null
+++ b/src/io.rs
@@ -0,0 +1,76 @@
+use std::collections::HashMap;
+use std::collections::HashSet;
+use std::fs::File;
+use std::io;
+use std::io::Read;
+
+use crate::address::Address;
+use crate::edge::Edge;
+use crate::u256::U256;
+
+pub fn read_edges_binary(path: &String) -> Result<HashSet<Edge>, io::Error> {
+    let mut f = File::open(path)?;
+    let address_index = read_address_index(&mut f)?;
+    read_edges(&mut f, &address_index)
+}
+
+fn read_address_index(file: &mut File) -> Result<HashMap<u32, Address>, io::Error> {
+    let address_count = read_u32(file)?;
+    let mut addresses = HashMap::new();
+    for i in 0..address_count {
+        let mut buf = [0; 20];
+        file.read_exact(&mut buf)?;
+        addresses.insert(i, Address::from(buf));
+    }
+    Ok(addresses)
+}
+
+fn read_u32(file: &mut File) -> Result<u32, io::Error> {
+    let mut buf = [0; 4];
+    file.read_exact(&mut buf)?;
+    Ok(u32::from_be_bytes(buf))
+}
+
+fn read_u8(file: &mut File) -> Result<u8, io::Error> {
+    let mut buf = [0; 1];
+    file.read_exact(&mut buf)?;
+    Ok(u8::from_be_bytes(buf))
+}
+
+fn read_address(
+    file: &mut File,
+    address_index: &HashMap<u32, Address>,
+) -> Result<Address, io::Error> {
+    let index = read_u32(file)?;
+    Ok(address_index[&index])
+}
+
+fn read_u256(file: &mut File) -> Result<U256, io::Error> {
+    let length = read_u8(file)? as usize;
+    let mut bytes = [0u8; 32];
+    file.read_exact(&mut bytes[32 - length..32])?;
+    let high = u128::from_be_bytes(*<&[u8; 16]>::try_from(&bytes[0..16]).unwrap());
+    let low = u128::from_be_bytes(*<&[u8; 16]>::try_from(&bytes[0..16]).unwrap());
+    Ok(U256::new(high, low))
+}
+
+fn read_edges(
+    file: &mut File,
+    address_index: &HashMap<u32, Address>,
+) -> Result<HashSet<Edge>, io::Error> {
+    let edge_count = read_u32(file)?;
+    let mut edges = HashSet::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 {
+            from,
+            to,
+            token,
+            capacity,
+        });
+    }
+    Ok(edges)
+}
diff --git a/src/main.rs b/src/main.rs
index c392296..322fe9f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,5 +1,9 @@
+mod address;
+mod edge;
+mod io;
 mod u256;
 
 fn main() {
-    println!("Hello, world!");
+    let edges = io::read_edges_binary(&String::from("./edges.dat")).expect("Error loading edges.");
+    print!("Read {} edges", edges.len());
 }
diff --git a/src/u256.rs b/src/u256.rs
index 62bac9d..576a16f 100644
--- a/src/u256.rs
+++ b/src/u256.rs
@@ -5,7 +5,11 @@ use std::ops::{Add, AddAssign};
 #[derive(Clone, Copy, Debug, Default, Hash, PartialEq, PartialOrd)]
 pub struct U256([u128; 2]);
 
-impl U256 {}
+impl U256 {
+    pub fn new(high: u128, low: u128) -> U256 {
+        U256([high, low])
+    }
+}
 
 impl From<u128> for U256 {
     fn from(item: u128) -> Self {
@@ -25,14 +29,22 @@ impl From<&str> for U256 {
             assert!(
                 low_hex.as_bytes().get(0) != Some(&54) && low_hex.as_bytes().get(0) != Some(&43)
             );
-            let low = if low_hex.is_empty() { 0 } else {u128::from_str_radix(low_hex, 16).unwrap()};
+            let low = if low_hex.is_empty() {
+                0
+            } else {
+                u128::from_str_radix(low_hex, 16).unwrap()
+            };
             let high_start = if len >= 2 + 32 + 32 { len - 64 } else { 2 };
             let high_hex = &item[high_start..low_start];
             // disallow + and - prefixes
             assert!(
                 high_hex.as_bytes().get(0) != Some(&54) && high_hex.as_bytes().get(0) != Some(&43)
             );
-            let high = if high_hex.is_empty() { 0 } else { u128::from_str_radix(high_hex, 16).unwrap() };
+            let high = if high_hex.is_empty() {
+                0
+            } else {
+                u128::from_str_radix(high_hex, 16).unwrap()
+            };
             U256([high, low])
         } else {
             todo!("Decimal import");
@@ -71,7 +83,7 @@ impl Display for U256 {
 }
 
 mod test {
-    use super::*;
+    use super::U256;
     #[test]
     fn to_string() {
         assert_eq!(format!("{}", U256::from(0)), "0x0");
@@ -116,5 +128,4 @@ mod test {
             U256::from(u128::MAX) + U256::from(1)
         );
     }
-
 }