rustc_mir_transform/
add_call_guards.rs1use rustc_data_structures::thin_vec::ThinVec;
19use rustc_index::{Idx, IndexVec};
20use rustc_middle::mir::*;
21use rustc_middle::ty::TyCtxt;
22use tracing::debug;
23
24#[derive(PartialEq)]
25pub(super) enum AddCallGuards {
26 AllCallEdges,
27 CriticalCallEdges,
28}
29pub(super) use self::AddCallGuards::*;
30
31impl<'tcx> crate::MirPass<'tcx> for AddCallGuards {
32 fn run_pass(&self, _tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
33 let mut pred_count = IndexVec::from_elem(0u8, &body.basic_blocks);
34 for (_, data) in body.basic_blocks.iter_enumerated() {
35 for succ in data.terminator().successors() {
36 pred_count[succ] = pred_count[succ].saturating_add(1);
37 }
38 }
39
40 enum Action {
41 Call,
42 Asm { target_index: usize },
43 }
44
45 let mut work = Vec::with_capacity(body.basic_blocks.len());
46 for (bb, block) in body.basic_blocks.iter_enumerated() {
47 let term = block.terminator();
48 match term.kind {
49 TerminatorKind::Call { target: Some(destination), unwind, .. }
50 if pred_count[destination] > 1
51 && (generates_invoke(unwind) || self == &AllCallEdges) =>
52 {
53 work.push((bb, Action::Call));
55 }
56 TerminatorKind::InlineAsm {
57 asm_macro: InlineAsmMacro::Asm,
58 ref targets,
59 ref operands,
60 unwind,
61 ..
62 } if self == &CriticalCallEdges => {
63 let has_outputs = operands.iter().any(|op| {
64 matches!(op, InlineAsmOperand::InOut { .. } | InlineAsmOperand::Out { .. })
65 });
66 let has_labels =
67 operands.iter().any(|op| matches!(op, InlineAsmOperand::Label { .. }));
68 if has_outputs && (has_labels || generates_invoke(unwind)) {
69 for (target_index, target) in targets.iter().enumerate() {
70 if pred_count[*target] > 1 {
71 work.push((bb, Action::Asm { target_index }));
72 }
73 }
74 }
75 }
76 _ => {}
77 }
78 }
79
80 if work.is_empty() {
81 return;
82 }
83
84 let mut new_blocks = Vec::with_capacity(work.len());
86
87 let cur_len = body.basic_blocks.len();
88 let mut new_block = |source_info: SourceInfo, is_cleanup: bool, target: BasicBlock| {
89 let block = BasicBlockData::new(
90 Some(Terminator {
91 source_info,
92 kind: TerminatorKind::Goto { target },
93 attributes: ThinVec::new(),
94 }),
95 is_cleanup,
96 );
97 let idx = cur_len + new_blocks.len();
98 new_blocks.push(block);
99 BasicBlock::new(idx)
100 };
101
102 let basic_blocks = body.basic_blocks.as_mut();
103 for (source, action) in work {
104 let block = &mut basic_blocks[source];
105 let is_cleanup = block.is_cleanup;
106 let term = block.terminator_mut();
107 let source_info = term.source_info;
108 let destination = match action {
109 Action::Call => {
110 let TerminatorKind::Call { target: Some(ref mut destination), .. } = term.kind
111 else {
112 unreachable!()
113 };
114 destination
115 }
116 Action::Asm { target_index } => {
117 let TerminatorKind::InlineAsm { ref mut targets, .. } = term.kind else {
118 unreachable!()
119 };
120 &mut targets[target_index]
121 }
122 };
123 *destination = new_block(source_info, is_cleanup, *destination);
124 }
125
126 debug!("Broke {} N edges", new_blocks.len());
127 basic_blocks.extend(new_blocks);
128 }
129
130 fn is_required(&self) -> bool {
131 true
132 }
133}
134
135fn generates_invoke(unwind: UnwindAction) -> bool {
137 match unwind {
138 UnwindAction::Continue | UnwindAction::Unreachable => false,
139 UnwindAction::Cleanup(_) | UnwindAction::Terminate(_) => true,
140 }
141}