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, GetItem, Hierarchy, Scope, ScopeRef, 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, SIter, F, S>
9where
10    HIter: Iterator<Item = VarRef> + 'w,
11    SIter: Iterator<Item = ScopeRef> + 'w,
12    F: Fn(&'w Scope) -> HIter,
13    S: Fn(&'w Scope) -> SIter,
14{
15    hierarchy: &'w Hierarchy,
16    get_iter_of_scope: F,
17    get_siter_of_scope: S,
18    scopes: Vec<&'w Scope>,
19    iter: HIter,
20    siter: SIter,
21}
22
23impl<'w, HIter, SIter, F, S> Iterator for VarRefIterator<'w, HIter, SIter, F, S>
24where
25    HIter: Iterator<Item = VarRef> + 'w,
26    SIter: Iterator<Item = ScopeRef> + 'w,
27    F: Fn(&'w Scope) -> HIter,
28    S: Fn(&'w Scope) -> SIter,
29{
30    type Item = VarRef;
31    fn next(&mut self) -> Option<Self::Item> {
32        match self.iter.next() {
33            Some(v) => Some(v),
34            None => {
35                if let Some(scope) = self.scopes.pop() {
36                    self.iter = (self.get_iter_of_scope)(scope);
37                    self.siter = (self.get_siter_of_scope)(scope);
38                    self.scopes
39                        .extend(self.siter.by_ref().map(|s| self.hierarchy.get(s)));
40                    self.next()
41                } else {
42                    None
43                }
44            }
45        }
46    }
47}
48
49pub trait VarRefsIter {
50    fn var_refs_iter<'s>(&'s self) -> impl Iterator<Item = VarRef> + 's;
51}
52
53impl VarRefsIter for Hierarchy {
54    fn var_refs_iter<'s>(&'s self) -> impl Iterator<Item = VarRef> + 's {
55        VarRefIterator {
56            hierarchy: &self,
57            scopes: self.scopes().map(|s| self.get(s)).collect(),
58            get_iter_of_scope: |s: &Scope| s.vars(self),
59            get_siter_of_scope: |s: &Scope| s.scopes(self),
60            iter: self.vars(),
61            siter: self.scopes(),
62        }
63    }
64}