use alloc::{
collections::{BTreeMap, BTreeSet},
vec::Vec,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct OverlayEdge {
pub source: u32,
pub target: u32,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct OverlayState {
added_edges: Vec<OverlayEdge>,
by_source: BTreeMap<u32, Vec<u32>>,
by_target: BTreeMap<u32, Vec<u32>>,
tombstoned_edges: BTreeSet<u32>,
tombstoned_nodes: BTreeSet<u32>,
}
impl OverlayState {
#[must_use]
pub fn from_edges(edges: impl IntoIterator<Item = OverlayEdge>) -> Self {
let mut overlay = Self::default();
for edge in edges {
overlay.push_edge(edge);
}
overlay
}
#[must_use]
pub fn added_edges(&self) -> &[OverlayEdge] {
&self.added_edges
}
#[must_use]
pub const fn overlay_edge_count(&self) -> usize {
self.added_edges.len()
}
#[must_use]
pub fn tombstoned_edge_count(&self) -> usize {
self.tombstoned_edges.len()
}
#[must_use]
pub fn tombstoned_node_count(&self) -> usize {
self.tombstoned_nodes.len()
}
pub fn tombstone_edge(&mut self, edge_id: u32) {
self.tombstoned_edges.insert(edge_id);
}
pub fn tombstone_node(&mut self, node_id: u32) {
self.tombstoned_nodes.insert(node_id);
}
pub fn clear(&mut self) {
self.added_edges.clear();
self.by_source.clear();
self.by_target.clear();
self.tombstoned_edges.clear();
self.tombstoned_nodes.clear();
}
#[must_use]
pub fn has_edge_tombstones(&self) -> bool {
!self.tombstoned_edges.is_empty()
}
#[must_use]
pub fn has_node_tombstones(&self) -> bool {
!self.tombstoned_nodes.is_empty()
}
#[must_use]
pub fn edge_visible(&self, edge_id: u32) -> bool {
!self.tombstoned_edges.contains(&edge_id)
}
#[must_use]
pub fn node_visible(&self, node_id: u32) -> bool {
!self.tombstoned_nodes.contains(&node_id)
}
pub fn push_edge(&mut self, edge: OverlayEdge) {
self.added_edges.push(edge);
self.by_source
.entry(edge.source)
.or_default()
.push(edge.target);
self.by_target
.entry(edge.target)
.or_default()
.push(edge.source);
}
#[must_use]
pub fn overlay_targets(&self, source: u32) -> &[u32] {
self.by_source
.get(&source)
.map_or(&[] as &[u32], Vec::as_slice)
}
#[must_use]
pub fn overlay_sources(&self, target: u32) -> &[u32] {
self.by_target
.get(&target)
.map_or(&[] as &[u32], Vec::as_slice)
}
pub fn remove_edge(&mut self, source: u32, target: u32) {
self.added_edges
.retain(|edge| edge.source != source || edge.target != target);
self.rebuild_indexes();
}
pub(crate) fn rebuild_indexes(&mut self) {
self.by_source.clear();
self.by_target.clear();
for edge in &self.added_edges {
self.by_source
.entry(edge.source)
.or_default()
.push(edge.target);
self.by_target
.entry(edge.target)
.or_default()
.push(edge.source);
}
}
}
#[cfg(test)]
mod tests {
use super::{OverlayEdge, OverlayState};
#[test]
fn indexes_match_linear_scan() {
let mut overlay = OverlayState::default();
overlay.push_edge(OverlayEdge {
source: 1,
target: 2,
});
overlay.push_edge(OverlayEdge {
source: 1,
target: 3,
});
overlay.push_edge(OverlayEdge {
source: 4,
target: 1,
});
assert_eq!(overlay.overlay_targets(1), &[2, 3]);
assert_eq!(overlay.overlay_sources(1), &[4]);
overlay.clear();
assert!(overlay.overlay_targets(1).is_empty());
}
#[test]
fn from_edges_builds_indexes() {
let overlay = OverlayState::from_edges([
OverlayEdge {
source: 0,
target: 1,
},
OverlayEdge {
source: 2,
target: 0,
},
]);
assert_eq!(overlay.overlay_targets(0), &[1]);
assert_eq!(overlay.overlay_sources(0), &[2]);
}
}