Skip to main content

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 *raw* (graph-direct) neighbors of a component.
56///
57/// Returned by [`ComponentGraph::raw_predecessors`] and
58/// [`ComponentGraph::raw_successors`]. Yields every node connected by an
59/// edge, including pass-through categories.
60pub struct RawNeighbors<'a, N>
61where
62    N: Node,
63{
64    pub(crate) graph: &'a DiGraph<N, ()>,
65    pub(crate) iter: petgraph::graph::Neighbors<'a, ()>,
66}
67
68impl<'a, N> Iterator for RawNeighbors<'a, N>
69where
70    N: Node,
71{
72    type Item = &'a N;
73
74    fn next(&mut self) -> Option<Self::Item> {
75        self.iter.next().map(|i| &self.graph[i])
76    }
77}
78
79/// An iterator over the *effective* (pass-through-aware) neighbors of a
80/// component.
81///
82/// Returned by [`ComponentGraph::predecessors`] and
83/// [`ComponentGraph::successors`]. Yields only non-pass-through ancestors
84/// or descendants, walking transparently past pass-through nodes in the
85/// chain. Eagerly collected at construction time.
86pub struct Neighbors<'a, N>
87where
88    N: Node,
89{
90    pub(crate) iter: IntoIter<&'a N>,
91}
92
93impl<'a, N> Iterator for Neighbors<'a, N>
94where
95    N: Node,
96{
97    type Item = &'a N;
98
99    fn next(&mut self) -> Option<Self::Item> {
100        self.iter.next()
101    }
102}
103
104/// An iterator over the siblings of a component in a `ComponentGraph`.
105pub struct Siblings<'a, N>
106where
107    N: Node,
108{
109    pub(crate) component_id: u64,
110    pub(crate) iter: Flatten<IntoIter<Neighbors<'a, N>>>,
111    visited: HashSet<u64>,
112}
113
114impl<'a, N> Siblings<'a, N>
115where
116    N: Node,
117{
118    pub(crate) fn new(component_id: u64, iter: Flatten<IntoIter<Neighbors<'a, N>>>) -> Self {
119        Siblings {
120            component_id,
121            iter,
122            visited: HashSet::new(),
123        }
124    }
125}
126
127impl<'a, N> Iterator for Siblings<'a, N>
128where
129    N: Node,
130{
131    type Item = &'a N;
132
133    fn next(&mut self) -> Option<Self::Item> {
134        for i in self.iter.by_ref() {
135            if i.component_id() == self.component_id || !self.visited.insert(i.component_id()) {
136                continue;
137            }
138            return Some(i);
139        }
140        None
141    }
142}