avian3d/data_structures/
graph.rs

1//! A stripped down version of petgraph's `UnGraph`.
2//!
3//! - Index types always use `u32`.
4//! - Edge iteration order after serialization/deserialization is preserved.
5//! - Fewer iterators and helpers, and a few new ones.
6
7use core::cmp::max;
8use core::ops::{Index, IndexMut};
9
10use derive_more::derive::From;
11
12/// A node identifier for a graph structure.
13#[derive(Clone, Copy, Debug, Default, PartialEq, PartialOrd, Eq, Ord, Hash, From)]
14#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
15pub struct NodeIndex(pub u32);
16
17impl NodeIndex {
18    /// A special index used to denote the final node,
19    /// for example at the end of an adjacency list.
20    ///
21    /// Equivalent to `NodeIndex(u32::MAX)`.
22    pub const END: NodeIndex = NodeIndex(u32::MAX);
23
24    /// Returns the inner `u32` value as a `usize`.
25    pub fn index(self) -> usize {
26        self.0 as usize
27    }
28
29    /// Returns `true` if this is the special `END` index.
30    pub fn is_end(self) -> bool {
31        self == NodeIndex::END
32    }
33}
34
35/// An edge identifier for a graph structure.
36#[derive(Clone, Copy, Debug, Default, PartialEq, PartialOrd, Eq, Ord, Hash, From)]
37#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
38pub struct EdgeIndex(pub u32);
39
40impl EdgeIndex {
41    /// A special index used to denote the abscence of an edge,
42    /// for example at the end of an adjacency list.
43    ///
44    /// Equivalent to `EdgeIndex(u32::MAX)`.
45    pub const END: EdgeIndex = EdgeIndex(u32::MAX);
46
47    /// Returns the inner `u32` value as a `usize`.
48    pub fn index(self) -> usize {
49        self.0 as usize
50    }
51
52    /// Returns `true` if this is the special `END` index.
53    pub fn is_end(self) -> bool {
54        self == EdgeIndex::END
55    }
56}
57
58/// The direction of a graph edge.
59#[derive(Clone, Copy, Debug, PartialEq, PartialOrd, Ord, Eq, Hash)]
60#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
61#[repr(usize)]
62pub enum EdgeDirection {
63    /// An `Outgoing` edge is an outward edge *from* the current node.
64    Outgoing = 0,
65    /// An `Incoming` edge is an inbound edge *to* the current node.
66    Incoming = 1,
67}
68
69impl EdgeDirection {
70    /// The two possible directions for an edge.
71    pub const ALL: [EdgeDirection; 2] = [EdgeDirection::Outgoing, EdgeDirection::Incoming];
72
73    /// Returns the opposite direction.
74    #[inline]
75    fn opposite(self) -> EdgeDirection {
76        match self {
77            EdgeDirection::Outgoing => EdgeDirection::Incoming,
78            EdgeDirection::Incoming => EdgeDirection::Outgoing,
79        }
80    }
81}
82
83/// The node type for a graph structure.
84#[derive(Clone, Copy, Debug)]
85#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
86pub struct Node<N> {
87    /// Associated node data.
88    pub weight: N,
89    /// Next edge in outgoing and incoming edge lists.
90    next: [EdgeIndex; 2],
91}
92
93/// The edge type for a graph structure.
94#[derive(Clone, Copy, Debug)]
95#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
96pub struct Edge<E> {
97    /// Associated edge data.
98    pub weight: E,
99    /// Next edge in outgoing and incoming edge lists.
100    next: [EdgeIndex; 2],
101    /// Start and End node index
102    node: [NodeIndex; 2],
103}
104
105impl<E> Edge<E> {
106    /// Return the source node index.
107    pub fn source(&self) -> NodeIndex {
108        self.node[0]
109    }
110
111    /// Return the target node index.
112    pub fn target(&self) -> NodeIndex {
113        self.node[1]
114    }
115}
116
117/// A graph with undirected edges.
118///
119/// For example, an edge between *1* and *2* is equivalent to an edge between
120/// *2* and *1*.
121#[derive(Clone, Debug)]
122#[cfg_attr(feature = "serialize", derive(serde::Serialize, serde::Deserialize))]
123pub struct UnGraph<N, E> {
124    nodes: Vec<Node<N>>,
125    edges: Vec<Edge<E>>,
126}
127
128impl<N, E> Default for UnGraph<N, E> {
129    fn default() -> Self {
130        Self {
131            nodes: Vec::new(),
132            edges: Vec::new(),
133        }
134    }
135}
136
137enum Pair<T> {
138    Both(T, T),
139    One(T),
140    None,
141}
142
143/// Get mutable references at index `a` and `b`.
144fn index_twice<T>(arr: &mut [T], a: usize, b: usize) -> Pair<&mut T> {
145    if max(a, b) >= arr.len() {
146        Pair::None
147    } else if a == b {
148        Pair::One(&mut arr[max(a, b)])
149    } else {
150        // safe because `a` and `b` are in bounds and distinct.
151        unsafe {
152            let ar = &mut *(arr.get_unchecked_mut(a) as *mut _);
153            let br = &mut *(arr.get_unchecked_mut(b) as *mut _);
154            Pair::Both(ar, br)
155        }
156    }
157}
158
159impl<N, E> UnGraph<N, E> {
160    /// Creates a new [`UnGraph`] with estimated capacity.
161    pub fn with_capacity(nodes: usize, edges: usize) -> Self {
162        UnGraph {
163            nodes: Vec::with_capacity(nodes),
164            edges: Vec::with_capacity(edges),
165        }
166    }
167
168    /// Returns the number of nodes in the graph.
169    ///
170    /// Computes in **O(1)** time.
171    pub fn node_count(&self) -> usize {
172        self.nodes.len()
173    }
174
175    /// Returns the number of edges in the graph.
176    ///
177    /// Computes in **O(1)** time.
178    pub fn edge_count(&self) -> usize {
179        self.edges.len()
180    }
181
182    /// Adds a node (also called vertex) with associated data `weight` to the graph.
183    ///
184    /// Computes in **O(1)** time.
185    ///
186    /// Returns the index of the new node.
187    ///
188    /// # Panics
189    ///
190    /// Panics if the graph is at the maximum number of nodes.
191    pub fn add_node(&mut self, weight: N) -> NodeIndex {
192        let node = Node {
193            weight,
194            next: [EdgeIndex::END, EdgeIndex::END],
195        };
196        assert!(self.nodes.len() != usize::MAX);
197        let node_idx = NodeIndex(self.nodes.len() as u32);
198        self.nodes.push(node);
199        node_idx
200    }
201
202    /// Accesses the weight for node `a`.
203    ///
204    /// Also available with indexing syntax: `&graph[a]`.
205    pub fn node_weight(&self, a: NodeIndex) -> Option<&N> {
206        self.nodes.get(a.index()).map(|n| &n.weight)
207    }
208
209    /// Adds an edge from `a` to `b` to the graph, with its associated
210    /// data `weight`.
211    ///
212    /// Returns the index of the new edge.
213    ///
214    /// Computes in **O(1)** time.
215    ///
216    /// **Note:** `UnGraph` allows adding parallel (“duplicate”) edges. If you want
217    /// to avoid this, use [`.update_edge(a, b, weight)`](#method.update_edge) instead.
218    ///
219    /// # Panics
220    ///
221    /// Panics if any of the nodes don't exist or the graph is at the maximum number of edges.
222    pub fn add_edge(&mut self, a: NodeIndex, b: NodeIndex, weight: E) -> EdgeIndex {
223        assert!(self.edges.len() != usize::MAX);
224        let edge_idx = EdgeIndex(self.edges.len() as u32);
225        let mut edge = Edge {
226            weight,
227            node: [a, b],
228            next: [EdgeIndex::END; 2],
229        };
230        match index_twice(&mut self.nodes, a.index(), b.index()) {
231            Pair::None => panic!("`UnGraph::add_edge`: node indices out of bounds"),
232            Pair::One(an) => {
233                edge.next = an.next;
234                an.next[0] = edge_idx;
235                an.next[1] = edge_idx;
236            }
237            Pair::Both(an, bn) => {
238                // `a` and `b` are different indices
239                edge.next = [an.next[0], bn.next[1]];
240                an.next[0] = edge_idx;
241                bn.next[1] = edge_idx;
242            }
243        }
244        self.edges.push(edge);
245        edge_idx
246    }
247
248    /// Adds or updates an edge from `a` to `b`.
249    /// If the edge already exists, its weight is updated.
250    ///
251    /// Returns the index of the affected edge.
252    ///
253    /// Computes in **O(e')** time, where **e'** is the number of edges
254    /// connected to `a` (and `b`, if the graph edges are undirected).
255    ///
256    /// # Panics
257    ///
258    /// Panics if any of the nodes doesn't exist.
259    pub fn update_edge(&mut self, a: NodeIndex, b: NodeIndex, weight: E) -> EdgeIndex {
260        if let Some(ix) = self.find_edge(a, b) {
261            if let Some(ed) = self.edge_weight_mut(ix) {
262                *ed = weight;
263                return ix;
264            }
265        }
266        self.add_edge(a, b, weight)
267    }
268
269    /// Accesses the weight for edge `e`.
270    ///
271    /// Also available with indexing syntax: `&graph[e]`.
272    pub fn edge_weight(&self, e: EdgeIndex) -> Option<&E> {
273        self.edges.get(e.index()).map(|e| &e.weight)
274    }
275
276    /// Accesses the weight for edge `e` mutably.
277    ///
278    /// Also available with indexing syntax: `&mut graph[e]`.
279    pub fn edge_weight_mut(&mut self, e: EdgeIndex) -> Option<&mut E> {
280        self.edges.get_mut(e.index()).map(|e| &mut e.weight)
281    }
282
283    /// Accesses the source and target nodes for `e`.
284    pub fn edge_endpoints(&self, e: EdgeIndex) -> Option<(NodeIndex, NodeIndex)> {
285        self.edges
286            .get(e.index())
287            .map(|ed| (ed.source(), ed.target()))
288    }
289
290    /// Removes `a` from the graph if it exists, calling `edge_callback` for each of its edges,
291    /// and returns its weight. If it doesn't exist in the graph, returns `None`.
292    ///
293    /// Apart from `a`, this invalidates the last node index in the graph
294    /// (that node will adopt the removed node index). Edge indices are
295    /// invalidated as they would be following the removal of each edge
296    /// with an endpoint in `a`.
297    ///
298    /// Computes in **O(e')** time, where **e'** is the number of affected
299    /// edges, including *n* calls to [`remove_edge`](Self::remove_edge),
300    /// where *n* is the number of edges with an endpoint in `a`,
301    /// and including the edges with an endpoint in the displaced node.
302    pub fn remove_node_with<F>(&mut self, a: NodeIndex, mut edge_callback: F) -> Option<N>
303    where
304        F: FnMut(E),
305    {
306        self.nodes.get(a.index())?;
307        for d in EdgeDirection::ALL {
308            let k = d as usize;
309
310            // Remove all edges from and to this node.
311            loop {
312                let next = self.nodes[a.index()].next[k];
313                if next == EdgeIndex::END {
314                    break;
315                }
316                let edge = self.remove_edge(next).expect("edge not found for removal");
317                edge_callback(edge);
318            }
319        }
320
321        // Use `swap_remove` -- only the swapped-in node is going to change
322        // `NodeIndex`, so we only have to walk its edges and update them.
323
324        let node = self.nodes.swap_remove(a.index());
325
326        // Find the edge lists of the node that had to relocate.
327        // It may be that no node had to relocate, then we are done already.
328        let swap_edges = match self.nodes.get(a.index()) {
329            None => return Some(node.weight),
330            Some(ed) => ed.next,
331        };
332
333        // The swapped element's old index
334        let old_index = NodeIndex(self.nodes.len() as u32);
335        let new_index = a;
336
337        // Adjust the starts of the out edges, and ends of the in edges.
338        for d in EdgeDirection::ALL {
339            let k = d as usize;
340            let mut edges = EdgesWalkerMut {
341                edges: &mut self.edges,
342                next: swap_edges[k],
343                dir: d,
344            };
345            while let Some(curedge) = edges.next_edge() {
346                debug_assert!(curedge.node[k] == old_index);
347                curedge.node[k] = new_index;
348            }
349        }
350        Some(node.weight)
351    }
352
353    /// For edge `e` with endpoints `edge_node`, replaces links to it
354    /// with links to `edge_next`.
355    fn change_edge_links(
356        &mut self,
357        edge_node: [NodeIndex; 2],
358        e: EdgeIndex,
359        edge_next: [EdgeIndex; 2],
360    ) {
361        for d in EdgeDirection::ALL {
362            let k = d as usize;
363            let node = match self.nodes.get_mut(edge_node[k].index()) {
364                Some(r) => r,
365                None => {
366                    debug_assert!(
367                        false,
368                        "Edge's endpoint dir={:?} index={:?} not found",
369                        d, edge_node[k]
370                    );
371                    return;
372                }
373            };
374            let fst = node.next[k];
375            if fst == e {
376                node.next[k] = edge_next[k];
377            } else {
378                let mut edges = EdgesWalkerMut {
379                    edges: &mut self.edges,
380                    next: fst,
381                    dir: d,
382                };
383                while let Some(curedge) = edges.next_edge() {
384                    if curedge.next[k] == e {
385                        curedge.next[k] = edge_next[k];
386                        // The edge can only be present once in the list.
387                        break;
388                    }
389                }
390            }
391        }
392    }
393
394    /// Removes an edge and returns its edge weight, or `None` if it didn't exist.
395    ///
396    /// Apart from `e`, this invalidates the last edge index in the graph
397    /// (that edge will adopt the removed edge index).
398    ///
399    /// Computes in **O(e')** time, where **e'** is the size of four particular edge lists, for
400    /// the vertices of `e` and the vertices of another affected edge.
401    pub fn remove_edge(&mut self, e: EdgeIndex) -> Option<E> {
402        // Every edge is part of two lists, outgoing and incoming edges.
403        // Remove it from both.
404        let (edge_node, edge_next) = match self.edges.get(e.index()) {
405            None => return None,
406            Some(x) => (x.node, x.next),
407        };
408
409        // Remove the edge from its in and out lists by replacing it with
410        // a link to the next in the list.
411        self.change_edge_links(edge_node, e, edge_next);
412        self.remove_edge_adjust_indices(e)
413    }
414
415    fn remove_edge_adjust_indices(&mut self, e: EdgeIndex) -> Option<E> {
416        // `swap_remove` the edge -- only the removed edge
417        // and the edge swapped into place are affected and need updating
418        // indices.
419        let edge = self.edges.swap_remove(e.index());
420        let swap = match self.edges.get(e.index()) {
421            // No element needed to be swapped.
422            None => return Some(edge.weight),
423            Some(ed) => ed.node,
424        };
425        let swapped_e = EdgeIndex(self.edges.len() as u32);
426
427        // Update the edge lists by replacing links to the old index by references to the new
428        // edge index.
429        self.change_edge_links(swap, swapped_e, [e, e]);
430
431        Some(edge.weight)
432    }
433
434    /// Returns an iterator of all nodes with an edge connected to `a`.
435    ///
436    /// Produces an empty iterator if the node doesn't exist.
437    ///
438    /// The iterator element type is `NodeIndex`.
439    pub fn neighbors(&self, a: NodeIndex) -> Neighbors<E> {
440        Neighbors {
441            skip_start: a,
442            edges: &self.edges,
443            next: match self.nodes.get(a.index()) {
444                None => [EdgeIndex::END, EdgeIndex::END],
445                Some(n) => n.next,
446            },
447        }
448    }
449
450    /// Returns an iterator of all edges connected to `a`.
451    ///
452    /// Produces an empty iterator if the node doesn't exist.
453    ///
454    /// The iterator element type is `EdgeReference<E>`.
455    pub fn edges(&self, a: NodeIndex) -> Edges<E> {
456        Edges {
457            skip_start: a,
458            edges: &self.edges,
459            direction: EdgeDirection::Outgoing,
460            next: match self.nodes.get(a.index()) {
461                None => [EdgeIndex::END, EdgeIndex::END],
462                Some(n) => n.next,
463            },
464        }
465    }
466
467    /// Returns a mutable iterator of all edges connected to `a`.
468    ///
469    /// Produces an empty iterator if the node doesn't exist.
470    ///
471    /// The iterator element type is `EdgeReference<E>`.
472    pub fn edges_mut(&mut self, a: NodeIndex) -> EdgesMut<N, E> {
473        let incoming_edge = self.first_edge(a, EdgeDirection::Incoming);
474        let outgoing_edge = self.first_edge(a, EdgeDirection::Outgoing);
475
476        EdgesMut {
477            graph: self,
478            incoming_edge,
479            outgoing_edge,
480        }
481    }
482
483    /// Looks up if there is an edge from `a` to `b`.
484    ///
485    /// Computes in **O(e')** time, where **e'** is the number of edges
486    /// connected to `a` (and `b`, if the graph edges are undirected).
487    pub fn contains_edge(&self, a: NodeIndex, b: NodeIndex) -> bool {
488        self.find_edge(a, b).is_some()
489    }
490
491    /// Looks up an edge from `a` to `b`.
492    ///
493    /// Computes in **O(e')** time, where **e'** is the number of edges
494    /// connected to `a` (and `b`, if the graph edges are undirected).
495    pub fn find_edge(&self, a: NodeIndex, b: NodeIndex) -> Option<EdgeIndex> {
496        let node = self.nodes.get(a.index())?;
497        for &d in &EdgeDirection::ALL {
498            let k = d as usize;
499            let mut edix = node.next[k];
500            while let Some(edge) = self.edges.get(edix.index()) {
501                if edge.node[1 - k] == b {
502                    return Some(edix);
503                }
504                edix = edge.next[k];
505            }
506        }
507        None
508    }
509
510    /// Returns an iterator over the node indices of the graph.
511    pub fn node_indices(&self) -> impl DoubleEndedIterator<Item = NodeIndex> {
512        (0..self.node_count()).map(|i| NodeIndex(i as u32))
513    }
514
515    /// Returns an iterator over the edge indices of the graph
516    pub fn edge_indices(&self) -> impl DoubleEndedIterator<Item = EdgeIndex> {
517        (0..self.edge_count()).map(|i| EdgeIndex(i as u32))
518    }
519
520    /// Returns an iterator yielding immutable access to edge weights for edges from or to `a`.
521    pub fn edge_weights(&self, a: NodeIndex) -> EdgeWeights<E> {
522        EdgeWeights {
523            skip_start: a,
524            edges: &self.edges,
525            direction: EdgeDirection::Outgoing,
526            next: match self.nodes.get(a.index()) {
527                None => [EdgeIndex::END, EdgeIndex::END],
528                Some(n) => n.next,
529            },
530        }
531    }
532
533    /// Returns an iterator yielding mutable access to edge weights for edges from or to `a`.
534    pub fn edge_weights_mut(&mut self, a: NodeIndex) -> EdgeWeightsMut<N, E> {
535        let incoming_edge = self.first_edge(a, EdgeDirection::Incoming);
536        let outgoing_edge = self.first_edge(a, EdgeDirection::Outgoing);
537
538        EdgeWeightsMut {
539            graph: self,
540            incoming_edge,
541            outgoing_edge,
542        }
543    }
544
545    /// Returns an iterator yielding immutable access to all edge weights.
546    ///
547    /// The order in which weights are yielded matches the order of their
548    /// edge indices.
549    pub fn all_edge_weights(&self) -> AllEdgeWeights<E> {
550        AllEdgeWeights {
551            edges: self.edges.iter(),
552        }
553    }
554
555    /// Returns an iterator yielding mutable access to all edge weights.
556    ///
557    /// The order in which weights are yielded matches the order of their
558    /// edge indices.
559    pub fn all_edge_weights_mut(&mut self) -> AllEdgeWeightsMut<E> {
560        AllEdgeWeightsMut {
561            edges: self.edges.iter_mut(),
562        }
563    }
564
565    /// Accesses the internal node array.
566    pub fn raw_nodes(&self) -> &[Node<N>] {
567        &self.nodes
568    }
569
570    /// Accesses the internal node array mutably.
571    pub fn raw_nodes_mut(&mut self) -> &mut [Node<N>] {
572        &mut self.nodes
573    }
574
575    /// Accesses the internal edge array.
576    pub fn raw_edges(&self) -> &[Edge<E>] {
577        &self.edges
578    }
579
580    /// Accesses the internal edge array mutably.
581    pub fn raw_edges_mut(&mut self) -> &mut [Edge<E>] {
582        &mut self.edges
583    }
584
585    /// Accessor for data structure internals: returns the first edge in the given direction.
586    pub fn first_edge(&self, a: NodeIndex, dir: EdgeDirection) -> Option<EdgeIndex> {
587        match self.nodes.get(a.index()) {
588            None => None,
589            Some(node) => {
590                let edix = node.next[dir as usize];
591                if edix == EdgeIndex::END {
592                    None
593                } else {
594                    Some(edix)
595                }
596            }
597        }
598    }
599
600    /// Accessor for data structure internals: returns the next edge for the given direction.
601    pub fn next_edge(&self, e: EdgeIndex, dir: EdgeDirection) -> Option<EdgeIndex> {
602        match self.edges.get(e.index()) {
603            None => None,
604            Some(node) => {
605                let edix = node.next[dir as usize];
606                if edix == EdgeIndex::END {
607                    None
608                } else {
609                    Some(edix)
610                }
611            }
612        }
613    }
614
615    /// Removes all nodes and edges.
616    pub fn clear(&mut self) {
617        self.nodes.clear();
618        self.edges.clear();
619    }
620
621    /// Removes all edges.
622    pub fn clear_edges(&mut self) {
623        self.edges.clear();
624        for node in &mut self.nodes {
625            node.next = [EdgeIndex::END, EdgeIndex::END];
626        }
627    }
628
629    /// Returns the current node capacity of the graph.
630    pub fn nodes_capacity(&self) -> usize {
631        self.nodes.capacity()
632    }
633
634    /// Returns the current edge capacity of the graph.
635    pub fn edges_capacity(&self) -> usize {
636        self.edges.capacity()
637    }
638
639    /// Reserves capacity for at least `additional` more nodes to be inserted in
640    /// the graph. Graph may reserve more space to avoid frequent reallocations.
641    ///
642    /// # Panics
643    ///
644    /// Panics if the new capacity overflows `usize`.
645    pub fn reserve_nodes(&mut self, additional: usize) {
646        self.nodes.reserve(additional);
647    }
648
649    /// Reserves capacity for at least `additional` more edges to be inserted in
650    /// the graph. Graph may reserve more space to avoid frequent reallocations.
651    ///
652    /// # Panics
653    ///
654    /// Panics if the new capacity overflows `usize`.
655    pub fn reserve_edges(&mut self, additional: usize) {
656        self.edges.reserve(additional);
657    }
658
659    /// Keep all nodes that return `true` from the `visit` closure,
660    /// remove the others.
661    ///
662    /// `visit` is provided a proxy reference to the graph, so that
663    /// the graph can be walked and associated data modified.
664    ///
665    /// The order nodes are visited is not specified.
666    pub fn retain_nodes<F>(&mut self, mut visit: F)
667    where
668        F: FnMut(&Self, NodeIndex) -> bool,
669    {
670        for index in self.node_indices().rev() {
671            if !visit(self, index) {
672                let ret = self.remove_node_with(index, |_| ());
673                debug_assert!(ret.is_some());
674                let _ = ret;
675            }
676        }
677    }
678
679    /// Keep all edges that return `true` from the `visit` closure,
680    /// remove the others.
681    ///
682    /// `visit` is provided a proxy reference to the graph, so that
683    /// the graph can be walked and associated data modified.
684    ///
685    /// The order edges are visited is not specified.
686    pub fn retain_edges<F>(&mut self, mut visit: F)
687    where
688        F: FnMut(&Self, EdgeIndex) -> bool,
689    {
690        for index in self.edge_indices().rev() {
691            if !visit(self, index) {
692                let ret = self.remove_edge(index);
693                debug_assert!(ret.is_some());
694                let _ = ret;
695            }
696        }
697    }
698}
699
700/// An iterator over the neighbors of a node.
701///
702/// The iterator element type is `NodeIndex`.
703#[derive(Debug)]
704pub struct Neighbors<'a, E: 'a> {
705    /// The starting node to skip over.
706    skip_start: NodeIndex,
707
708    /// The edges to iterate over.
709    edges: &'a [Edge<E>],
710
711    /// The next edge to visit.
712    next: [EdgeIndex; 2],
713}
714
715impl<E> Iterator for Neighbors<'_, E> {
716    type Item = NodeIndex;
717
718    fn next(&mut self) -> Option<NodeIndex> {
719        // First any outgoing edges.
720        match self.edges.get(self.next[0].index()) {
721            None => {}
722            Some(edge) => {
723                self.next[0] = edge.next[0];
724                return Some(edge.node[1]);
725            }
726        }
727
728        // Then incoming edges.
729        // For an "undirected" iterator, make sure we don't double
730        // count self-loops by skipping them in the incoming list.
731        while let Some(edge) = self.edges.get(self.next[1].index()) {
732            self.next[1] = edge.next[1];
733            if edge.node[0] != self.skip_start {
734                return Some(edge.node[0]);
735            }
736        }
737        None
738    }
739}
740
741impl<E> Clone for Neighbors<'_, E> {
742    fn clone(&self) -> Self {
743        Neighbors {
744            skip_start: self.skip_start,
745            edges: self.edges,
746            next: self.next,
747        }
748    }
749}
750
751/// An iterator over edges from or to a node.
752pub struct Edges<'a, E: 'a> {
753    /// The starting node to skip over.
754    skip_start: NodeIndex,
755
756    /// The edges to iterate over.
757    edges: &'a [Edge<E>],
758
759    /// The next edge to visit.
760    next: [EdgeIndex; 2],
761
762    /// The direction of edges.
763    direction: EdgeDirection,
764}
765
766impl<'a, E> Iterator for Edges<'a, E> {
767    type Item = EdgeReference<'a, E>;
768
769    fn next(&mut self) -> Option<Self::Item> {
770        //      type        direction    |    iterate over    reverse
771        // ------------------------------|------------------------------
772        //    Directed      Outgoing     |      outgoing        no
773        //    Directed      Incoming     |      incoming        no
774        //   Undirected     Outgoing     |        both       incoming
775        //   Undirected     Incoming     |        both       outgoing
776
777        // For `iterate_over`, "both" is represented as `None`.
778        // For `reverse`, "no" is represented as `None`.
779        let (iterate_over, _reverse) = (None, Some(self.direction.opposite()));
780
781        if iterate_over.unwrap_or(EdgeDirection::Outgoing) == EdgeDirection::Outgoing {
782            let i = self.next[0].index();
783            if let Some(Edge { weight, next, .. }) = self.edges.get(i) {
784                self.next[0] = next[0];
785                return Some(EdgeReference {
786                    index: EdgeIndex(i as u32),
787                    weight,
788                });
789            }
790        }
791
792        if iterate_over.unwrap_or(EdgeDirection::Incoming) == EdgeDirection::Incoming {
793            while let Some(Edge { node, weight, next }) = self.edges.get(self.next[1].index()) {
794                let edge_index = self.next[1];
795                self.next[1] = next[1];
796
797                // In any of the "both" situations, self-loops would be iterated over twice.
798                // Skip them here.
799                if iterate_over.is_none() && node[0] == self.skip_start {
800                    continue;
801                }
802
803                return Some(EdgeReference {
804                    index: edge_index,
805                    weight,
806                });
807            }
808        }
809
810        None
811    }
812}
813
814impl<E> Clone for Edges<'_, E> {
815    fn clone(&self) -> Self {
816        Edges {
817            skip_start: self.skip_start,
818            edges: self.edges,
819            next: self.next,
820            direction: self.direction,
821        }
822    }
823}
824
825/// An iterator over mutable references to all edges from or to a node.
826pub struct EdgesMut<'a, N, E> {
827    graph: &'a mut UnGraph<N, E>,
828    incoming_edge: Option<EdgeIndex>,
829    outgoing_edge: Option<EdgeIndex>,
830}
831
832impl<'a, N: Copy, E> Iterator for EdgesMut<'a, N, E> {
833    type Item = EdgeMut<'a, E>;
834
835    #[inline]
836    fn next(&mut self) -> Option<EdgeMut<'a, E>> {
837        if let Some(edge) = self.incoming_edge {
838            self.incoming_edge = self.graph.next_edge(edge, EdgeDirection::Incoming);
839            let weights = &mut self.graph[edge];
840            return Some(EdgeMut {
841                index: edge,
842                weight: unsafe { core::mem::transmute::<&mut E, &'a mut E>(weights) },
843            });
844        }
845
846        let edge = self.outgoing_edge?;
847        self.outgoing_edge = self.graph.next_edge(edge, EdgeDirection::Outgoing);
848        let weights = &mut self.graph[edge];
849        Some(EdgeMut {
850            index: edge,
851            weight: unsafe { core::mem::transmute::<&mut E, &'a mut E>(weights) },
852        })
853    }
854}
855
856// TODO: Reduce duplication between `Edges` variants.
857/// An iterator over edge weights for edges from or to a node.
858pub struct EdgeWeights<'a, E: 'a> {
859    /// The starting node to skip over.
860    skip_start: NodeIndex,
861
862    /// The edges to iterate over.
863    edges: &'a [Edge<E>],
864
865    /// The next edge to visit.
866    next: [EdgeIndex; 2],
867
868    /// The direction of edges.
869    direction: EdgeDirection,
870}
871
872impl<'a, E> Iterator for EdgeWeights<'a, E> {
873    type Item = &'a E;
874
875    fn next(&mut self) -> Option<Self::Item> {
876        //      type        direction    |    iterate over    reverse
877        // ------------------------------|------------------------------
878        //    Directed      Outgoing     |      outgoing        no
879        //    Directed      Incoming     |      incoming        no
880        //   Undirected     Outgoing     |        both       incoming
881        //   Undirected     Incoming     |        both       outgoing
882
883        // For `iterate_over`, "both" is represented as `None`.
884        // For `reverse`, "no" is represented as `None`.
885        let (iterate_over, _reverse) = (None, Some(self.direction.opposite()));
886
887        if iterate_over.unwrap_or(EdgeDirection::Outgoing) == EdgeDirection::Outgoing {
888            let i = self.next[0].index();
889            if let Some(Edge { weight, next, .. }) = self.edges.get(i) {
890                self.next[0] = next[0];
891                return Some(weight);
892            }
893        }
894
895        if iterate_over.unwrap_or(EdgeDirection::Incoming) == EdgeDirection::Incoming {
896            while let Some(Edge { node, weight, next }) = self.edges.get(self.next[1].index()) {
897                self.next[1] = next[1];
898
899                // In any of the "both" situations, self-loops would be iterated over twice.
900                // Skip them here.
901                if iterate_over.is_none() && node[0] == self.skip_start {
902                    continue;
903                }
904
905                return Some(weight);
906            }
907        }
908
909        None
910    }
911}
912
913impl<E> Clone for EdgeWeights<'_, E> {
914    fn clone(&self) -> Self {
915        EdgeWeights {
916            skip_start: self.skip_start,
917            edges: self.edges,
918            next: self.next,
919            direction: self.direction,
920        }
921    }
922}
923
924/// An iterator over mutable references to all edge weights from or to a node.
925pub struct EdgeWeightsMut<'a, N, E> {
926    /// A mutable reference to the graph.
927    pub graph: &'a mut UnGraph<N, E>,
928
929    /// The next incoming edge to visit.
930    pub incoming_edge: Option<EdgeIndex>,
931
932    /// The next outgoing edge to visit.
933    pub outgoing_edge: Option<EdgeIndex>,
934}
935
936impl<'a, N: Copy, E> Iterator for EdgeWeightsMut<'a, N, E> {
937    type Item = &'a mut E;
938
939    #[inline]
940    fn next(&mut self) -> Option<&'a mut E> {
941        if let Some(edge) = self.incoming_edge {
942            self.incoming_edge = self.graph.next_edge(edge, EdgeDirection::Incoming);
943            let weight = &mut self.graph[edge];
944            return Some(unsafe { core::mem::transmute::<&mut E, &'a mut E>(weight) });
945        }
946
947        let edge = self.outgoing_edge?;
948        self.outgoing_edge = self.graph.next_edge(edge, EdgeDirection::Outgoing);
949        let weight = &mut self.graph[edge];
950        Some(unsafe { core::mem::transmute::<&mut E, &'a mut E>(weight) })
951    }
952}
953
954/// An iterator yielding immutable access to all edge weights.
955pub struct AllEdgeWeights<'a, E: 'a> {
956    edges: core::slice::Iter<'a, Edge<E>>,
957}
958
959impl<'a, E> Iterator for AllEdgeWeights<'a, E> {
960    type Item = &'a E;
961
962    fn next(&mut self) -> Option<&'a E> {
963        self.edges.next().map(|edge| &edge.weight)
964    }
965
966    fn size_hint(&self) -> (usize, Option<usize>) {
967        self.edges.size_hint()
968    }
969}
970
971/// An iterator yielding mutable access to all edge weights.
972#[derive(Debug)]
973pub struct AllEdgeWeightsMut<'a, E: 'a> {
974    edges: core::slice::IterMut<'a, Edge<E>>,
975}
976
977impl<'a, E> Iterator for AllEdgeWeightsMut<'a, E> {
978    type Item = &'a mut E;
979
980    fn next(&mut self) -> Option<&'a mut E> {
981        self.edges.next().map(|edge| &mut edge.weight)
982    }
983
984    fn size_hint(&self) -> (usize, Option<usize>) {
985        self.edges.size_hint()
986    }
987}
988
989struct EdgesWalkerMut<'a, E: 'a> {
990    edges: &'a mut [Edge<E>],
991    next: EdgeIndex,
992    dir: EdgeDirection,
993}
994
995impl<E> EdgesWalkerMut<'_, E> {
996    fn next_edge(&mut self) -> Option<&mut Edge<E>> {
997        self.next().map(|t| t.1)
998    }
999
1000    fn next(&mut self) -> Option<(EdgeIndex, &mut Edge<E>)> {
1001        let this_index = self.next;
1002        let k = self.dir as usize;
1003        match self.edges.get_mut(self.next.index()) {
1004            None => None,
1005            Some(edge) => {
1006                self.next = edge.next[k];
1007                Some((this_index, edge))
1008            }
1009        }
1010    }
1011}
1012
1013/// Indexes the `UnGraph` by `NodeIndex` to access node weights.
1014///
1015/// # Panics
1016///
1017/// Panics if the node doesn't exist.
1018impl<N, E> Index<NodeIndex> for UnGraph<N, E> {
1019    type Output = N;
1020    fn index(&self, index: NodeIndex) -> &N {
1021        &self.nodes[index.index()].weight
1022    }
1023}
1024
1025/// Indexes the `UnGraph` by `NodeIndex` to access node weights.
1026///
1027/// # Panics
1028///
1029/// Panics if the node doesn't exist.
1030impl<N, E> IndexMut<NodeIndex> for UnGraph<N, E> {
1031    fn index_mut(&mut self, index: NodeIndex) -> &mut N {
1032        &mut self.nodes[index.index()].weight
1033    }
1034}
1035
1036/// Indexes the `UnGraph` by `EdgeIndex` to access edge weights.
1037///
1038/// # Panics
1039///
1040/// Panics if the edge doesn't exist.
1041impl<N, E> Index<EdgeIndex> for UnGraph<N, E> {
1042    type Output = E;
1043    fn index(&self, index: EdgeIndex) -> &E {
1044        &self.edges[index.index()].weight
1045    }
1046}
1047
1048/// Indexes the `UnGraph` by `EdgeIndex` to access edge weights.
1049///
1050/// # Panics
1051///
1052/// Panics if the edge doesn't exist.
1053impl<N, E> IndexMut<EdgeIndex> for UnGraph<N, E> {
1054    fn index_mut(&mut self, index: EdgeIndex) -> &mut E {
1055        &mut self.edges[index.index()].weight
1056    }
1057}
1058
1059/// A reference to a graph edge.
1060#[derive(Debug)]
1061pub struct EdgeReference<'a, E: 'a> {
1062    index: EdgeIndex,
1063    weight: &'a E,
1064}
1065
1066impl<'a, E: 'a> EdgeReference<'a, E> {
1067    /// Returns the index of the edge.
1068    #[inline]
1069    pub fn index(&self) -> EdgeIndex {
1070        self.index
1071    }
1072
1073    /// Returns the weight of the edge.
1074    #[inline]
1075    pub fn weight(&self) -> &'a E {
1076        self.weight
1077    }
1078}
1079
1080impl<E> Clone for EdgeReference<'_, E> {
1081    fn clone(&self) -> Self {
1082        *self
1083    }
1084}
1085
1086impl<E> Copy for EdgeReference<'_, E> {}
1087
1088impl<E> PartialEq for EdgeReference<'_, E>
1089where
1090    E: PartialEq,
1091{
1092    fn eq(&self, rhs: &Self) -> bool {
1093        self.index == rhs.index && self.weight == rhs.weight
1094    }
1095}
1096
1097/// A mutable reference to a graph edge.
1098#[derive(Debug)]
1099pub struct EdgeMut<'a, E: 'a> {
1100    index: EdgeIndex,
1101    weight: &'a mut E,
1102}
1103
1104impl<E> EdgeMut<'_, E> {
1105    /// Returns the index of the edge.
1106    #[inline]
1107    pub fn index(&self) -> EdgeIndex {
1108        self.index
1109    }
1110
1111    /// Returns the weight of the edge.
1112    #[inline]
1113    pub fn weight(&self) -> &E {
1114        self.weight
1115    }
1116
1117    /// Returns the weight of the edge mutably.
1118    #[inline]
1119    pub fn weight_mut(&mut self) -> &mut E {
1120        self.weight
1121    }
1122}
1123
1124impl<E> PartialEq for EdgeMut<'_, E>
1125where
1126    E: PartialEq,
1127{
1128    fn eq(&self, rhs: &Self) -> bool {
1129        self.index == rhs.index && self.weight == rhs.weight
1130    }
1131}