fe_mir/ir/
body_order.rs

1use fxhash::FxHashMap;
2
3use super::{basic_block::BasicBlockId, function::BodyDataStore, inst::InstId};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
6/// Represents basic block order and instruction order.
7pub struct BodyOrder {
8    blocks: FxHashMap<BasicBlockId, BlockNode>,
9    insts: FxHashMap<InstId, InstNode>,
10    entry_block: BasicBlockId,
11    last_block: BasicBlockId,
12}
13impl BodyOrder {
14    pub fn new(entry_block: BasicBlockId) -> Self {
15        let entry_block_node = BlockNode::default();
16        let mut blocks = FxHashMap::default();
17        blocks.insert(entry_block, entry_block_node);
18
19        Self {
20            blocks,
21            insts: FxHashMap::default(),
22            entry_block,
23            last_block: entry_block,
24        }
25    }
26
27    /// Returns an entry block of a function body.
28    pub fn entry(&self) -> BasicBlockId {
29        self.entry_block
30    }
31
32    /// Returns a last block of a function body.
33    pub fn last_block(&self) -> BasicBlockId {
34        self.last_block
35    }
36
37    /// Returns `true` if a block doesn't contain any block.
38    pub fn is_block_empty(&self, block: BasicBlockId) -> bool {
39        self.first_inst(block).is_none()
40    }
41
42    /// Returns `true` if a function body contains a given `block`.
43    pub fn is_block_inserted(&self, block: BasicBlockId) -> bool {
44        self.blocks.contains_key(&block)
45    }
46
47    /// Returns a number of block in a function.
48    pub fn block_num(&self) -> usize {
49        self.blocks.len()
50    }
51
52    /// Returns a previous block of a given block.
53    ///
54    /// # Panics
55    /// Panics if `block` is not inserted yet.
56    pub fn prev_block(&self, block: BasicBlockId) -> Option<BasicBlockId> {
57        self.blocks[&block].prev
58    }
59
60    /// Returns a next block of a given block.
61    ///
62    /// # Panics
63    /// Panics if `block` is not inserted yet.
64    pub fn next_block(&self, block: BasicBlockId) -> Option<BasicBlockId> {
65        self.blocks[&block].next
66    }
67
68    /// Returns `true` is a given `inst` is inserted.
69    pub fn is_inst_inserted(&self, inst: InstId) -> bool {
70        self.insts.contains_key(&inst)
71    }
72
73    /// Returns first instruction of a block if exists.
74    ///
75    /// # Panics
76    /// Panics if `block` is not inserted yet.
77    pub fn first_inst(&self, block: BasicBlockId) -> Option<InstId> {
78        self.blocks[&block].first_inst
79    }
80
81    /// Returns a terminator instruction of a block.
82    ///
83    /// # Panics
84    /// Panics if
85    /// 1. `block` is not inserted yet.
86    pub fn terminator(&self, store: &BodyDataStore, block: BasicBlockId) -> Option<InstId> {
87        let last_inst = self.last_inst(block)?;
88        if store.is_terminator(last_inst) {
89            Some(last_inst)
90        } else {
91            None
92        }
93    }
94
95    /// Returns `true` if a `block` is terminated.
96    pub fn is_terminated(&self, store: &BodyDataStore, block: BasicBlockId) -> bool {
97        self.terminator(store, block).is_some()
98    }
99
100    /// Returns a last instruction of a block.
101    ///
102    /// # Panics
103    /// Panics if `block` is not inserted yet.
104    pub fn last_inst(&self, block: BasicBlockId) -> Option<InstId> {
105        self.blocks[&block].last_inst
106    }
107
108    /// Returns a previous instruction of a given `inst`.
109    ///
110    /// # Panics
111    /// Panics if `inst` is not inserted yet.
112    pub fn prev_inst(&self, inst: InstId) -> Option<InstId> {
113        self.insts[&inst].prev
114    }
115
116    /// Returns a next instruction of a given `inst`.
117    ///
118    /// # Panics
119    /// Panics if `inst` is not inserted yet.
120    pub fn next_inst(&self, inst: InstId) -> Option<InstId> {
121        self.insts[&inst].next
122    }
123
124    /// Returns a block to which a given `inst` belongs.
125    ///
126    /// # Panics
127    /// Panics if `inst` is not inserted yet.
128    pub fn inst_block(&self, inst: InstId) -> BasicBlockId {
129        self.insts[&inst].block
130    }
131
132    /// Returns an iterator which iterates all basic blocks in a function body
133    /// in pre-order.
134    pub fn iter_block(&self) -> impl Iterator<Item = BasicBlockId> + '_ {
135        BlockIter {
136            next: Some(self.entry_block),
137            blocks: &self.blocks,
138        }
139    }
140
141    /// Returns an iterator which iterates all instruction in a given `block` in
142    /// pre-order.
143    ///
144    /// # Panics
145    /// Panics if `block` is not inserted yet.
146    pub fn iter_inst(&self, block: BasicBlockId) -> impl Iterator<Item = InstId> + '_ {
147        InstIter {
148            next: self.blocks[&block].first_inst,
149            insts: &self.insts,
150        }
151    }
152
153    /// Appends a given `block` to a function body.
154    ///
155    /// # Panics
156    /// Panics if a given `block` is already inserted to a function.
157    pub fn append_block(&mut self, block: BasicBlockId) {
158        debug_assert!(!self.is_block_inserted(block));
159
160        let mut block_node = BlockNode::default();
161        let last_block = self.last_block;
162        let last_block_node = &mut self.block_mut(last_block);
163        last_block_node.next = Some(block);
164        block_node.prev = Some(last_block);
165
166        self.blocks.insert(block, block_node);
167        self.last_block = block;
168    }
169
170    /// Inserts a given `block` before a `before` block.
171    ///
172    /// # Panics
173    /// Panics if
174    /// 1. a given `block` is already inserted.
175    /// 2. a given `before` block is NOTE inserted yet.
176    pub fn insert_block_before_block(&mut self, block: BasicBlockId, before: BasicBlockId) {
177        debug_assert!(self.is_block_inserted(before));
178        debug_assert!(!self.is_block_inserted(block));
179
180        let mut block_node = BlockNode::default();
181
182        match self.blocks[&before].prev {
183            Some(prev) => {
184                block_node.prev = Some(prev);
185                self.block_mut(prev).next = Some(block);
186            }
187            None => self.entry_block = block,
188        }
189
190        block_node.next = Some(before);
191        self.block_mut(before).prev = Some(block);
192        self.blocks.insert(block, block_node);
193    }
194
195    /// Inserts a given `block` after a `after` block.
196    ///
197    /// # Panics
198    /// Panics if
199    /// 1. a given `block` is already inserted.
200    /// 2. a given `after` block is NOTE inserted yet.
201    pub fn insert_block_after_block(&mut self, block: BasicBlockId, after: BasicBlockId) {
202        debug_assert!(self.is_block_inserted(after));
203        debug_assert!(!self.is_block_inserted(block));
204
205        let mut block_node = BlockNode::default();
206
207        match self.blocks[&after].next {
208            Some(next) => {
209                block_node.next = Some(next);
210                self.block_mut(next).prev = Some(block);
211            }
212            None => self.last_block = block,
213        }
214        block_node.prev = Some(after);
215        self.block_mut(after).next = Some(block);
216        self.blocks.insert(block, block_node);
217    }
218
219    /// Remove a given `block` from a function. All instructions in a block are
220    /// also removed.
221    ///
222    /// # Panics
223    /// Panics if
224    /// 1. a given `block` is NOT inserted.
225    /// 2. a `block` is the last one block in a function.
226    pub fn remove_block(&mut self, block: BasicBlockId) {
227        debug_assert!(self.is_block_inserted(block));
228        debug_assert!(self.block_num() > 1);
229
230        // Remove all insts in a `block`.
231        let mut next_inst = self.first_inst(block);
232        while let Some(inst) = next_inst {
233            next_inst = self.next_inst(inst);
234            self.remove_inst(inst);
235        }
236
237        // Remove `block`.
238        let block_node = &self.blocks[&block];
239        let prev_block = block_node.prev;
240        let next_block = block_node.next;
241        match (prev_block, next_block) {
242            // `block` is in the middle of a function.
243            (Some(prev), Some(next)) => {
244                self.block_mut(prev).next = Some(next);
245                self.block_mut(next).prev = Some(prev);
246            }
247            // `block` is the last block of a function.
248            (Some(prev), None) => {
249                self.block_mut(prev).next = None;
250                self.last_block = prev;
251            }
252            // `block` is the first block of a function.
253            (None, Some(next)) => {
254                self.block_mut(next).prev = None;
255                self.entry_block = next
256            }
257            (None, None) => {
258                unreachable!()
259            }
260        }
261
262        self.blocks.remove(&block);
263    }
264
265    /// Appends `inst` to the end of a `block`
266    ///
267    /// # Panics
268    /// Panics if
269    /// 1. a given `block` is NOT inserted.
270    /// 2. a given `inst` is already inserted.
271    pub fn append_inst(&mut self, inst: InstId, block: BasicBlockId) {
272        debug_assert!(self.is_block_inserted(block));
273        debug_assert!(!self.is_inst_inserted(inst));
274
275        let mut inst_node = InstNode::new(block);
276
277        if let Some(last_inst) = self.blocks[&block].last_inst {
278            inst_node.prev = Some(last_inst);
279            self.inst_mut(last_inst).next = Some(inst);
280        } else {
281            self.block_mut(block).first_inst = Some(inst);
282        }
283
284        self.block_mut(block).last_inst = Some(inst);
285        self.insts.insert(inst, inst_node);
286    }
287
288    /// Prepends `inst` to the beginning of a `block`
289    ///
290    /// # Panics
291    /// Panics if
292    /// 1. a given `block` is NOT inserted.
293    /// 2. a given `inst` is already inserted.
294    pub fn prepend_inst(&mut self, inst: InstId, block: BasicBlockId) {
295        debug_assert!(self.is_block_inserted(block));
296        debug_assert!(!self.is_inst_inserted(inst));
297
298        let mut inst_node = InstNode::new(block);
299
300        if let Some(first_inst) = self.blocks[&block].first_inst {
301            inst_node.next = Some(first_inst);
302            self.inst_mut(first_inst).prev = Some(inst);
303        } else {
304            self.block_mut(block).last_inst = Some(inst);
305        }
306
307        self.block_mut(block).first_inst = Some(inst);
308        self.insts.insert(inst, inst_node);
309    }
310
311    /// Insert `inst` before `before` inst.
312    ///
313    /// # Panics
314    /// Panics if
315    /// 1. a given `before` is NOT inserted.
316    /// 2. a given `inst` is already inserted.
317    pub fn insert_inst_before_inst(&mut self, inst: InstId, before: InstId) {
318        debug_assert!(self.is_inst_inserted(before));
319        debug_assert!(!self.is_inst_inserted(inst));
320
321        let before_inst_node = &self.insts[&before];
322        let block = before_inst_node.block;
323        let mut inst_node = InstNode::new(block);
324
325        match before_inst_node.prev {
326            Some(prev) => {
327                inst_node.prev = Some(prev);
328                self.inst_mut(prev).next = Some(inst);
329            }
330            None => self.block_mut(block).first_inst = Some(inst),
331        }
332        inst_node.next = Some(before);
333        self.inst_mut(before).prev = Some(inst);
334        self.insts.insert(inst, inst_node);
335    }
336
337    /// Insert `inst` after `after` inst.
338    ///
339    /// # Panics
340    /// Panics if
341    /// 1. a given `after` is NOT inserted.
342    /// 2. a given `inst` is already inserted.
343    pub fn insert_inst_after(&mut self, inst: InstId, after: InstId) {
344        debug_assert!(self.is_inst_inserted(after));
345        debug_assert!(!self.is_inst_inserted(inst));
346
347        let after_inst_node = &self.insts[&after];
348        let block = after_inst_node.block;
349        let mut inst_node = InstNode::new(block);
350
351        match after_inst_node.next {
352            Some(next) => {
353                inst_node.next = Some(next);
354                self.inst_mut(next).prev = Some(inst);
355            }
356            None => self.block_mut(block).last_inst = Some(inst),
357        }
358        inst_node.prev = Some(after);
359        self.inst_mut(after).next = Some(inst);
360        self.insts.insert(inst, inst_node);
361    }
362
363    /// Remove instruction from the function body.
364    ///
365    /// # Panics
366    /// Panics if a given `inst` is not inserted.
367    pub fn remove_inst(&mut self, inst: InstId) {
368        debug_assert!(self.is_inst_inserted(inst));
369
370        let inst_node = &self.insts[&inst];
371        let inst_block = inst_node.block;
372        let prev_inst = inst_node.prev;
373        let next_inst = inst_node.next;
374        match (prev_inst, next_inst) {
375            (Some(prev), Some(next)) => {
376                self.inst_mut(prev).next = Some(next);
377                self.inst_mut(next).prev = Some(prev);
378            }
379            (Some(prev), None) => {
380                self.inst_mut(prev).next = None;
381                self.block_mut(inst_block).last_inst = Some(prev);
382            }
383            (None, Some(next)) => {
384                self.inst_mut(next).prev = None;
385                self.block_mut(inst_block).first_inst = Some(next);
386            }
387            (None, None) => {
388                let block_node = self.block_mut(inst_block);
389                block_node.first_inst = None;
390                block_node.last_inst = None;
391            }
392        }
393
394        self.insts.remove(&inst);
395    }
396
397    fn block_mut(&mut self, block: BasicBlockId) -> &mut BlockNode {
398        self.blocks.get_mut(&block).unwrap()
399    }
400
401    fn inst_mut(&mut self, inst: InstId) -> &mut InstNode {
402        self.insts.get_mut(&inst).unwrap()
403    }
404}
405
406struct BlockIter<'a> {
407    next: Option<BasicBlockId>,
408    blocks: &'a FxHashMap<BasicBlockId, BlockNode>,
409}
410
411impl<'a> Iterator for BlockIter<'a> {
412    type Item = BasicBlockId;
413
414    fn next(&mut self) -> Option<BasicBlockId> {
415        let next = self.next?;
416        self.next = self.blocks[&next].next;
417        Some(next)
418    }
419}
420
421struct InstIter<'a> {
422    next: Option<InstId>,
423    insts: &'a FxHashMap<InstId, InstNode>,
424}
425
426impl<'a> Iterator for InstIter<'a> {
427    type Item = InstId;
428
429    fn next(&mut self) -> Option<InstId> {
430        let next = self.next?;
431        self.next = self.insts[&next].next;
432        Some(next)
433    }
434}
435
436#[derive(Default, Debug, Clone, PartialEq, Eq)]
437/// A helper struct to track a basic block order in a function body.
438struct BlockNode {
439    /// A previous block.
440    prev: Option<BasicBlockId>,
441
442    /// A next block.
443    next: Option<BasicBlockId>,
444
445    /// A first instruction of a block.
446    first_inst: Option<InstId>,
447
448    /// A last instruction of a block.
449    last_inst: Option<InstId>,
450}
451
452#[derive(Debug, Clone, PartialEq, Eq)]
453/// A helper struct to track a instruction order in a basic block.
454struct InstNode {
455    /// An block to which a inst belongs.
456    block: BasicBlockId,
457
458    /// A previous instruction.
459    prev: Option<InstId>,
460
461    /// A next instruction.
462    next: Option<InstId>,
463}
464
465impl InstNode {
466    fn new(block: BasicBlockId) -> Self {
467        Self {
468            block,
469            prev: None,
470            next: None,
471        }
472    }
473}