use core::cmp::max;
use core::ops::{Index, IndexMut};
use derive_more::derive::From;
#[derive(Clone, Copy, Debug, Default, PartialEq, PartialOrd, Eq, Ord, Hash, From)]
#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
pub struct NodeIndex(pub u32);
impl NodeIndex {
pub const END: NodeIndex = NodeIndex(u32::MAX);
pub fn index(self) -> usize {
self.0 as usize
}
pub fn is_end(self) -> bool {
self == NodeIndex::END
}
}
#[derive(Clone, Copy, Debug, Default, PartialEq, PartialOrd, Eq, Ord, Hash, From)]
#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
pub struct EdgeIndex(pub u32);
impl EdgeIndex {
pub const END: EdgeIndex = EdgeIndex(u32::MAX);
pub fn index(self) -> usize {
self.0 as usize
}
pub fn is_end(self) -> bool {
self == EdgeIndex::END
}
}
#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Ord, Eq, Hash)]
#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
#[repr(usize)]
pub enum EdgeDirection {
Outgoing = 0,
Incoming = 1,
}
impl EdgeDirection {
pub const ALL: [EdgeDirection; 2] = [EdgeDirection::Outgoing, EdgeDirection::Incoming];
}
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
pub struct Node<N> {
pub weight: N,
pub(super) next: [EdgeIndex; 2],
}
#[derive(Clone, Copy, Debug)]
#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
pub struct Edge<E> {
pub weight: E,
pub(super) next: [EdgeIndex; 2],
pub(super) node: [NodeIndex; 2],
}
impl<E> Edge<E> {
pub fn source(&self) -> NodeIndex {
self.node[0]
}
pub fn target(&self) -> NodeIndex {
self.node[1]
}
}
#[derive(Clone, Debug)]
#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
pub struct UnGraph<N, E> {
pub(super) nodes: Vec<Node<N>>,
pub(super) edges: Vec<Edge<E>>,
}
impl<N, E> Default for UnGraph<N, E> {
fn default() -> Self {
Self {
nodes: Vec::new(),
edges: Vec::new(),
}
}
}
pub(super) enum Pair<T> {
Both(T, T),
One(T),
None,
}
pub(super) fn index_twice<T>(arr: &mut [T], a: usize, b: usize) -> Pair<&mut T> {
if max(a, b) >= arr.len() {
Pair::None
} else if a == b {
Pair::One(&mut arr[max(a, b)])
} else {
unsafe {
let ar = &mut *(arr.get_unchecked_mut(a) as *mut _);
let br = &mut *(arr.get_unchecked_mut(b) as *mut _);
Pair::Both(ar, br)
}
}
}
impl<N, E> UnGraph<N, E> {
pub fn with_capacity(nodes: usize, edges: usize) -> Self {
UnGraph {
nodes: Vec::with_capacity(nodes),
edges: Vec::with_capacity(edges),
}
}
pub fn node_count(&self) -> usize {
self.nodes.len()
}
pub fn edge_count(&self) -> usize {
self.edges.len()
}
pub fn add_node(&mut self, weight: N) -> NodeIndex {
let node = Node {
weight,
next: [EdgeIndex::END, EdgeIndex::END],
};
assert!(self.nodes.len() as u32 != u32::MAX);
let node_idx = NodeIndex(self.nodes.len() as u32);
self.nodes.push(node);
node_idx
}
pub fn node_weight(&self, a: NodeIndex) -> Option<&N> {
self.nodes.get(a.index()).map(|n| &n.weight)
}
pub fn add_edge(&mut self, a: NodeIndex, b: NodeIndex, weight: E) -> EdgeIndex {
assert!(self.edges.len() as u32 != u32::MAX);
let edge_idx = EdgeIndex(self.edges.len() as u32);
let mut edge = Edge {
weight,
node: [a, b],
next: [EdgeIndex::END; 2],
};
match index_twice(&mut self.nodes, a.index(), b.index()) {
Pair::None => panic!("`UnGraph::add_edge`: node indices out of bounds"),
Pair::One(an) => {
edge.next = an.next;
an.next[0] = edge_idx;
an.next[1] = edge_idx;
}
Pair::Both(an, bn) => {
edge.next = [an.next[0], bn.next[1]];
an.next[0] = edge_idx;
bn.next[1] = edge_idx;
}
}
self.edges.push(edge);
edge_idx
}
pub fn update_edge(&mut self, a: NodeIndex, b: NodeIndex, weight: E) -> EdgeIndex {
if let Some(ix) = self.find_edge(a, b)
&& let Some(ed) = self.edge_weight_mut(ix)
{
*ed = weight;
return ix;
}
self.add_edge(a, b, weight)
}
pub fn edge_weight(&self, e: EdgeIndex) -> Option<&E> {
self.edges.get(e.index()).map(|e| &e.weight)
}
pub fn edge_weight_mut(&mut self, e: EdgeIndex) -> Option<&mut E> {
self.edges.get_mut(e.index()).map(|e| &mut e.weight)
}
pub fn edge_endpoints(&self, e: EdgeIndex) -> Option<(NodeIndex, NodeIndex)> {
self.edges
.get(e.index())
.map(|ed| (ed.source(), ed.target()))
}
pub fn remove_node_with<F>(&mut self, a: NodeIndex, mut edge_callback: F) -> Option<N>
where
F: FnMut(E),
{
if a.index() >= self.nodes.len() {
return None;
}
for d in EdgeDirection::ALL {
let k = d as usize;
loop {
let next = self.nodes[a.index()].next[k];
if next == EdgeIndex::END {
break;
}
let edge = self.remove_edge(next).expect("edge not found for removal");
edge_callback(edge);
}
}
let node = self.nodes.swap_remove(a.index());
let swap_edges = match self.nodes.get(a.index()) {
None => return Some(node.weight),
Some(ed) => ed.next,
};
let old_index = NodeIndex(self.nodes.len() as u32);
let new_index = a;
for d in EdgeDirection::ALL {
let k = d as usize;
let mut edges = EdgesWalkerMut {
edges: &mut self.edges,
next: swap_edges[k],
dir: d,
};
while let Some(curedge) = edges.next_edge() {
debug_assert!(curedge.node[k] == old_index);
curedge.node[k] = new_index;
}
}
Some(node.weight)
}
pub(super) fn change_edge_links(
&mut self,
edge_node: [NodeIndex; 2],
e: EdgeIndex,
edge_next: [EdgeIndex; 2],
) {
for d in EdgeDirection::ALL {
let k = d as usize;
let node = match self.nodes.get_mut(edge_node[k].index()) {
Some(r) => r,
None => {
debug_assert!(
false,
"Edge's endpoint dir={:?} index={:?} not found",
d, edge_node[k]
);
return;
}
};
let fst = node.next[k];
if fst == e {
node.next[k] = edge_next[k];
} else {
let mut edges = EdgesWalkerMut {
edges: &mut self.edges,
next: fst,
dir: d,
};
while let Some(curedge) = edges.next_edge() {
if curedge.next[k] == e {
curedge.next[k] = edge_next[k];
break;
}
}
}
}
}
pub fn remove_edge(&mut self, e: EdgeIndex) -> Option<E> {
let (edge_node, edge_next) = match self.edges.get(e.index()) {
None => return None,
Some(x) => (x.node, x.next),
};
self.change_edge_links(edge_node, e, edge_next);
self.remove_edge_adjust_indices(e)
}
fn remove_edge_adjust_indices(&mut self, e: EdgeIndex) -> Option<E> {
let edge = self.edges.swap_remove(e.index());
let swap = match self.edges.get(e.index()) {
None => return Some(edge.weight),
Some(ed) => ed.node,
};
let swapped_e = EdgeIndex(self.edges.len() as u32);
self.change_edge_links(swap, swapped_e, [e, e]);
Some(edge.weight)
}
pub fn neighbors(&self, a: NodeIndex) -> Neighbors<'_, E> {
Neighbors {
skip_start: a,
edges: &self.edges,
next: match self.nodes.get(a.index()) {
None => [EdgeIndex::END, EdgeIndex::END],
Some(n) => n.next,
},
}
}
pub fn edges(&self, a: NodeIndex) -> Edges<'_, E> {
Edges {
skip_start: a,
edges: &self.edges,
direction: EdgeDirection::Outgoing,
next: match self.nodes.get(a.index()) {
None => [EdgeIndex::END, EdgeIndex::END],
Some(n) => n.next,
},
}
}
pub fn edges_mut(&mut self, a: NodeIndex) -> EdgesMut<'_, N, E> {
let incoming_edge = self.first_edge(a, EdgeDirection::Incoming);
let outgoing_edge = self.first_edge(a, EdgeDirection::Outgoing);
EdgesMut {
graph: self,
incoming_edge,
outgoing_edge,
}
}
pub fn edges_between(&self, a: NodeIndex, b: NodeIndex) -> EdgesBetween<'_, E> {
EdgesBetween {
target_node: b,
edges: self.edges(a),
}
}
pub fn contains_edge(&self, a: NodeIndex, b: NodeIndex) -> bool {
self.find_edge(a, b).is_some()
}
pub fn find_edge(&self, a: NodeIndex, b: NodeIndex) -> Option<EdgeIndex> {
self.find_edge_from_node(self.nodes.get(a.index())?, b)
}
pub(super) fn find_edge_from_node(&self, node: &Node<N>, b: NodeIndex) -> Option<EdgeIndex> {
for &d in &EdgeDirection::ALL {
let k = d as usize;
let mut edix = node.next[k];
while let Some(edge) = self.edges.get(edix.index()) {
if edge.node[1 - k] == b {
return Some(edix);
}
edix = edge.next[k];
}
}
None
}
pub fn edge_weights(&self, a: NodeIndex) -> EdgeWeights<'_, E> {
EdgeWeights {
skip_start: a,
edges: &self.edges,
direction: EdgeDirection::Outgoing,
next: match self.nodes.get(a.index()) {
None => [EdgeIndex::END, EdgeIndex::END],
Some(n) => n.next,
},
}
}
pub fn edge_weights_mut(&mut self, a: NodeIndex) -> EdgeWeightsMut<'_, N, E> {
let incoming_edge = self.first_edge(a, EdgeDirection::Incoming);
let outgoing_edge = self.first_edge(a, EdgeDirection::Outgoing);
EdgeWeightsMut {
graph: self,
incoming_edge,
outgoing_edge,
}
}
pub fn all_edge_weights(&self) -> AllEdgeWeights<'_, E> {
AllEdgeWeights {
edges: self.edges.iter(),
}
}
pub fn all_edge_weights_mut(&mut self) -> AllEdgeWeightsMut<'_, E> {
AllEdgeWeightsMut {
edges: self.edges.iter_mut(),
}
}
pub fn raw_nodes(&self) -> &[Node<N>] {
&self.nodes
}
pub fn raw_nodes_mut(&mut self) -> &mut [Node<N>] {
&mut self.nodes
}
pub fn raw_edges(&self) -> &[Edge<E>] {
&self.edges
}
pub fn raw_edges_mut(&mut self) -> &mut [Edge<E>] {
&mut self.edges
}
pub fn first_edge(&self, a: NodeIndex, dir: EdgeDirection) -> Option<EdgeIndex> {
match self.nodes.get(a.index()) {
None => None,
Some(node) => {
let edix = node.next[dir as usize];
if edix == EdgeIndex::END {
None
} else {
Some(edix)
}
}
}
}
pub fn next_edge(&self, e: EdgeIndex, dir: EdgeDirection) -> Option<EdgeIndex> {
match self.edges.get(e.index()) {
None => None,
Some(node) => {
let edix = node.next[dir as usize];
if edix == EdgeIndex::END {
None
} else {
Some(edix)
}
}
}
}
pub fn clear(&mut self) {
self.nodes.clear();
self.edges.clear();
}
pub fn clear_edges(&mut self) {
self.edges.clear();
for node in &mut self.nodes {
node.next = [EdgeIndex::END, EdgeIndex::END];
}
}
pub fn nodes_capacity(&self) -> usize {
self.nodes.capacity()
}
pub fn edges_capacity(&self) -> usize {
self.edges.capacity()
}
pub fn reserve_nodes(&mut self, additional: usize) {
self.nodes.reserve(additional);
}
pub fn reserve_edges(&mut self, additional: usize) {
self.edges.reserve(additional);
}
}
#[derive(Debug)]
pub struct Neighbors<'a, E: 'a> {
skip_start: NodeIndex,
edges: &'a [Edge<E>],
next: [EdgeIndex; 2],
}
impl<E> Iterator for Neighbors<'_, E> {
type Item = NodeIndex;
fn next(&mut self) -> Option<NodeIndex> {
match self.edges.get(self.next[0].index()) {
None => {}
Some(edge) => {
self.next[0] = edge.next[0];
return Some(edge.node[1]);
}
}
while let Some(edge) = self.edges.get(self.next[1].index()) {
self.next[1] = edge.next[1];
if edge.node[0] != self.skip_start {
return Some(edge.node[0]);
}
}
None
}
}
impl<E> Clone for Neighbors<'_, E> {
fn clone(&self) -> Self {
Neighbors {
skip_start: self.skip_start,
edges: self.edges,
next: self.next,
}
}
}
pub struct Edges<'a, E: 'a> {
skip_start: NodeIndex,
edges: &'a [Edge<E>],
next: [EdgeIndex; 2],
direction: EdgeDirection,
}
impl<'a, E> Iterator for Edges<'a, E> {
type Item = EdgeReference<'a, E>;
fn next(&mut self) -> Option<Self::Item> {
let i = self.next[0].index();
if let Some(Edge {
node, weight, next, ..
}) = self.edges.get(i)
{
self.next[0] = next[0];
return Some(EdgeReference {
index: EdgeIndex(i as u32),
node: if self.direction == EdgeDirection::Incoming {
swap_pair(*node)
} else {
*node
},
weight,
});
}
while let Some(Edge { node, weight, next }) = self.edges.get(self.next[1].index()) {
let edge_index = self.next[1];
self.next[1] = next[1];
if node[0] == self.skip_start {
continue;
}
return Some(EdgeReference {
index: edge_index,
node: if self.direction == EdgeDirection::Outgoing {
swap_pair(*node)
} else {
*node
},
weight,
});
}
None
}
}
impl<E> Clone for Edges<'_, E> {
fn clone(&self) -> Self {
Edges {
skip_start: self.skip_start,
edges: self.edges,
next: self.next,
direction: self.direction,
}
}
}
pub struct EdgesMut<'a, N, E> {
graph: &'a mut UnGraph<N, E>,
incoming_edge: Option<EdgeIndex>,
outgoing_edge: Option<EdgeIndex>,
}
impl<'a, N: Copy, E> Iterator for EdgesMut<'a, N, E> {
type Item = EdgeMut<'a, E>;
#[inline]
fn next(&mut self) -> Option<EdgeMut<'a, E>> {
if let Some(edge) = self.incoming_edge {
self.incoming_edge = self.graph.next_edge(edge, EdgeDirection::Incoming);
let weights = &mut self.graph[edge];
return Some(EdgeMut {
index: edge,
weight: unsafe { core::mem::transmute::<&mut E, &'a mut E>(weights) },
});
}
let edge = self.outgoing_edge?;
self.outgoing_edge = self.graph.next_edge(edge, EdgeDirection::Outgoing);
let weights = &mut self.graph[edge];
Some(EdgeMut {
index: edge,
weight: unsafe { core::mem::transmute::<&mut E, &'a mut E>(weights) },
})
}
}
#[derive(Clone)]
pub struct EdgesBetween<'a, E: 'a> {
target_node: NodeIndex,
edges: Edges<'a, E>,
}
impl<'a, E> Iterator for EdgesBetween<'a, E> {
type Item = EdgeReference<'a, E>;
fn next(&mut self) -> Option<EdgeReference<'a, E>> {
let target_node = self.target_node;
self.edges
.by_ref()
.find(|&edge| edge.target() == target_node)
}
fn size_hint(&self) -> (usize, Option<usize>) {
let (_, upper) = self.edges.size_hint();
(0, upper)
}
}
fn swap_pair<T>(mut x: [T; 2]) -> [T; 2] {
x.swap(0, 1);
x
}
pub struct EdgeWeights<'a, E: 'a> {
skip_start: NodeIndex,
edges: &'a [Edge<E>],
next: [EdgeIndex; 2],
direction: EdgeDirection,
}
impl<'a, E> Iterator for EdgeWeights<'a, E> {
type Item = &'a E;
fn next(&mut self) -> Option<Self::Item> {
let i = self.next[0].index();
if let Some(Edge { weight, next, .. }) = self.edges.get(i) {
self.next[0] = next[0];
return Some(weight);
}
while let Some(Edge { node, weight, next }) = self.edges.get(self.next[1].index()) {
self.next[1] = next[1];
if node[0] == self.skip_start {
continue;
}
return Some(weight);
}
None
}
}
impl<E> Clone for EdgeWeights<'_, E> {
fn clone(&self) -> Self {
EdgeWeights {
skip_start: self.skip_start,
edges: self.edges,
next: self.next,
direction: self.direction,
}
}
}
pub struct EdgeWeightsMut<'a, N, E> {
pub graph: &'a mut UnGraph<N, E>,
pub incoming_edge: Option<EdgeIndex>,
pub outgoing_edge: Option<EdgeIndex>,
}
impl<'a, N: Copy, E> Iterator for EdgeWeightsMut<'a, N, E> {
type Item = &'a mut E;
#[inline]
fn next(&mut self) -> Option<&'a mut E> {
if let Some(edge) = self.incoming_edge {
self.incoming_edge = self.graph.next_edge(edge, EdgeDirection::Incoming);
let weight = &mut self.graph[edge];
return Some(unsafe { core::mem::transmute::<&mut E, &'a mut E>(weight) });
}
let edge = self.outgoing_edge?;
self.outgoing_edge = self.graph.next_edge(edge, EdgeDirection::Outgoing);
let weight = &mut self.graph[edge];
Some(unsafe { core::mem::transmute::<&mut E, &'a mut E>(weight) })
}
}
pub struct AllEdgeWeights<'a, E: 'a> {
edges: core::slice::Iter<'a, Edge<E>>,
}
impl<'a, E> Iterator for AllEdgeWeights<'a, E> {
type Item = &'a E;
fn next(&mut self) -> Option<&'a E> {
self.edges.next().map(|edge| &edge.weight)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.edges.size_hint()
}
}
#[derive(Debug)]
pub struct AllEdgeWeightsMut<'a, E: 'a> {
edges: core::slice::IterMut<'a, Edge<E>>,
}
impl<'a, E> Iterator for AllEdgeWeightsMut<'a, E> {
type Item = &'a mut E;
fn next(&mut self) -> Option<&'a mut E> {
self.edges.next().map(|edge| &mut edge.weight)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.edges.size_hint()
}
}
struct EdgesWalkerMut<'a, E: 'a> {
edges: &'a mut [Edge<E>],
next: EdgeIndex,
dir: EdgeDirection,
}
impl<E> EdgesWalkerMut<'_, E> {
fn next_edge(&mut self) -> Option<&mut Edge<E>> {
self.next().map(|t| t.1)
}
fn next(&mut self) -> Option<(EdgeIndex, &mut Edge<E>)> {
let this_index = self.next;
let k = self.dir as usize;
match self.edges.get_mut(self.next.index()) {
None => None,
Some(edge) => {
self.next = edge.next[k];
Some((this_index, edge))
}
}
}
}
impl<N, E> Index<NodeIndex> for UnGraph<N, E> {
type Output = N;
fn index(&self, index: NodeIndex) -> &N {
&self.nodes[index.index()].weight
}
}
impl<N, E> IndexMut<NodeIndex> for UnGraph<N, E> {
fn index_mut(&mut self, index: NodeIndex) -> &mut N {
&mut self.nodes[index.index()].weight
}
}
impl<N, E> Index<EdgeIndex> for UnGraph<N, E> {
type Output = E;
fn index(&self, index: EdgeIndex) -> &E {
&self.edges[index.index()].weight
}
}
impl<N, E> IndexMut<EdgeIndex> for UnGraph<N, E> {
fn index_mut(&mut self, index: EdgeIndex) -> &mut E {
&mut self.edges[index.index()].weight
}
}
#[derive(Debug)]
pub struct EdgeReference<'a, E: 'a> {
index: EdgeIndex,
node: [NodeIndex; 2],
weight: &'a E,
}
impl<'a, E: 'a> EdgeReference<'a, E> {
#[inline]
pub fn index(&self) -> EdgeIndex {
self.index
}
#[inline]
pub fn source(&self) -> NodeIndex {
self.node[0]
}
#[inline]
pub fn target(&self) -> NodeIndex {
self.node[1]
}
#[inline]
pub fn weight(&self) -> &'a E {
self.weight
}
}
impl<E> Clone for EdgeReference<'_, E> {
fn clone(&self) -> Self {
*self
}
}
impl<E> Copy for EdgeReference<'_, E> {}
impl<E> PartialEq for EdgeReference<'_, E>
where
E: PartialEq,
{
fn eq(&self, rhs: &Self) -> bool {
self.index == rhs.index && self.weight == rhs.weight
}
}
#[derive(Debug)]
pub struct EdgeMut<'a, E: 'a> {
index: EdgeIndex,
weight: &'a mut E,
}
impl<E> EdgeMut<'_, E> {
#[inline]
pub fn index(&self) -> EdgeIndex {
self.index
}
#[inline]
pub fn weight(&self) -> &E {
self.weight
}
#[inline]
pub fn weight_mut(&mut self) -> &mut E {
self.weight
}
}
impl<E> PartialEq for EdgeMut<'_, E>
where
E: PartialEq,
{
fn eq(&self, rhs: &Self) -> bool {
self.index == rhs.index && self.weight == rhs.weight
}
}