Impl node_weights_mut and edge_weights_mut methods for StableGraph
diff --git a/src/graph_impl/stable_graph/mod.rs b/src/graph_impl/stable_graph/mod.rs
index 6abfa99..142aaeb 100644
--- a/src/graph_impl/stable_graph/mod.rs
+++ b/src/graph_impl/stable_graph/mod.rs
@@ -474,6 +474,16 @@
}
}
+ /// Return an iterator yielding mutable access to all node weights.
+ ///
+ /// The order in which weights are yielded matches the order of their node
+ /// indices.
+ pub fn node_weights_mut(&mut self) -> impl Iterator<Item = &mut N> {
+ self.g
+ .node_weights_mut()
+ .flat_map(|maybe_node| maybe_node.iter_mut())
+ }
+
/// Return an iterator over the node indices of the graph
pub fn node_indices(&self) -> NodeIndices<N, Ix> {
NodeIndices {
@@ -501,6 +511,16 @@
}
}
+ /// Return an iterator yielding mutable access to all edge weights.
+ ///
+ /// The order in which weights are yielded matches the order of their edge
+ /// indices.
+ pub fn edge_weights_mut(&mut self) -> impl Iterator<Item = &mut E> {
+ self.g
+ .edge_weights_mut()
+ .flat_map(|maybe_edge| maybe_edge.iter_mut())
+ }
+
/// Access the source and target nodes for `e`.
pub fn edge_endpoints(&self, e: EdgeIndex<Ix>) -> Option<(NodeIndex<Ix>, NodeIndex<Ix>)> {
match self.g.edges.get(e.index()) {
diff --git a/tests/stable_graph.rs b/tests/stable_graph.rs
index 9ac401e..2ca04ba 100644
--- a/tests/stable_graph.rs
+++ b/tests/stable_graph.rs
@@ -375,3 +375,35 @@
}
let _ = StableGraph::<(), (), Undirected, usize>::from_elements(min_spanning_tree(&g));
}
+
+#[test]
+fn weights_mut_iterator() {
+ let mut gr = StableGraph::new();
+ let a = gr.add_node(1);
+ let b = gr.add_node(2);
+ let c = gr.add_node(3);
+ let e1 = gr.add_edge(a, a, 10);
+ let e2 = gr.add_edge(a, b, 20);
+ let e3 = gr.add_edge(b, c, 30);
+ let e4 = gr.add_edge(a, c, 40);
+
+ for n in gr.node_weights_mut() {
+ *n += 1;
+ }
+ assert_eq!(gr[a], 2);
+ assert_eq!(gr[b], 3);
+ assert_eq!(gr[c], 4);
+
+ for e in gr.edge_weights_mut() {
+ *e -= 1;
+ }
+ assert_eq!(gr[e1], 9);
+ assert_eq!(gr[e2], 19);
+ assert_eq!(gr[e3], 29);
+ assert_eq!(gr[e4], 39);
+
+ // test on deletion
+ gr.remove_node(b);
+ assert_eq!(gr.node_weights_mut().count(), gr.node_count());
+ assert_eq!(gr.edge_weights_mut().count(), gr.edge_count());
+}
