Fix bfs_search and other search methods panicking with invalid sources (Qiskit#1388)

* Validate source nodes in search functions

* Add tests for graph searches

* Add tests for digraph searches

* Black

* Add release notes

* Address clippy

Author: IvanIsCoding

releasenotes/notes/fix-search-panic-fd54b704f7ee2101.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    Fixed a panic when passing invalid source nodes to search methods
+    such as :func:`~rustworkx.bfs_search`. See
+    `#1386 <https://github.com/Qiskit/rustworkx/issues/1386>`__ for more details.

src/traversal/mod.rs

@@ -30,13 +30,30 @@ use std::convert::TryFrom;
 
 use hashbrown::HashSet;
 
-use pyo3::exceptions::PyTypeError;
+use pyo3::exceptions::{PyIndexError, PyTypeError};
 use pyo3::prelude::*;
 use pyo3::Python;
 
 use petgraph::graph::NodeIndex;
+use petgraph::EdgeType;
 
 use crate::iterators::EdgeList;
+use crate::StablePyGraph;
+
+fn validate_source_nodes<Ty: EdgeType>(
+    graph: &StablePyGraph<Ty>,
+    starts: &[NodeIndex],
+) -> PyResult<()> {
+    for index in starts.iter() {
+        if !graph.contains_node(*index) {
+            return Err(PyIndexError::new_err(format!(
+                "Node source index \"{}\" out of graph bound",
+                index.index()
+            )));
+        }
+    }
+    Ok(())
+}
 
 /// Get an edge list of the tree edges from a depth-first traversal
 ///
@@ -386,6 +403,8 @@ pub fn digraph_bfs_search(
         None => graph.graph.node_indices().collect(),
     };
 
+    validate_source_nodes(&graph.graph, &starts)?;
+
     breadth_first_search(&graph.graph, starts, |event| {
         bfs_handler(py, &visitor, event)
     })?;
@@ -530,6 +549,8 @@ pub fn graph_bfs_search(
         None => graph.graph.node_indices().collect(),
     };
 
+    validate_source_nodes(&graph.graph, &starts)?;
+
     breadth_first_search(&graph.graph, starts, |event| {
         bfs_handler(py, &visitor, event)
     })?;
@@ -644,6 +665,8 @@ pub fn digraph_dfs_search(
         None => graph.graph.node_indices().collect(),
     };
 
+    validate_source_nodes(&graph.graph, &starts)?;
+
     depth_first_search(&graph.graph, starts, |event| {
         dfs_handler(py, &visitor, event)
     })?;
@@ -758,6 +781,8 @@ pub fn graph_dfs_search(
         None => graph.graph.node_indices().collect(),
     };
 
+    validate_source_nodes(&graph.graph, &starts)?;
+
     depth_first_search(&graph.graph, starts, |event| {
         dfs_handler(py, &visitor, event)
     })?;
@@ -895,6 +920,8 @@ pub fn digraph_dijkstra_search(
         None => graph.graph.node_indices().collect(),
     };
 
+    validate_source_nodes(&graph.graph, &starts)?;
+
     let edge_cost_fn = CostFn::try_from((weight_fn, 1.0))?;
     dijkstra_search(
         &graph.graph,
@@ -1036,6 +1063,8 @@ pub fn graph_dijkstra_search(
         None => graph.graph.node_indices().collect(),
     };
 
+    validate_source_nodes(&graph.graph, &starts)?;
+
     let edge_cost_fn = CostFn::try_from((weight_fn, 1.0))?;
     dijkstra_search(
         &graph.graph,

tests/digraph/test_bfs_search.py

@@ -160,3 +160,8 @@ def gray_target_edge(self, _):
 
         vis = PruneGrayTargetEdge()
         rustworkx.digraph_bfs_search(self.graph, [0], vis)
+
+    def test_invalid_source(self):
+        graph = rustworkx.PyDiGraph()
+        with self.assertRaises(IndexError):
+            rustworkx.bfs_search(graph, [1], rustworkx.visit.BFSVisitor())

tests/digraph/test_dfs_search.py

@@ -104,3 +104,8 @@ def reconstruct_path(self):
         except rustworkx.visit.StopSearch:
             pass
         self.assertEqual(vis.reconstruct_path(), [0, 2, 5, 3])
+
+    def test_invalid_source(self):
+        graph = rustworkx.PyDiGraph()
+        with self.assertRaises(IndexError):
+            rustworkx.dfs_search(graph, [1], rustworkx.visit.DFSVisitor())

tests/digraph/test_dijkstra_search.py

@@ -187,3 +187,8 @@ def edge_not_relaxed(self, _):
 
         vis = PruneEdgeNotRelaxed()
         rustworkx.digraph_dijkstra_search(self.graph, [0], float, vis)
+
+    def test_invalid_source(self):
+        graph = rustworkx.PyDiGraph()
+        with self.assertRaises(IndexError):
+            rustworkx.dijkstra_search(graph, [1], float, rustworkx.visit.DijkstraVisitor())

tests/graph/test_bfs_search.py

@@ -160,3 +160,8 @@ def gray_target_edge(self, _):
 
         vis = PruneGrayTargetEdge()
         rustworkx.graph_bfs_search(self.graph, [0], vis)
+
+    def test_invalid_source(self):
+        graph = rustworkx.PyGraph()
+        with self.assertRaises(IndexError):
+            rustworkx.bfs_search(graph, [1], rustworkx.visit.BFSVisitor())

tests/graph/test_dfs_search.py

@@ -104,3 +104,8 @@ def reconstruct_path(self):
         except rustworkx.visit.StopSearch:
             pass
         self.assertEqual(vis.reconstruct_path(), [0, 2, 5, 3])
+
+    def test_invalid_source(self):
+        graph = rustworkx.PyGraph()
+        with self.assertRaises(IndexError):
+            rustworkx.dfs_search(graph, [1], rustworkx.visit.DFSVisitor())

tests/graph/test_dijkstra_search.py

@@ -187,3 +187,8 @@ def edge_not_relaxed(self, _):
 
         vis = PruneEdgeNotRelaxed()
         rustworkx.graph_dijkstra_search(self.graph, [0], float, vis)
+
+    def test_invalid_source(self):
+        graph = rustworkx.PyGraph()
+        with self.assertRaises(IndexError):
+            rustworkx.dijkstra_search(graph, [1], float, rustworkx.visit.DijkstraVisitor())