92 lines
2.4 KiB
Rust
92 lines
2.4 KiB
Rust
use std::collections::HashMap;
|
|
|
|
use crate::types::Address;
|
|
use crate::types::U256;
|
|
|
|
#[derive(Clone, Copy, Debug, Default, Hash, PartialEq, Eq, Ord, 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.
|
|
|
|
pub fn eq_up_to_capacity(e1: &Edge, e2: &Edge) -> bool {
|
|
e1.from == e2.from && e1.to == e2.to && e1.token == e2.token
|
|
}
|
|
|
|
#[derive(Debug, Default, Clone)]
|
|
pub struct EdgeDB {
|
|
edges: Vec<Edge>,
|
|
outgoing: HashMap<Address, Vec<usize>>,
|
|
incoming: HashMap<Address, Vec<usize>>,
|
|
}
|
|
|
|
impl EdgeDB {
|
|
pub fn new(edges: Vec<Edge>) -> EdgeDB {
|
|
let outgoing = outgoing_index(&edges);
|
|
let incoming = incoming_index(&edges);
|
|
EdgeDB {
|
|
edges,
|
|
outgoing,
|
|
incoming,
|
|
}
|
|
}
|
|
|
|
pub fn edge_count(&self) -> usize {
|
|
self.edges.len()
|
|
}
|
|
|
|
pub fn update(&mut self, update: Edge) {
|
|
match self.index_of(&update) {
|
|
Some(i) => self.edges[i].capacity = update.capacity,
|
|
None => {
|
|
let i = self.edges.len();
|
|
self.outgoing.entry(update.from).or_default().push(i);
|
|
self.incoming.entry(update.to).or_default().push(i);
|
|
self.edges.push(update);
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn outgoing(&self, source: &Address) -> Vec<&Edge> {
|
|
match self.outgoing.get(source) {
|
|
Some(out) => out
|
|
.iter()
|
|
.map(|i| self.edges.get(*i).unwrap())
|
|
.filter(|e| e.capacity != U256::from(0))
|
|
.collect(),
|
|
None => vec![],
|
|
}
|
|
}
|
|
|
|
fn index_of(&self, e: &Edge) -> Option<usize> {
|
|
self.outgoing.get(&e.from).and_then(|out| {
|
|
for i in out {
|
|
if eq_up_to_capacity(&self.edges[*i], e) {
|
|
return Some(*i);
|
|
}
|
|
}
|
|
None
|
|
})
|
|
}
|
|
}
|
|
|
|
fn outgoing_index(edges: &[Edge]) -> HashMap<Address, Vec<usize>> {
|
|
let mut index: HashMap<Address, Vec<usize>> = HashMap::new();
|
|
for (i, e) in edges.iter().enumerate() {
|
|
index.entry(e.from).or_default().push(i)
|
|
}
|
|
index
|
|
}
|
|
|
|
fn incoming_index(edges: &[Edge]) -> HashMap<Address, Vec<usize>> {
|
|
let mut index: HashMap<Address, Vec<usize>> = HashMap::new();
|
|
for (i, e) in edges.iter().enumerate() {
|
|
index.entry(e.to).or_default().push(i)
|
|
}
|
|
index
|
|
}
|