Merge pull request #286 from petgraph/graphmap-neighbors-directed
Include self loops in incoming edges
diff --git a/src/graphmap.rs b/src/graphmap.rs
index 8a4d97c..bb76e5e 100644
--- a/src/graphmap.rs
+++ b/src/graphmap.rs
@@ -348,6 +348,7 @@
Some(neigh) => neigh.iter(),
None => [].iter(),
},
+ start_node: a,
dir: dir,
ty: self.ty,
}
@@ -531,6 +532,7 @@
Ty: EdgeType,
{
iter: Iter<'a, (N, CompactDirection)>,
+ start_node: N,
dir: Direction,
ty: PhantomData<Ty>,
}
@@ -543,8 +545,9 @@
fn next(&mut self) -> Option<N> {
if Ty::is_directed() {
let self_dir = self.dir;
+ let start_node = self.start_node;
(&mut self.iter)
- .filter_map(move |&(n, dir)| if dir == self_dir {
+ .filter_map(move |&(n, dir)| if dir == self_dir || n == start_node {
Some(n)
} else { None })
.next()
diff --git a/tests/graphmap.rs b/tests/graphmap.rs
index 51b37e6..02362ef 100644
--- a/tests/graphmap.rs
+++ b/tests/graphmap.rs
@@ -298,4 +298,41 @@
for (start, end, weight) in graph.all_edges() {
assert_eq!((start + end) * 2, *weight);
}
-}
\ No newline at end of file
+}
+
+#[test]
+fn neighbors_incoming_includes_self_loops() {
+ let mut graph = DiGraphMap::new();
+
+ let node = graph.add_node(());
+ graph.add_edge(node, node, ());
+
+ let mut neighbors = graph.neighbors_directed(node, Incoming);
+ assert_eq!(neighbors.next(), Some(node));
+ assert_eq!(neighbors.next(), None);
+}
+
+#[test]
+fn undirected_neighbors_includes_self_loops() {
+ let mut graph = UnGraphMap::new();
+
+ let node = graph.add_node(());
+ graph.add_edge(node, node, ());
+
+ let mut neighbors = graph.neighbors(node);
+ assert_eq!(neighbors.next(), Some(node));
+ assert_eq!(neighbors.next(), None);
+}
+
+#[test]
+fn self_loops_can_be_removed() {
+ let mut graph = DiGraphMap::new();
+
+ let node = graph.add_node(());
+ graph.add_edge(node, node, ());
+
+ graph.remove_edge(node, node);
+
+ assert_eq!(graph.neighbors_directed(node, Outgoing).next(), None);
+ assert_eq!(graph.neighbors_directed(node, Incoming).next(), None);
+}
