Binary safe database import.

src/bin/convert.rs

@@ -1,61 +1,60 @@
 use std::env;
 
-use pathfinder2::{
-    io::{read_edges_binary, read_edges_csv, write_edges_binary, write_edges_csv},
-    safe_db::safes_json::import_from_safes_json,
-};
+use pathfinder2::io::*;
+use pathfinder2::safe_db::safes_json::import_from_safes_json;
 
 fn main() {
-    let operation = env::args().nth(1).and_then(|op| {
+    let input_format = env::args().nth(1).and_then(|op| {
         if matches!(
             op.as_str(),
-            "--safes-json-to-edges-bin"
-                | "--safes-json-to-edges-csv"
-                | "--edges-csv-to-edges-bin"
-                | "--edges-bin-to-edges-csv"
+            "--safes-json" | "--safes-bin" | "--edges-csv" | "--edges-bin"
         ) {
             Some(op)
         } else {
             None
         }
     });
-    if env::args().len() != 4 || operation.is_none() {
-        println!("Usage: convert --safes-json-to-edges-bin <safes.json> <edges.dat>");
-        println!("Usage: convert --safes-json-to-edges-csv <safes.json> <edges.csv>");
-        println!("Usage: convert --edges-csv-to-edges-bin <edges.csv> <edges.dat>");
-        println!("Usage: convert --edges-bin-to-edges-csv <edges.dat> <edges.csv>");
+    let output_format = env::args().nth(3).and_then(|op| {
+        if matches!(op.as_str(), "--edges-csv" | "--edges-bin") {
+            Some(op)
+        } else {
+            None
+        }
+    });
+    if env::args().len() != 5 || input_format.is_none() || output_format.is_none() {
+        println!("Usage: convert <input> <input_file> <output> <output_file>");
+        println!("  Where <input> is one of:");
+        println!("    --safes-json");
+        println!("    --safes-bin");
+        println!("    --edges-csv");
+        println!("    --edges-bin");
+        println!("  and <output>is one of:");
+        println!("    --edges-csv");
+        println!("    --edges-bin");
         return;
     }
 
-    let input = env::args().nth(2).unwrap();
-    let output = env::args().nth(3).unwrap();
-    match operation.unwrap().as_str() {
-        "--safes-json-to-edges-bin" => {
-            let safes = import_from_safes_json(&input);
-            let edges = safes.edges();
-            println!("Imported {} edges.", edges.edge_count());
-            write_edges_binary(edges, &output).unwrap();
-            println!("Export done.");
+    let input_file = env::args().nth(2).unwrap();
+    let edges = match input_format.unwrap().as_str() {
+        "--safes-json" => {
+            let safes = import_from_safes_json(&input_file);
+            safes.edges().clone()
         }
-        "--safes-json-to-edges-csv" => {
-            let safes = import_from_safes_json(&input);
-            let edges = safes.edges();
-            println!("Imported {} edges.", edges.edge_count());
-            write_edges_csv(edges, &output).unwrap();
-            println!("Export done.");
-        }
-        "--edges-csv-to-edges-bin" => {
-            let edges = &read_edges_csv(&input).unwrap();
-            println!("Imported {} edges.", edges.edge_count());
-            write_edges_binary(edges, &output).unwrap();
-            println!("Export done.");
-        }
-        "--edges-bin-to-edges-csv" => {
-            let edges = &read_edges_binary(&input).unwrap();
-            println!("Imported {} edges.", edges.edge_count());
-            write_edges_csv(edges, &output).unwrap();
-            println!("Export done.");
+        "--safes-bin" => {
+            let safes = import_from_safes_binary(&input_file).unwrap();
+            safes.edges().clone()
         }
+        "--edges-csv" => read_edges_csv(&input_file).unwrap(),
+        "--edges-bin" => read_edges_binary(&input_file).unwrap(),
+        _ => unreachable!(),
+    };
+    println!("Imported {} edges.", edges.edge_count());
+
+    let output_file = env::args().nth(4).unwrap();
+    match output_format.unwrap().as_str() {
+        "--edges-csv" => write_edges_csv(&edges, &output_file).unwrap(),
+        "--edges-bin" => write_edges_binary(&edges, &output_file).unwrap(),
         _ => unreachable!(),
     }
+    println!("Export done.");
 }

src/io.rs

@@ -1,11 +1,12 @@
-use std::collections::BTreeSet;
+use std::collections::{BTreeMap, BTreeSet};
 use std::fs::File;
 use std::io::{self, BufRead};
 use std::io::{Read, Write};
 use std::{collections::HashMap, io::BufReader};
 
+use crate::safe_db::db::DB;
 use crate::types::edge::EdgeDB;
-use crate::types::{Address, Edge, U256};
+use crate::types::{Address, Edge, Safe, U256};
 
 pub fn read_edges_binary(path: &String) -> Result<EdgeDB, io::Error> {
     let mut f = File::open(path)?;
@@ -65,6 +66,63 @@ pub fn write_edges_csv(edges: &EdgeDB, path: &String) -> Result<(), io::Error> {
     Ok(())
 }
 
+pub fn import_from_safes_binary(path: &str) -> Result<DB, io::Error> {
+    let mut f = File::open(path)?;
+
+    let mut safes: BTreeMap<Address, Safe> = Default::default();
+
+    let address_index = read_address_index(&mut f)?;
+
+    // organizations
+    for _ in 0..read_u32(&mut f)? {
+        let org_address = read_address(&mut f, &address_index)?;
+        safes.entry(org_address).or_default().organization = true;
+    }
+
+    // trust edges
+    for _ in 0..read_u32(&mut f)? {
+        let user = read_address(&mut f, &address_index)?;
+        assert!(user != Address::default());
+        let send_to = read_address(&mut f, &address_index)?;
+        assert!(send_to != Address::default());
+        let limit_percentage = read_u8(&mut f)?;
+        assert!(limit_percentage <= 100);
+
+        if send_to != user && limit_percentage > 0 {
+            safes
+                .entry(user)
+                .or_default()
+                .limit_percentage
+                .insert(send_to, limit_percentage);
+        }
+    }
+
+    // balances
+    for _ in 0..read_u32(&mut f)? {
+        let user = read_address(&mut f, &address_index)?;
+        assert!(user != Address::default());
+        let token_owner = read_address(&mut f, &address_index)?;
+        assert!(token_owner != Address::default());
+        let balance = read_u256(&mut f)?;
+        if balance != U256::from(0) {
+            safes
+                .entry(user)
+                .or_default()
+                .balances
+                .insert(token_owner, balance);
+        }
+    }
+
+    // we use the safe address as token address
+    let mut token_owner = BTreeMap::default();
+    for (addr, safe) in &mut safes {
+        safe.token_address = *addr;
+        token_owner.insert(*addr, *addr);
+    }
+
+    Ok(DB::new(safes, token_owner))
+}
+
 fn read_address_index(file: &mut File) -> Result<HashMap<u32, Address>, io::Error> {
     let address_count = read_u32(file)?;
     let mut addresses = HashMap::new();

src/types/safe.rs

@@ -4,6 +4,8 @@ use super::{Address, U256};
 
 #[derive(Default, Debug)]
 pub struct Safe {
+    /// The address of the token, or the address of the safe if
+    /// the database does not use the distinction.
     pub token_address: Address,
     pub balances: BTreeMap<Address, U256>,
     /// Limit percentage in "send to" direction