use crate::core::id::{Eid, Vid};
use fxhash::FxBuildHasher;
use std::collections::HashMap;
#[derive(Clone, Copy, Debug)]
pub struct EdgeEntry {
pub eid: Eid,
pub src_vid: Vid,
pub dst_vid: Vid,
pub edge_type: u32,
}
type FxHashMap<K, V> = HashMap<K, V, FxBuildHasher>;
#[derive(Clone, Debug)]
pub struct SimpleGraph {
vertices: FxHashMap<Vid, ()>,
outgoing: FxHashMap<Vid, Vec<EdgeEntry>>,
incoming: FxHashMap<Vid, Vec<EdgeEntry>>,
edge_map: FxHashMap<Eid, EdgeEntry>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Direction {
Outgoing,
Incoming,
}
impl Default for SimpleGraph {
fn default() -> Self {
Self {
vertices: HashMap::with_hasher(FxBuildHasher::default()),
outgoing: HashMap::with_hasher(FxBuildHasher::default()),
incoming: HashMap::with_hasher(FxBuildHasher::default()),
edge_map: HashMap::with_hasher(FxBuildHasher::default()),
}
}
}
impl SimpleGraph {
pub fn new() -> Self {
Self::default()
}
pub fn with_capacity(vertices: usize, edges: usize) -> Self {
Self {
vertices: HashMap::with_capacity_and_hasher(vertices, FxBuildHasher::default()),
outgoing: HashMap::with_capacity_and_hasher(vertices, FxBuildHasher::default()),
incoming: HashMap::with_capacity_and_hasher(vertices, FxBuildHasher::default()),
edge_map: HashMap::with_capacity_and_hasher(edges, FxBuildHasher::default()),
}
}
pub fn add_vertex(&mut self, vid: Vid) -> bool {
if self.vertices.contains_key(&vid) {
return false;
}
self.vertices.insert(vid, ());
true
}
pub fn remove_vertex(&mut self, vid: Vid) {
self.vertices.remove(&vid);
if let Some(edges) = self.outgoing.remove(&vid) {
for edge in edges {
self.edge_map.remove(&edge.eid);
if let Some(incoming_list) = self.incoming.get_mut(&edge.dst_vid) {
incoming_list.retain(|e| e.eid != edge.eid);
}
}
}
if let Some(edges) = self.incoming.remove(&vid) {
for edge in edges {
self.edge_map.remove(&edge.eid);
if let Some(outgoing_list) = self.outgoing.get_mut(&edge.src_vid) {
outgoing_list.retain(|e| e.eid != edge.eid);
}
}
}
}
pub fn contains_vertex(&self, vid: Vid) -> bool {
self.vertices.contains_key(&vid)
}
pub fn vertex_count(&self) -> usize {
self.vertices.len()
}
pub fn add_edge(&mut self, src_vid: Vid, dst_vid: Vid, eid: Eid, edge_type: u32) {
self.add_vertex(src_vid);
self.add_vertex(dst_vid);
self.add_edge_unchecked(src_vid, dst_vid, eid, edge_type);
}
#[inline]
pub fn add_edge_unchecked(&mut self, src_vid: Vid, dst_vid: Vid, eid: Eid, edge_type: u32) {
let entry = EdgeEntry {
eid,
src_vid,
dst_vid,
edge_type,
};
self.outgoing.entry(src_vid).or_default().push(entry);
self.incoming.entry(dst_vid).or_default().push(entry);
self.edge_map.insert(eid, entry);
}
pub fn remove_edge(&mut self, eid: Eid) -> Option<EdgeEntry> {
let entry = self.edge_map.remove(&eid)?;
if let Some(outgoing_list) = self.outgoing.get_mut(&entry.src_vid) {
outgoing_list.retain(|e| e.eid != eid);
}
if let Some(incoming_list) = self.incoming.get_mut(&entry.dst_vid) {
incoming_list.retain(|e| e.eid != eid);
}
Some(entry)
}
pub fn neighbors(&self, vid: Vid, direction: Direction) -> &[EdgeEntry] {
let map = match direction {
Direction::Outgoing => &self.outgoing,
Direction::Incoming => &self.incoming,
};
map.get(&vid).map(|v| v.as_slice()).unwrap_or(&[])
}
pub fn edges(&self) -> impl Iterator<Item = &EdgeEntry> {
self.outgoing.values().flat_map(|v| v.iter())
}
pub fn edge(&self, eid: Eid) -> Option<&EdgeEntry> {
self.edge_map.get(&eid)
}
pub fn vertex(&self, vid: Vid) -> Option<Vid> {
if self.vertices.contains_key(&vid) {
Some(vid)
} else {
None
}
}
pub fn edge_count(&self) -> usize {
self.outgoing.values().map(|v| v.len()).sum()
}
pub fn vertices(&self) -> impl Iterator<Item = Vid> + '_ {
self.vertices.keys().copied()
}
pub fn clear(&mut self) {
self.vertices.clear();
self.outgoing.clear();
self.incoming.clear();
self.edge_map.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_add_remove_vertex() {
let mut g = SimpleGraph::new();
let v1 = Vid::new(1);
let v2 = Vid::new(2);
assert!(g.add_vertex(v1));
assert!(!g.add_vertex(v1)); assert!(g.add_vertex(v2));
assert_eq!(g.vertex_count(), 2);
assert!(g.contains_vertex(v1));
g.remove_vertex(v1);
assert_eq!(g.vertex_count(), 1);
assert!(!g.contains_vertex(v1));
}
#[test]
fn test_add_remove_edge() {
let mut g = SimpleGraph::new();
let v1 = Vid::new(1);
let v2 = Vid::new(2);
let v3 = Vid::new(3);
let e1 = Eid::new(101);
let e2 = Eid::new(102);
g.add_edge(v1, v2, e1, 1);
g.add_edge(v1, v3, e2, 1);
assert_eq!(g.edge_count(), 2);
let out = g.neighbors(v1, Direction::Outgoing);
assert_eq!(out.len(), 2);
let inc = g.neighbors(v2, Direction::Incoming);
assert_eq!(inc.len(), 1);
assert_eq!(inc[0].src_vid, v1);
let removed = g.remove_edge(e1);
assert!(removed.is_some());
assert_eq!(removed.unwrap().dst_vid, v2);
assert_eq!(g.edge_count(), 1);
assert_eq!(g.neighbors(v1, Direction::Outgoing).len(), 1);
assert_eq!(g.neighbors(v2, Direction::Incoming).len(), 0);
}
#[test]
fn test_remove_vertex_with_edges() {
let mut g = SimpleGraph::new();
let v1 = Vid::new(1);
let v2 = Vid::new(2);
let v3 = Vid::new(3);
let e1 = Eid::new(101);
let e2 = Eid::new(102);
g.add_edge(v1, v2, e1, 1);
g.add_edge(v2, v3, e2, 1);
g.remove_vertex(v2);
assert_eq!(g.vertex_count(), 2);
assert_eq!(g.edge_count(), 0);
assert_eq!(g.neighbors(v1, Direction::Outgoing).len(), 0);
assert_eq!(g.neighbors(v3, Direction::Incoming).len(), 0);
}
#[test]
fn test_edge_iteration() {
let mut g = SimpleGraph::new();
let v1 = Vid::new(1);
let v2 = Vid::new(2);
let v3 = Vid::new(3);
let e1 = Eid::new(101);
let e2 = Eid::new(102);
g.add_edge(v1, v2, e1, 1);
g.add_edge(v2, v3, e2, 2);
let edges: Vec<_> = g.edges().collect();
assert_eq!(edges.len(), 2);
}
#[test]
fn test_neighbor_iteration_after_deletion() {
let mut g = SimpleGraph::new();
let v1 = Vid::new(1);
let v2 = Vid::new(2);
let v3 = Vid::new(3);
let e1 = Eid::new(101);
let e2 = Eid::new(102);
g.add_edge(v1, v2, e1, 1);
g.add_edge(v1, v3, e2, 1);
g.remove_edge(e1);
let neighbors = g.neighbors(v1, Direction::Outgoing);
assert_eq!(neighbors.len(), 1);
assert_eq!(neighbors[0].dst_vid, v3);
}
#[test]
fn test_edge_lookup_o1() {
let mut g = SimpleGraph::new();
let v1 = Vid::new(1);
let v2 = Vid::new(2);
let v3 = Vid::new(3);
let e1 = Eid::new(101);
let e2 = Eid::new(102);
g.add_edge(v1, v2, e1, 1);
g.add_edge(v2, v3, e2, 2);
let edge = g.edge(e1);
assert!(edge.is_some());
let edge = edge.unwrap();
assert_eq!(edge.src_vid, v1);
assert_eq!(edge.dst_vid, v2);
assert_eq!(edge.edge_type, 1);
let non_existent = Eid::new(999);
assert!(g.edge(non_existent).is_none());
g.remove_edge(e1);
assert!(g.edge(e1).is_none());
assert!(g.edge(e2).is_some()); }
}