From 03f0ff8d382f746e7b22cd8266cf8fc834038276 Mon Sep 17 00:00:00 2001
From: chriseth <chris@ethereum.org>
Date: Fri, 28 Oct 2022 22:39:47 +0200
Subject: [PATCH] Support importing csv files.

---
 src/io.rs     | 44 ++++++++++++++++++++++++++++++++++++++++++--
 src/server.rs | 21 ++++++++++++++++++++-
 2 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/src/io.rs b/src/io.rs
index b7f860d..1b3e7f3 100644
--- a/src/io.rs
+++ b/src/io.rs
@@ -1,7 +1,7 @@
-use std::collections::HashMap;
 use std::fs::File;
-use std::io;
 use std::io::Read;
+use std::io::{self, BufRead};
+use std::{collections::HashMap, io::BufReader};
 
 use crate::types::{Address, Edge, U256};
 
@@ -11,6 +11,36 @@ pub fn read_edges_binary(path: &String) -> Result<HashMap<Address, Vec<Edge>>, i
     read_edges(&mut f, &address_index)
 }
 
+pub fn read_edges_csv(path: &String) -> Result<HashMap<Address, Vec<Edge>>, io::Error> {
+    let mut result = HashMap::<Address, Vec<Edge>>::new();
+    let f = BufReader::new(File::open(path)?);
+    for line in f.lines() {
+        let line = line?;
+        match &line.split(',').collect::<Vec<_>>()[..] {
+            [] => continue,
+            [from, to, token, capacity] => {
+                let from = Address::from(unescape(from));
+                let to = Address::from(unescape(to));
+                let token = Address::from(unescape(token));
+                let capacity = U256::from(unescape(capacity));
+                result.entry(from).or_default().push(Edge {
+                    from,
+                    to,
+                    token,
+                    capacity,
+                });
+            }
+            _ => {
+                return Result::Err(io::Error::new(
+                    io::ErrorKind::Other,
+                    format!("Expected from,to,token,capacity, but got {line}"),
+                ))
+            }
+        }
+    }
+    Ok(result)
+}
+
 fn read_address_index(file: &mut File) -> Result<HashMap<u32, Address>, io::Error> {
     let address_count = read_u32(file)?;
     let mut addresses = HashMap::new();
@@ -71,3 +101,13 @@ fn read_edges(
     }
     Ok(edges)
 }
+
+fn unescape(input: &str) -> &str {
+    match input.chars().next() {
+        Some('"') | Some('\'') => {
+            assert!(input.len() >= 2 && input.chars().last() == input.chars().next());
+            &input[1..input.len() - 1]
+        }
+        _ => input,
+    }
+}
diff --git a/src/server.rs b/src/server.rs
index 158f717..1cb1192 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -1,5 +1,5 @@
 use crate::graph;
-use crate::io::read_edges_binary;
+use crate::io::{read_edges_binary, read_edges_csv};
 use crate::types::{Address, Edge, U256};
 use json::JsonValue;
 use std::collections::HashMap;
@@ -70,6 +70,15 @@ fn handle_connection(
             };
             socket.write_all(response.as_bytes())?;
         }
+        "load_edges_csv" => {
+            let response = match load_edges_csv(edges, &request.params["file"].to_string()) {
+                Ok(len) => jsonrpc_response(request.id, len),
+                Err(e) => {
+                    jsonrpc_error_response(request.id, -32000, &format!("Error loading edges: {e}"))
+                }
+            };
+            socket.write_all(response.as_bytes())?;
+        }
         "compute_transfer" => {
             println!("Computing flow");
             let e = edges.read().unwrap().clone();
@@ -91,6 +100,16 @@ fn load_edges_binary(
     Ok(len)
 }
 
+fn load_edges_csv(
+    edges: &RwLock<Arc<HashMap<Address, Vec<Edge>>>>,
+    file: &String,
+) -> Result<usize, Box<dyn Error>> {
+    let updated_edges = read_edges_csv(file)?;
+    let len = updated_edges.len();
+    *edges.write().unwrap() = Arc::new(updated_edges);
+    Ok(len)
+}
+
 fn compute_transfer(
     request: JsonRpcRequest,
     edges: &HashMap<Address, Vec<Edge>>,