pathfinder2/src/types/edge.rs
2022-12-06 22:16:36 +01:00

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
}