cranelift_codegen/legalizer/
branch_to_trap.rs

1//! Rewrite branch-to-unconditional-trap into conditional trap instructions.
2//!
3//! Given this instruction:
4//!
5//! ```clif
6//! brif v0, block1, block2
7//! ```
8//!
9//! If we know that `block1` does nothing but immediately trap then we can
10//! rewrite that `brif` into the following:
11//!
12//! ```clif
13//! trapz v0, <trapcode>
14//! jump block2
15//! ```
16//!
17//! (And we can do the equivalent with `trapz` if `block2` immediately traps).
18//!
19//! This transformation allows for the conditional trap instructions to be GVN'd
20//! and for our egraphs mid-end to generally better optimize the program. We
21//! additionally have better codegen in our backends for `trapz` than branches
22//! to unconditional traps.
23
24use super::*;
25
26#[derive(Default)]
27pub struct BranchToTrap {
28    /// The set of blocks that contain exactly one unconditional trap
29    /// instruction.
30    just_trap_blocks: EntitySet<ir::Block>,
31}
32
33impl BranchToTrap {
34    /// Analyze the given block.
35    ///
36    /// The `block` must be terminated by a `trap` instruction.
37    pub fn analyze_trapping_block(&mut self, func: &ir::Function, block: ir::Block) {
38        if func.layout.block_contains_exactly_one_inst(block) {
39            self.just_trap_blocks.insert(block);
40        }
41    }
42
43    fn just_trap_block_code(&self, func: &ir::Function, block: ir::Block) -> ir::TrapCode {
44        debug_assert!(self.just_trap_blocks.contains(block));
45        debug_assert!(func.layout.block_contains_exactly_one_inst(block));
46        let inst = func.layout.first_inst(block).unwrap();
47        match func.dfg.insts[inst] {
48            InstructionData::Trap { code, .. } => code,
49            _ => unreachable!(),
50        }
51    }
52
53    /// Process a `brif` instruction, potentially performing our rewrite.
54    ///
55    /// The `inst` must be a `brif` containing the given `arg` and `blocks`.
56    pub fn process_brif(
57        &self,
58        func: &mut ir::Function,
59        inst: ir::Inst,
60        arg: ir::Value,
61        blocks: [ir::BlockCall; 2],
62    ) {
63        let consequent = blocks[0].block(&func.dfg.value_lists);
64        let alternative = blocks[1].block(&func.dfg.value_lists);
65
66        if self.just_trap_blocks.contains(consequent) {
67            let mut pos = FuncCursor::new(func);
68            pos.use_srcloc(
69                pos.func
70                    .layout
71                    .first_inst(consequent)
72                    .expect("just-trap blocks have exactly one inst"),
73            );
74            pos.goto_inst(inst);
75
76            let code = self.just_trap_block_code(pos.func, consequent);
77            pos.ins().trapnz(arg, code);
78
79            let args: SmallVec<[_; 8]> = blocks[1].args(&pos.func.dfg.value_lists).collect();
80            pos.func.dfg.replace(inst).jump(alternative, &args);
81        } else if self.just_trap_blocks.contains(alternative) {
82            let mut pos = FuncCursor::new(func);
83            pos.use_srcloc(
84                pos.func
85                    .layout
86                    .first_inst(alternative)
87                    .expect("just-trap blocks have exactly one inst"),
88            );
89            pos.goto_inst(inst);
90
91            let code = self.just_trap_block_code(pos.func, alternative);
92            pos.ins().trapz(arg, code);
93
94            let args: SmallVec<[_; 8]> = blocks[0].args(&pos.func.dfg.value_lists).collect();
95            pos.func.dfg.replace(inst).jump(consequent, &args);
96        }
97    }
98}
cranelift_codegen/legalizer/branch_to_trap.rs

cranelift_codegen/legalizer/
branch_to_trap.rs
