Move barbell_graph generator to rustworkx-core

rustworkx-core/src/generators/barbell_graph.rs

@@ -0,0 +1,239 @@
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations
+// under the License.
+
+use std::hash::Hash;
+
+use petgraph::data::{Build, Create};
+use petgraph::visit::{Data, NodeIndexable, NodeRef};
+
+use super::utils::get_num_nodes;
+use super::utils::pairwise;
+use super::InvalidInputError;
+
+/// Generate an undirected barbell graph where two identical mesh graphs are
+/// connected by a path.
+///
+/// .. note::
+///
+///   If neither `num_path_nodes` nor `path_weights` (described below) is not specified
+///   then this is equivalent to two mesh graphs joined together.
+///
+/// Arguments:
+///
+/// * `num_mesh_nodes` - The number of nodes to generate the mesh graph
+///     with. Node weights will be None if this is specified. If both
+///     `num_mesh_nodes` and `mesh_weights` are set this will be ignored and
+///     `mesh_weights` will be used.
+/// * `num_path_nodes` - The number of nodes to generate the path
+///     with. Node weights will be None if this is specified. If both
+///     `num_path_nodes` and `path_weights` are set this will be ignored and
+///     `path_weights` will be used.
+/// * `mesh_weights` - A list of node weights for the mesh graph. If both
+///     `num_mesh_nodes` and `mesh_weights` are set `num_mesh_nodes` will
+///     be ignored and `mesh_weights` will be used.
+/// * `path_weights` - A list of node weights for the path. If both
+///     `num_path_nodes` and `path_weights` are set `num_path_nodes` will
+///     be ignored and `path_weights` will be used.
+/// * `default_node_weight` - A callable that will return the weight to use
+///     for newly created nodes. This is ignored if weights are specified,
+///     as the weights from that argument will be used instead.
+/// * `default_edge_weight` - A callable that will return the weight object
+///     to use for newly created edges.
+///
+/// # Example
+/// ```rust
+/// use rustworkx_core::petgraph;
+/// use rustworkx_core::generators::barbell_graph;
+/// use rustworkx_core::petgraph::visit::EdgeRef;
+///
+/// let expected_edge_list = vec![
+///     (0, 1),
+///     (0, 2),
+///     (0, 3),
+///     (1, 2),
+///     (1, 3),
+///     (2, 3),
+///     (3, 4),
+///     (4, 5),
+///     (5, 6),
+///     (6, 7),
+///     (7, 8),
+///     (7, 9),
+///     (7, 10),
+///     (8, 9),
+///     (8, 10),
+///     (9, 10),
+/// ];
+/// let g: petgraph::graph::UnGraph<(), ()> = barbell_graph(
+///     Some(4),
+///     Some(3),
+///     None,
+///     None,
+///     || {()},
+///     || {()},
+/// ).unwrap();
+/// assert_eq!(
+///     expected_edge_list,
+///     g.edge_references()
+///         .map(|edge| (edge.source().index(), edge.target().index()))
+///         .collect::<Vec<(usize, usize)>>(),
+/// )
+/// ```
+pub fn barbell_graph<G, T, F, H, M>(
+    num_mesh_nodes: Option<usize>,
+    num_path_nodes: Option<usize>,
+    mesh_weights: Option<Vec<T>>,
+    path_weights: Option<Vec<T>>,
+    mut default_node_weight: F,
+    mut default_edge_weight: H,
+) -> Result<G, InvalidInputError>
+where
+    G: Build + Create + Data<NodeWeight = T, EdgeWeight = M> + NodeIndexable,
+    F: FnMut() -> T,
+    H: FnMut() -> M,
+    G::NodeId: Eq + Hash + NodeRef,
+    T: Clone,
+{
+    if num_mesh_nodes.is_none() && mesh_weights.is_none() {
+        return Err(InvalidInputError {});
+    }
+    let num_nodes = get_num_nodes(&num_mesh_nodes, &mesh_weights);
+    let num_edges = (num_nodes * (num_nodes - 1)) / 2;
+    let mesh_weights_2 = mesh_weights.clone();
+
+    let mut graph = G::with_capacity(num_nodes, num_edges);
+
+    // Create left mesh graph
+    let mesh_nodes: Vec<G::NodeId> = match mesh_weights {
+        Some(mesh_weights) => mesh_weights
+            .into_iter()
+            .map(|weight| graph.add_node(weight))
+            .collect(),
+        None => (0..num_mesh_nodes.unwrap())
+            .map(|_| graph.add_node(default_node_weight()))
+            .collect(),
+    };
+
+    let meshlen = mesh_nodes.len();
+    for i in 0..meshlen - 1 {
+        for j in i + 1..meshlen {
+            graph.add_edge(mesh_nodes[i], mesh_nodes[j], default_edge_weight());
+        }
+    }
+
+    // Add path nodes and edges
+    let path_nodes: Vec<G::NodeId> = match path_weights {
+        Some(path_weights) => path_weights
+            .into_iter()
+            .map(|weight| graph.add_node(weight))
+            .collect(),
+        None => {
+            if let Some(num_path) = num_path_nodes {
+                (0..num_path)
+                    .map(|_| graph.add_node(default_node_weight()))
+                    .collect()
+            } else {
+                vec![]
+            }
+        }
+    };
+    let pathlen = path_nodes.len();
+    if pathlen > 0 {
+        graph.add_edge(
+            graph.from_index(meshlen - 1),
+            graph.from_index(meshlen),
+            default_edge_weight(),
+        );
+        for (node_a, node_b) in pairwise(path_nodes) {
+            match node_a {
+                Some(node_a) => graph.add_edge(node_a, node_b, default_edge_weight()),
+                None => continue,
+            };
+        }
+    }
+
+    // Add right mesh graph
+    let mesh_nodes: Vec<G::NodeId> = match mesh_weights_2 {
+        Some(mesh_weights_2) => mesh_weights_2
+            .into_iter()
+            .map(|weight| graph.add_node(weight))
+            .collect(),
+        None => (0..num_mesh_nodes.unwrap())
+            .map(|_| graph.add_node(default_node_weight()))
+            .collect(),
+    };
+    let meshlen = mesh_nodes.len();
+
+    graph.add_edge(
+        graph.from_index(meshlen + pathlen - 1),
+        graph.from_index(meshlen + pathlen),
+        default_edge_weight(),
+    );
+    for i in 0..meshlen - 1 {
+        for j in i + 1..meshlen {
+            graph.add_edge(mesh_nodes[i], mesh_nodes[j], default_edge_weight());
+        }
+    }
+    Ok(graph)
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::generators::barbell_graph;
+    use crate::generators::InvalidInputError;
+    use crate::petgraph::visit::EdgeRef;
+
+    #[test]
+    fn test_barbell_mesh_path() {
+        let expected_edge_list = vec![
+            (0, 1),
+            (0, 2),
+            (0, 3),
+            (1, 2),
+            (1, 3),
+            (2, 3),
+            (3, 4),
+            (4, 5),
+            (5, 6),
+            (6, 7),
+            (7, 8),
+            (7, 9),
+            (7, 10),
+            (8, 9),
+            (8, 10),
+            (9, 10),
+        ];
+        let g: petgraph::graph::UnGraph<(), ()> =
+            barbell_graph(Some(4), Some(3), None, None, || (), || ()).unwrap();
+        assert_eq!(
+            expected_edge_list,
+            g.edge_references()
+                .map(|edge| (edge.source().index(), edge.target().index()))
+                .collect::<Vec<(usize, usize)>>(),
+        );
+    }
+
+    #[test]
+    fn test_barbell_none_mesh() {
+        match barbell_graph::<petgraph::graph::UnGraph<(), ()>, (), _, _, ()>(
+            None,
+            None,
+            None,
+            None,
+            || (),
+            || (),
+        ) {
+            Ok(_) => panic!("Returned a non-error"),
+            Err(e) => assert_eq!(e, InvalidInputError),
+        };
+    }
+}

rustworkx-core/src/generators/mod.rs

@@ -12,6 +12,7 @@
 
 //! This module contains generator functions for building graphs
 
+mod barbell_graph;
 mod binomial_tree_graph;
 mod cycle_graph;
 mod grid_graph;
@@ -39,6 +40,7 @@ impl fmt::Display for InvalidInputError {
     }
 }
 
+pub use barbell_graph::barbell_graph;
 pub use binomial_tree_graph::binomial_tree_graph;
 pub use cycle_graph::cycle_graph;
 pub use grid_graph::grid_graph;

rustworkx-core/src/generators/utils.rs

@@ -10,6 +10,8 @@
 // License for the specific language governing permissions and limitations
 // under the License.
 
+use std::iter;
+
 #[inline]
 pub fn get_num_nodes<T>(num_nodes: &Option<usize>, weights: &Option<Vec<T>>) -> usize {
     if weights.is_some() {
@@ -18,3 +20,11 @@ pub fn get_num_nodes<T>(num_nodes: &Option<usize>, weights: &Option<Vec<T>>) ->
         num_nodes.unwrap()
     }
 }
+
+pub fn pairwise<I>(right: I) -> impl Iterator<Item = (Option<I::Item>, I::Item)>
+where
+    I: IntoIterator + Clone,
+{
+    let left = iter::once(None).chain(right.clone().into_iter().map(Some));
+    left.zip(right)
+}

src/generators.rs

@@ -16,7 +16,6 @@ use std::iter;
 use petgraph::algo;
 use petgraph::graph::NodeIndex;
 use petgraph::prelude::*;
-use petgraph::visit::{EdgeRef, IntoEdgeReferences};
 use petgraph::Undirected;
 
 use pyo3::exceptions::{PyIndexError, PyOverflowError};
@@ -1456,6 +1455,12 @@ pub fn generalized_petersen_graph(
 ///     :class:`~rustworkx.PyGraph` object will not be not be a multigraph and
 ///     won't  allow parallel edges to be added. Instead
 ///     calls which would create a parallel edge will update the existing edge.
+/// :param list mesh_weights: A list of node weights for the mesh graph. If both
+///     ``num_mesh_nodes`` and ``mesh_weights`` are set ``num_mesh_nodes`` will
+///     be ignored and ``mesh_weights`` will be used.
+/// :param list path_weights: A list of node weights for the path. If both
+///     ``num_path_nodes`` and ``path_weights`` are set ``num_path_nodes`` will
+///     be ignored and ``path_weights`` will be used.
 ///
 /// :returns: The generated barbell graph
 /// :rtype: PyGraph
@@ -1476,62 +1481,27 @@ pub fn barbell_graph(
     num_mesh_nodes: Option<usize>,
     num_path_nodes: Option<usize>,
     multigraph: bool,
+    mesh_weights: Option<Vec<PyObject>>,
+    path_weights: Option<Vec<PyObject>>,
 ) -> PyResult<graph::PyGraph> {
-    if num_mesh_nodes.is_none() {
-        return Err(PyIndexError::new_err("num_mesh_nodes not specified"));
-    }
-
-    let mut left_mesh = StableUnGraph::<PyObject, PyObject>::default();
-    let mesh_nodes: Vec<NodeIndex> = (0..num_mesh_nodes.unwrap())
-        .map(|_| left_mesh.add_node(py.None()))
-        .collect();
-    let mut nodelen = mesh_nodes.len();
-    for i in 0..nodelen - 1 {
-        for j in i + 1..nodelen {
-            left_mesh.add_edge(mesh_nodes[i], mesh_nodes[j], py.None());
-        }
-    }
-
-    let right_mesh = left_mesh.clone();
-
-    if let Some(num_nodes) = num_path_nodes {
-        let path_nodes: Vec<NodeIndex> = (0..num_nodes)
-            .map(|_| left_mesh.add_node(py.None()))
-            .collect();
-        left_mesh.add_edge(
-            NodeIndex::new(nodelen - 1),
-            NodeIndex::new(nodelen),
-            py.None(),
-        );
-
-        nodelen += path_nodes.len();
-
-        for (node_a, node_b) in pairwise(path_nodes) {
-            match node_a {
-                Some(node_a) => left_mesh.add_edge(node_a, node_b, py.None()),
-                None => continue,
-            };
+    let default_fn = || py.None();
+    let graph: StablePyGraph<Undirected> = match core_generators::barbell_graph(
+        num_mesh_nodes,
+        num_path_nodes,
+        mesh_weights,
+        path_weights,
+        default_fn,
+        default_fn,
+    ) {
+        Ok(graph) => graph,
+        Err(_) => {
+            return Err(PyIndexError::new_err(
+                "num_nodes and weights list not specified",
+            ))
         }
-    }
-
-    for node in right_mesh.node_indices() {
-        let new_node = &right_mesh[node];
-        left_mesh.add_node(new_node.clone_ref(py));
-    }
-    left_mesh.add_edge(
-        NodeIndex::new(nodelen - 1),
-        NodeIndex::new(nodelen),
-        py.None(),
-    );
-    for edge in right_mesh.edge_references() {
-        let new_source = NodeIndex::new(nodelen + edge.source().index());
-        let new_target = NodeIndex::new(nodelen + edge.target().index());
-        let weight = edge.weight();
-        left_mesh.add_edge(new_source, new_target, weight.clone_ref(py));
-    }
-
+    };
     Ok(graph::PyGraph {
-        graph: left_mesh,
+        graph,
         node_removed: false,
         multigraph,
         attrs: py.None(),