diff --git a/src/bin/convert.rs b/src/bin/convert.rs index d9c3afc..093e36b 100644 --- a/src/bin/convert.rs +++ b/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 "); - println!("Usage: convert --safes-json-to-edges-csv "); - println!("Usage: convert --edges-csv-to-edges-bin "); - println!("Usage: convert --edges-bin-to-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 "); + println!(" Where is one of:"); + println!(" --safes-json"); + println!(" --safes-bin"); + println!(" --edges-csv"); + println!(" --edges-bin"); + println!(" and 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."); } diff --git a/src/io.rs b/src/io.rs index 4d58aca..8969d21 100644 --- a/src/io.rs +++ b/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 { 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 { + let mut f = File::open(path)?; + + let mut safes: BTreeMap = 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, io::Error> { let address_count = read_u32(file)?; let mut addresses = HashMap::new(); diff --git a/src/types/safe.rs b/src/types/safe.rs index 295b95d..967065b 100644 --- a/src/types/safe.rs +++ b/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, /// Limit percentage in "send to" direction