Add a direction parameter to the find_all method

Earlier the direction was hard coded, which made it unclear that the
method was only finding outgoing nodes.

Signed-off-by: Sahas Subramanian <[email protected]>

Author: shsms

src/graph/formulas/generators/battery.rs

@@ -29,7 +29,12 @@ where
         let inverter_ids = if let Some(battery_ids) = battery_ids {
             Self::find_inverter_ids(graph, &battery_ids)?
         } else {
-            graph.find_all(graph.root_id, |node| node.is_battery_inverter(), false)?
+            graph.find_all(
+                graph.root_id,
+                |node| node.is_battery_inverter(),
+                petgraph::Direction::Outgoing,
+                false,
+            )?
         };
         Ok(Self {
             graph,

src/graph/formulas/generators/chp.rs

@@ -29,7 +29,12 @@ where
         let chp_ids = if let Some(chp_ids) = chp_ids {
             chp_ids
         } else {
-            graph.find_all(graph.root_id, |node| node.is_chp(), false)?
+            graph.find_all(
+                graph.root_id,
+                |node| node.is_chp(),
+                petgraph::Direction::Outgoing,
+                false,
+            )?
         };
         Ok(Self { graph, chp_ids })
     }

src/graph/formulas/generators/consumer.rs

@@ -24,7 +24,12 @@ where
 {
     pub fn try_new(graph: &'a ComponentGraph<N, E>) -> Result<Self, Error> {
         Ok(Self {
-            unvisited_meters: graph.find_all(graph.root_id, |node| node.is_meter(), true)?,
+            unvisited_meters: graph.find_all(
+                graph.root_id,
+                |node| node.is_meter(),
+                petgraph::Direction::Outgoing,
+                true,
+            )?,
             graph,
         })
     }

src/graph/formulas/generators/ev_charger.rs

@@ -29,7 +29,12 @@ where
         let ev_charger_ids = if let Some(ev_charger_ids) = ev_charger_ids {
             ev_charger_ids
         } else {
-            graph.find_all(graph.root_id, |node| node.is_ev_charger(), false)?
+            graph.find_all(
+                graph.root_id,
+                |node| node.is_ev_charger(),
+                petgraph::Direction::Outgoing,
+                false,
+            )?
         };
         Ok(Self {
             graph,

src/graph/formulas/generators/producer.rs

@@ -41,6 +41,7 @@ where
                     || node.is_pv_inverter()
                     || node.is_chp()
             },
+            petgraph::Direction::Outgoing,
             false,
         )? {
             let comp_expr = Self::min_zero(self.graph.fallback_expr([component_id], false)?);

src/graph/formulas/generators/pv.rs

@@ -29,7 +29,12 @@ where
         let pv_inverter_ids = if let Some(pv_inverter_ids) = pv_inverter_ids {
             pv_inverter_ids
         } else {
-            graph.find_all(graph.root_id, |node| node.is_pv_inverter(), false)?
+            graph.find_all(
+                graph.root_id,
+                |node| node.is_pv_inverter(),
+                petgraph::Direction::Outgoing,
+                false,
+            )?
         };
         Ok(Self {
             graph,

src/graph/retrieval.rs

@@ -108,14 +108,15 @@ where
     }
 
     /// Returns a set of all components that match the given predicate, starting
-    /// from the component with the given `component_id`.
+    /// from the component with the given `component_id`, in the given direction.
     ///
     /// If `follow_after_match` is `true`, the search continues deeper beyond
     /// the matching components.
     pub(crate) fn find_all(
         &self,
         from: u64,
         mut pred: impl FnMut(&N) -> bool,
+        direction: petgraph::Direction,
         follow_after_match: bool,
     ) -> Result<BTreeSet<u64>, Error> {
         let index = self.node_indices.get(&from).ok_or_else(|| {
@@ -133,9 +134,7 @@ where
                 }
             }
 
-            let neighbors = self
-                .graph
-                .neighbors_directed(index, petgraph::Direction::Outgoing);
+            let neighbors = self.graph.neighbors_directed(index, direction);
             stack.extend(neighbors);
         }
 
@@ -353,37 +352,55 @@ mod tests {
         let (components, connections) = nodes_and_edges();
         let graph = ComponentGraph::try_new(components.clone(), connections.clone())?;
 
-        let found = graph.find_all(graph.root_id, |x| x.is_meter(), false)?;
+        let found = graph.find_all(
+            graph.root_id,
+            |x| x.is_meter(),
+            petgraph::Direction::Outgoing,
+            false,
+        )?;
         assert_eq!(found, [2].iter().cloned().collect());
 
-        let found = graph.find_all(graph.root_id, |x| x.is_meter(), true)?;
+        let found = graph.find_all(
+            graph.root_id,
+            |x| x.is_meter(),
+            petgraph::Direction::Outgoing,
+            true,
+        )?;
         assert_eq!(found, [2, 3, 6].iter().cloned().collect());
 
         let found = graph.find_all(
             graph.root_id,
             |x| !x.is_grid() && !graph.is_component_meter(x.component_id()).unwrap_or(false),
+            petgraph::Direction::Outgoing,
             true,
         )?;
         assert_eq!(found, [2, 4, 5, 7, 8].iter().cloned().collect());
 
         let found = graph.find_all(
             6,
             |x| !x.is_grid() && !graph.is_component_meter(x.component_id()).unwrap_or(false),
+            petgraph::Direction::Outgoing,
             true,
         )?;
         assert_eq!(found, [7, 8].iter().cloned().collect());
 
         let found = graph.find_all(
             graph.root_id,
             |x| !x.is_grid() && !graph.is_component_meter(x.component_id()).unwrap_or(false),
+            petgraph::Direction::Outgoing,
             false,
         )?;
         assert_eq!(found, [2].iter().cloned().collect());
 
-        let found = graph.find_all(graph.root_id, |_| true, false)?;
+        let found = graph.find_all(
+            graph.root_id,
+            |_| true,
+            petgraph::Direction::Outgoing,
+            false,
+        )?;
         assert_eq!(found, [1].iter().cloned().collect());
 
-        let found = graph.find_all(3, |_| true, true)?;
+        let found = graph.find_all(3, |_| true, petgraph::Direction::Outgoing, true)?;
         assert_eq!(found, [3, 4, 5].iter().cloned().collect());
 
         Ok(())