calyx_opt/traversal/
post_order.rs

1use calyx_ir::{self as ir, CellType};
2use calyx_utils::CalyxResult;
3use petgraph::algo;
4use petgraph::graph::{DiGraph, NodeIndex};
5use std::collections::HashMap;
6
7/// The order in which the components are traversed.
8#[derive(Default, PartialEq, Eq)]
9pub enum Order {
10    /// Use an arbitrary order.
11    #[default]
12    No,
13    /// Traverse components in pre-order.
14    Pre,
15    /// Traverse components in post-order.
16    Post,
17}
18
19/// Define traversal order of components: pre-order, post-order, or none.
20///
21/// ## No order
22/// Iterates over the components in any order
23///
24/// ## Post-order
25/// If a component `B` creates a cell of type `A` then component `A` is
26/// guaranteed to be visited before `B`.
27/// This is done by finding a topological order over a graph where `A` will
28/// have a directed edge to `B`.
29///
30/// Instead of constructing a new vector of components in a topological order,
31/// the implementation builds an `order` vector which contains indices into the
32/// original component vector.
33/// This way, we can return the components in the input order once we're done
34/// with the post order traversal.
35///
36/// ## Pre-order
37/// Reverse of post-order
38///
39/// ## Example
40/// ```rust
41/// let comps: Vec<ir::Component>;
42/// // Construct a post order.
43/// let post = PostOrder::new(comps, Order::Post);
44/// // Apply a mutable update to components.
45/// let upd: FnMut(&mut ir::Component) -> CalyxResult<()>;
46/// post.apply_update(upd);
47/// // Recover the components in original order.
48/// let new_comps = post.take();
49/// ```
50pub struct CompTraversal {
51    /// A topological ordering of the components.
52    order: Vec<NodeIndex>,
53    /// Vector of components in the original ordering.
54    comps: Vec<ir::Component>,
55}
56
57impl CompTraversal {
58    /// Returns a new instance the PostOrder iterator given a Vector of components.
59    ///
60    /// # Panics
61    /// Panics if there is no post-order traversal of the vectors possible.
62    pub fn new(comps: Vec<ir::Component>, order: Order) -> Self {
63        // If the order is not specified, return the components in the original order.
64        if order == Order::No {
65            return Self {
66                order: (0..comps.len()).map(NodeIndex::new).collect(),
67                comps,
68            };
69        }
70        let mut graph: DiGraph<usize, ()> = DiGraph::new();
71        // Reverse mapping from index to comps.
72        let rev_map: HashMap<ir::Id, NodeIndex> = comps
73            .iter()
74            .enumerate()
75            .map(|(idx, c)| (c.name, graph.add_node(idx)))
76            .collect::<HashMap<_, _>>();
77
78        // Construct a graph.
79        for comp in &comps {
80            for cell in comp.cells.iter() {
81                if let CellType::Component { name, .. } =
82                    &cell.borrow().prototype
83                {
84                    graph.add_edge(rev_map[name], rev_map[&comp.name], ());
85                }
86            }
87        }
88
89        // Build a topologically sorted ordering of the graph.
90        let mut topo = algo::toposort(&graph, None)
91            .expect("There is a cycle in definition of component cells");
92
93        // Reverse the order if a pre-order traversal is requested
94        if order == Order::Pre {
95            topo.reverse();
96        }
97        Self { order: topo, comps }
98    }
99
100    /// Traverses components in post-order and applies `upd`.
101    pub fn apply_update<F>(&mut self, mut upd: F) -> CalyxResult<()>
102    where
103        F: FnMut(&mut ir::Component, &Vec<ir::Component>) -> CalyxResult<()>,
104    {
105        for idx in self.order.iter() {
106            let mut comp = self.comps.remove(idx.index());
107            upd(&mut comp, &self.comps)?;
108            self.comps.insert(idx.index(), comp)
109        }
110
111        Ok(())
112    }
113
114    /// Returns the underlying component vector in original order.
115    pub fn take(self) -> Vec<ir::Component> {
116        self.comps
117    }
118}
calyx_opt/traversal/post_order.rs

calyx_opt/traversal/
post_order.rs
