1 //! Traversals over the IR.
2 
3 use crate::ir;
4 use alloc::vec::Vec;
5 use core::fmt::Debug;
6 use core::hash::Hash;
7 use cranelift_entity::EntitySet;
8 
9 /// A low-level DFS traversal event: either entering or exiting the traversal of
10 /// a block.
11 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
12 pub enum Event {
13     /// Entering traversal of a block.
14     ///
15     /// Processing a block upon this event corresponds to a pre-order,
16     /// depth-first traversal.
17     Enter,
18 
19     /// Exiting traversal of a block.
20     ///
21     /// Processing a block upon this event corresponds to a post-order,
22     /// depth-first traversal.
23     Exit,
24 }
25 
26 /// A depth-first traversal.
27 ///
28 /// This is a fairly low-level traversal type, and is generally intended to be
29 /// used as a building block for making specific pre-order or post-order
30 /// traversals for whatever problem is at hand.
31 ///
32 /// This type may be reused multiple times across different passes or functions
33 /// and will internally reuse any heap allocations its already made.
34 ///
35 /// Traversal is not recursive.
36 #[derive(Debug, Default, Clone)]
37 pub struct Dfs {
38     stack: Vec<(Event, ir::Block)>,
39     seen: EntitySet<ir::Block>,
40 }
41 
42 impl Dfs {
43     /// Construct a new depth-first traversal.
new() -> Self44     pub fn new() -> Self {
45         Self::default()
46     }
47 
48     /// Perform a depth-first traversal over the given function.
49     ///
50     /// Yields pairs of `(Event, ir::Block)`.
51     ///
52     /// This iterator can be used to perform either pre- or post-order
53     /// traversals, or a combination of the two.
iter<'a>(&'a mut self, func: &'a ir::Function) -> DfsIter<'a>54     pub fn iter<'a>(&'a mut self, func: &'a ir::Function) -> DfsIter<'a> {
55         self.clear();
56         if let Some(e) = func.layout.entry_block() {
57             self.stack.push((Event::Enter, e));
58         }
59         DfsIter { dfs: self, func }
60     }
61 
62     /// Perform a pre-order traversal over the given function.
63     ///
64     /// Yields `ir::Block` items.
pre_order_iter<'a>(&'a mut self, func: &'a ir::Function) -> DfsPreOrderIter<'a>65     pub fn pre_order_iter<'a>(&'a mut self, func: &'a ir::Function) -> DfsPreOrderIter<'a> {
66         DfsPreOrderIter(self.iter(func))
67     }
68 
69     /// Perform a post-order traversal over the given function.
70     ///
71     /// Yields `ir::Block` items.
post_order_iter<'a>(&'a mut self, func: &'a ir::Function) -> DfsPostOrderIter<'a>72     pub fn post_order_iter<'a>(&'a mut self, func: &'a ir::Function) -> DfsPostOrderIter<'a> {
73         DfsPostOrderIter(self.iter(func))
74     }
75 
76     /// Clear this DFS, but keep its allocations for future reuse.
clear(&mut self)77     pub fn clear(&mut self) {
78         let Dfs { stack, seen } = self;
79         stack.clear();
80         seen.clear();
81     }
82 }
83 
84 /// An iterator that yields pairs of `(Event, ir::Block)` items as it performs a
85 /// depth-first traversal over its associated function.
86 pub struct DfsIter<'a> {
87     dfs: &'a mut Dfs,
88     func: &'a ir::Function,
89 }
90 
91 impl Iterator for DfsIter<'_> {
92     type Item = (Event, ir::Block);
93 
next(&mut self) -> Option<(Event, ir::Block)>94     fn next(&mut self) -> Option<(Event, ir::Block)> {
95         loop {
96             let (event, block) = self.dfs.stack.pop()?;
97 
98             if event == Event::Enter {
99                 let first_time_seeing = self.dfs.seen.insert(block);
100                 if !first_time_seeing {
101                     continue;
102                 }
103 
104                 self.dfs.stack.push((Event::Exit, block));
105                 self.dfs.stack.extend(
106                     self.func
107                         .block_successors(block)
108                         // Heuristic: chase the children in reverse. This puts
109                         // the first successor block first in the postorder, all
110                         // other things being equal, which tends to prioritize
111                         // loop backedges over out-edges, putting the edge-block
112                         // closer to the loop body and minimizing live-ranges in
113                         // linear instruction space. This heuristic doesn't have
114                         // any effect on the computation of dominators, and is
115                         // purely for other consumers of the postorder we cache
116                         // here.
117                         .rev()
118                         // This is purely an optimization to avoid additional
119                         // iterations of the loop, and is not required; it's
120                         // merely inlining the check from the outer conditional
121                         // of this case to avoid the extra loop iteration. This
122                         // also avoids potential excess stack growth.
123                         .filter(|block| !self.dfs.seen.contains(*block))
124                         .map(|block| (Event::Enter, block)),
125                 );
126             }
127 
128             return Some((event, block));
129         }
130     }
131 }
132 
133 /// An iterator that yields `ir::Block` items during a depth-first, pre-order
134 /// traversal over its associated function.
135 pub struct DfsPreOrderIter<'a>(DfsIter<'a>);
136 
137 impl Iterator for DfsPreOrderIter<'_> {
138     type Item = ir::Block;
139 
next(&mut self) -> Option<Self::Item>140     fn next(&mut self) -> Option<Self::Item> {
141         loop {
142             match self.0.next()? {
143                 (Event::Enter, b) => return Some(b),
144                 (Event::Exit, _) => continue,
145             }
146         }
147     }
148 }
149 
150 /// An iterator that yields `ir::Block` items during a depth-first, post-order
151 /// traversal over its associated function.
152 pub struct DfsPostOrderIter<'a>(DfsIter<'a>);
153 
154 impl Iterator for DfsPostOrderIter<'_> {
155     type Item = ir::Block;
156 
next(&mut self) -> Option<Self::Item>157     fn next(&mut self) -> Option<Self::Item> {
158         loop {
159             match self.0.next()? {
160                 (Event::Exit, b) => return Some(b),
161                 (Event::Enter, _) => continue,
162             }
163         }
164     }
165 }
166 
167 #[cfg(test)]
168 mod tests {
169     use super::*;
170     use crate::cursor::{Cursor, FuncCursor};
171     use crate::ir::{Function, InstBuilder, TrapCode, types::I32};
172 
173     #[test]
test_dfs_traversal()174     fn test_dfs_traversal() {
175         let _ = env_logger::try_init();
176 
177         let mut func = Function::new();
178 
179         let block0 = func.dfg.make_block();
180         let v0 = func.dfg.append_block_param(block0, I32);
181         let block1 = func.dfg.make_block();
182         let block2 = func.dfg.make_block();
183         let block3 = func.dfg.make_block();
184 
185         let mut cur = FuncCursor::new(&mut func);
186 
187         // block0(v0):
188         //   br_if v0, block2, trap_block
189         cur.insert_block(block0);
190         cur.ins().brif(v0, block2, &[], block3, &[]);
191 
192         // block3:
193         //   trap user0
194         cur.insert_block(block3);
195         cur.ins().trap(TrapCode::unwrap_user(1));
196 
197         // block1:
198         //   v1 = iconst.i32 1
199         //   v2 = iadd v0, v1
200         //   jump block0(v2)
201         cur.insert_block(block1);
202         let v1 = cur.ins().iconst(I32, 1);
203         let v2 = cur.ins().iadd(v0, v1);
204         cur.ins().jump(block0, &[v2.into()]);
205 
206         // block2:
207         //   return v0
208         cur.insert_block(block2);
209         cur.ins().return_(&[v0]);
210 
211         let mut dfs = Dfs::new();
212 
213         assert_eq!(
214             dfs.iter(&func).collect::<Vec<_>>(),
215             vec![
216                 (Event::Enter, block0),
217                 (Event::Enter, block2),
218                 (Event::Exit, block2),
219                 (Event::Enter, block3),
220                 (Event::Exit, block3),
221                 (Event::Exit, block0)
222             ],
223         );
224     }
225 
226     #[test]
multiple_successors_to_the_same_block()227     fn multiple_successors_to_the_same_block() {
228         let _ = env_logger::try_init();
229 
230         let mut func = Function::new();
231 
232         let block0 = func.dfg.make_block();
233         let block1 = func.dfg.make_block();
234 
235         let mut cur = FuncCursor::new(&mut func);
236 
237         // block0(v0):
238         //   v1 = iconst.i32 36
239         //   v2 = iconst.i32 42
240         //   br_if v0, block1(v1), block1(v2)
241         cur.insert_block(block0);
242         let v0 = cur.func.dfg.append_block_param(block0, I32);
243         let v1 = cur.ins().iconst(ir::types::I32, 36);
244         let v2 = cur.ins().iconst(ir::types::I32, 42);
245         cur.ins()
246             .brif(v0, block1, &[v1.into()], block1, &[v2.into()]);
247 
248         // block1(v3: i32):
249         //   return v3
250         cur.insert_block(block1);
251         let v3 = cur.func.dfg.append_block_param(block1, I32);
252         cur.ins().return_(&[v3]);
253 
254         let mut dfs = Dfs::new();
255 
256         // We should only enter `block1` once.
257         assert_eq!(
258             dfs.iter(&func).collect::<Vec<_>>(),
259             vec![
260                 (Event::Enter, block0),
261                 (Event::Enter, block1),
262                 (Event::Exit, block1),
263                 (Event::Exit, block0),
264             ],
265         );
266 
267         // We should only iterate over `block1` once in a pre-order traversal.
268         assert_eq!(
269             dfs.pre_order_iter(&func).collect::<Vec<_>>(),
270             vec![block0, block1],
271         );
272 
273         // We should only iterate over `block1` once in a post-order traversal.
274         assert_eq!(
275             dfs.post_order_iter(&func).collect::<Vec<_>>(),
276             vec![block1, block0],
277         );
278     }
279 }
280