frequenz_microgrid_component_graph/graph/
iterators.rs

1// License: MIT
2// Copyright © 2024 Frequenz Energy-as-a-Service GmbH
3
4//! Iterators over components and connections in a `ComponentGraph`.
5
6use std::{collections::HashSet, iter::Flatten, vec::IntoIter};
7
8use petgraph::graph::DiGraph;
9
10use crate::{ComponentGraph, Edge, Node};
11
12/// An iterator over the components in a `ComponentGraph`.
13pub struct Components<'a, N>
14where
15    N: Node,
16{
17    pub(crate) iter: std::slice::Iter<'a, petgraph::graph::Node<N>>,
18}
19
20impl<'a, N> Iterator for Components<'a, N>
21where
22    N: Node,
23{
24    type Item = &'a N;
25
26    fn next(&mut self) -> Option<Self::Item> {
27        self.iter.next().map(|n| &n.weight)
28    }
29}
30
31/// An iterator over the connections in a `ComponentGraph`.
32pub struct Connections<'a, N, E>
33where
34    N: Node,
35    E: Edge,
36{
37    pub(crate) cg: &'a ComponentGraph<N, E>,
38    pub(crate) iter: std::slice::Iter<'a, petgraph::graph::Edge<()>>,
39}
40
41impl<'a, N, E> Iterator for Connections<'a, N, E>
42where
43    N: Node,
44    E: Edge,
45{
46    type Item = &'a E;
47
48    fn next(&mut self) -> Option<Self::Item> {
49        self.iter
50            .next()
51            .and_then(|e| self.cg.edges.get(&(e.source(), e.target())))
52    }
53}
54
55/// An iterator over the neighbors of a component in a `ComponentGraph`.
56pub struct Neighbors<'a, N>
57where
58    N: Node,
59{
60    pub(crate) graph: &'a DiGraph<N, ()>,
61    pub(crate) iter: petgraph::graph::Neighbors<'a, ()>,
62}
63
64impl<'a, N> Iterator for Neighbors<'a, N>
65where
66    N: Node,
67{
68    type Item = &'a N;
69
70    fn next(&mut self) -> Option<Self::Item> {
71        self.iter.next().map(|i| &self.graph[i])
72    }
73}
74
75/// An iterator over the siblings of a component in a `ComponentGraph`.
76pub struct Siblings<'a, N>
77where
78    N: Node,
79{
80    pub(crate) component_id: u64,
81    pub(crate) iter: Flatten<IntoIter<Neighbors<'a, N>>>,
82    visited: HashSet<u64>,
83}
84
85impl<'a, N> Siblings<'a, N>
86where
87    N: Node,
88{
89    pub(crate) fn new(component_id: u64, iter: Flatten<IntoIter<Neighbors<'a, N>>>) -> Self {
90        Siblings {
91            component_id,
92            iter,
93            visited: HashSet::new(),
94        }
95    }
96}
97
98impl<'a, N> Iterator for Siblings<'a, N>
99where
100    N: Node,
101{
102    type Item = &'a N;
103
104    fn next(&mut self) -> Option<Self::Item> {
105        for i in self.iter.by_ref() {
106            if i.component_id() == self.component_id || !self.visited.insert(i.component_id()) {
107                continue;
108            }
109            return Some(i);
110        }
111        None
112    }
113}