Skip to main content

libtrace2power/
util.rs

1// Copyright (c) 2024-2026 Antmicro <www.antmicro.com>
2// SPDX-License-Identifier: Apache-2.0
3
4use wellen::{self, Hierarchy, Scope, VarRef};
5
6// Wellen has no good iterator over ALL VarRefs, so I made one. The generics are here to deal with
7// hidden iterator types.
8struct VarRefIterator<'w, HIter, F>
9where
10    HIter: Iterator<Item = VarRef> + 'w,
11    F: Fn(&'w Scope) -> HIter,
12{
13    hierarchy: &'w Hierarchy,
14    get_iter_of_scope: F,
15    scopes: Vec<&'w Scope>,
16    iter: Option<HIter>,
17}
18
19impl<'w, HIter, F> Iterator for VarRefIterator<'w, HIter, F>
20where
21    HIter: Iterator<Item = VarRef> + 'w,
22    F: Fn(&'w Scope) -> HIter,
23{
24    type Item = VarRef;
25    fn next(&mut self) -> Option<Self::Item> {
26        match self.iter {
27            Some(ref mut it) => match it.next() {
28                Some(v) => Some(v),
29                None => {
30                    if let Some(scope) = self.scopes.pop() {
31                        self.iter = Some((self.get_iter_of_scope)(scope));
32                        let mut siter = scope.scopes(self.hierarchy);
33                        self.scopes
34                            .extend(siter.by_ref().map(|s| &self.hierarchy[s]));
35                        self.next()
36                    } else {
37                        None
38                    }
39                }
40            },
41            None => None,
42        }
43    }
44}
45
46pub trait VarRefsIter {
47    fn var_refs_iter<'s>(&'s self) -> impl Iterator<Item = VarRef> + 's;
48}
49
50impl VarRefsIter for Hierarchy {
51    fn var_refs_iter<'s>(&'s self) -> impl Iterator<Item = VarRef> + 's {
52        VarRefIterator {
53            hierarchy: &self,
54            scopes: self.scopes().map(|s| &self[s]).collect(),
55            get_iter_of_scope: |s: &Scope| s.vars(self),
56            iter: match self.first_scope() {
57                Some(s) => Some(s.vars(self)),
58                None => None,
59            },
60        }
61    }
62}