libreda_logic/algorithms/
visit.rs

1// SPDX-FileCopyrightText: 2022 Thomas Kramer <code@tkramer.ch>
2//
3// SPDX-License-Identifier: AGPL-3.0-or-later
4
5//! Traverse nodes of a logic network.
6
7use fnv::FnvHashSet;
8use std::hash::Hash;
9
10use crate::network::Network;
11
12/// Iterator over topologically sorted nodes.
13#[derive(Clone)]
14pub struct TopologicalIter<N>
15where
16    N: Network,
17{
18    network: N,
19    stack: Vec<N::NodeId>,
20    // TODO: Could be more efficient if N defines a 'VisitMap' type. This could also be an array instead of a HashSet.
21    visited: FnvHashSet<N::NodeId>,
22    /// Include primary inputs.
23    visit_primary_inputs: bool,
24    visit_constants: bool,
25}
26
27impl<N> TopologicalIter<N>
28where
29    N: Network,
30{
31    /// Create a new iterator which iterates over all nodes in the fan-in of the `top_nodes` including the `top_nodes`.
32    /// Nodes are visited in topological order, each node is visited at most once.
33    /// Includes primary inputs by default.
34    /// Excludes constants by default.
35    pub fn new(network: N, top_nodes: Vec<N::NodeId>) -> Self {
36        Self {
37            network,
38            stack: top_nodes,
39            visited: Default::default(),
40            visit_primary_inputs: true,
41            visit_constants: false,
42        }
43    }
44
45    /// Enable or disable visiting the primary inputs.
46    pub fn visit_primary_inputs(mut self, visit_primary_inputs: bool) -> Self {
47        self.visit_primary_inputs = visit_primary_inputs;
48        self
49    }
50
51    /// Enable or disable visiting constants.
52    pub fn visit_constants(mut self, visit_constants: bool) -> Self {
53        self.visit_constants = visit_constants;
54        self
55    }
56}
57
58impl<N> Iterator for TopologicalIter<N>
59where
60    N: Network,
61    N::NodeId: Hash + Eq,
62{
63    type Item = N::NodeId;
64
65    fn next(&mut self) -> Option<Self::Item> {
66        let net = &self.network;
67        loop {
68            if let Some(n) = self.stack.pop() {
69                let signal = net.get_node_output(&n);
70
71                if self.visited.contains(&n) {
72                    // Node already visited.
73                    continue;
74                } else if self.visit_constants && net.is_constant(signal) {
75                    self.visited.insert(n.clone());
76                    break Some(n);
77                } else if self.visit_primary_inputs && net.is_input(signal) {
78                    // Primary input, not yet visited.
79                    self.visited.insert(n.clone());
80                    break Some(n);
81                } else {
82                    let operands = (0..net.num_node_inputs(&n)).map(|i| net.get_node_input(&n, i));
83
84                    let mut values_complete = true;
85
86                    for in_sig in operands {
87                        // Skip constants.
88                        if !self.visit_constants && net.get_constant_value(in_sig).is_some() {
89                            continue;
90                        }
91
92                        // Skip primary inputs.
93                        if !self.visit_primary_inputs && net.is_input(in_sig) {
94                            continue;
95                        }
96
97                        let node = net.get_source_node(&in_sig);
98
99                        // Skip visited nodes.
100                        if self.visited.contains(&node) {
101                            continue;
102                        }
103
104                        // Node was not visited yet.
105
106                        if values_complete {
107                            // Need to visit this signal later.
108                            self.stack.push(n.clone());
109                        }
110
111                        values_complete = false;
112                        self.stack.push(node);
113                    }
114
115                    if values_complete {
116                        self.visited.insert(n.clone());
117                        break Some(n);
118                    }
119                }
120            } else {
121                break None;
122            }
123        }
124    }
125}
126
127#[test]
128fn test_topological_iter() {
129    use crate::network::*;
130    use crate::networks::mig::Mig;
131
132    let mut mig = Mig::new();
133
134    let [a, b, c, d] = mig.create_primary_inputs();
135
136    let anb = mig.create_and(a, b);
137    let bnc = mig.create_and(b, c);
138    let cnd = mig.create_and(c, d);
139    let or = mig.create_maj3(anb, bnc, cnd);
140
141    let out = mig.create_primary_output(or);
142
143    {
144        let iter = TopologicalIter::new(&mig, vec![out]);
145
146        let topo_sorted: Vec<_> = iter.collect();
147
148        assert_eq!(topo_sorted.last(), Some(&out));
149        assert_eq!(topo_sorted.len(), 8);
150
151        [a, b, c, d, anb, bnc, cnd, or]
152            .iter()
153            .for_each(|node| assert!(topo_sorted.contains(node)));
154    }
155
156    // Without primary inputs:
157    {
158        let iter = TopologicalIter::new(&mig, vec![out]).visit_primary_inputs(false);
159
160        let topo_sorted: Vec<_> = iter.collect();
161
162        assert_eq!(topo_sorted.last(), Some(&out));
163        assert_eq!(topo_sorted.len(), 4);
164
165        [anb, bnc, cnd, or]
166            .iter()
167            .for_each(|node| assert!(topo_sorted.contains(node)));
168    }
169}