|
| 1 | +use crate::MirPass; |
| 2 | +use rustc_data_structures::fx::{FxHashMap, FxHashSet}; |
| 3 | +use rustc_index::bit_set::BitSet; |
| 4 | +use rustc_middle::mir::visit::{NonUseContext, PlaceContext, Visitor}; |
| 5 | +use rustc_middle::mir::*; |
| 6 | +use rustc_middle::ty::TyCtxt; |
| 7 | + |
| 8 | +const SUCCESSOR_LIMIT: usize = 100; |
| 9 | + |
| 10 | +pub struct SinkConstAssignments; |
| 11 | + |
| 12 | +impl<'tcx> MirPass<'tcx> for SinkConstAssignments { |
| 13 | + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { |
| 14 | + sess.mir_opt_level() >= 1 |
| 15 | + } |
| 16 | + |
| 17 | + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 18 | + // The primary benefit of this pass is sinking assignments to drop flags and enabling |
| 19 | + // ConstGoto and SimplifyCfg to merge the drop flag check into existing control flow. |
| 20 | + // If we permit sinking assignments to any local, we will sometimes sink an assignment into |
| 21 | + // but not completely through a goto chain, preventing SimplifyCfg from removing the |
| 22 | + // blocks. |
| 23 | + let mut optimizable_locals = branched_locals(body); |
| 24 | + if optimizable_locals.is_empty() { |
| 25 | + return; |
| 26 | + } |
| 27 | + |
| 28 | + let borrowed_locals = rustc_mir_dataflow::impls::borrowed_locals(body); |
| 29 | + optimizable_locals.subtract(&borrowed_locals); |
| 30 | + if optimizable_locals.is_empty() { |
| 31 | + return; |
| 32 | + } |
| 33 | + |
| 34 | + 'outer: for block in 0..body.basic_blocks.len() { |
| 35 | + let block = block.into(); |
| 36 | + let block_data = &body.basic_blocks[block]; |
| 37 | + let Some(terminator) = &block_data.terminator else { continue; }; |
| 38 | + |
| 39 | + let mut successors = Vec::new(); |
| 40 | + for succ in terminator.successors() { |
| 41 | + // Successors which are just a Resume are okay |
| 42 | + if is_empty_resume(&body.basic_blocks[succ]) { |
| 43 | + continue; |
| 44 | + } |
| 45 | + if body.basic_blocks.predecessors()[succ].len() != 1 { |
| 46 | + debug!("Bailing from {block:?} because {succ:?} has multiple predecessors"); |
| 47 | + continue 'outer; |
| 48 | + } |
| 49 | + successors.push(succ); |
| 50 | + } |
| 51 | + |
| 52 | + if successors.len() > SUCCESSOR_LIMIT { |
| 53 | + debug!("Will not sink assignment, its basic block has too many successors"); |
| 54 | + } |
| 55 | + |
| 56 | + let mut local_uses = None; |
| 57 | + for statement_idx in 0..body.basic_blocks[block].statements.len() { |
| 58 | + let statement = &body.basic_blocks[block].statements[statement_idx]; |
| 59 | + if let StatementKind::Assign(box (place, Rvalue::Use(Operand::Constant(_)))) = |
| 60 | + &statement.kind |
| 61 | + { |
| 62 | + let local = place.local; |
| 63 | + if !place.projection.is_empty() { |
| 64 | + debug!("Nonempty place projection: {statement:?}"); |
| 65 | + continue; |
| 66 | + } |
| 67 | + if !optimizable_locals.contains(local) { |
| 68 | + continue; |
| 69 | + } |
| 70 | + |
| 71 | + let uses = match &local_uses { |
| 72 | + Some(uses) => uses, |
| 73 | + None => { |
| 74 | + let mut visitor = CountUsesVisitor::new(); |
| 75 | + visitor.visit_basic_block_data(block, &body.basic_blocks[block]); |
| 76 | + local_uses = Some(visitor); |
| 77 | + local_uses.as_ref().unwrap() |
| 78 | + } |
| 79 | + }; |
| 80 | + |
| 81 | + // If the local dies in this block, don't propagate it |
| 82 | + if uses.dead.contains(&local) { |
| 83 | + continue; |
| 84 | + } |
| 85 | + if !uses.is_used_once(local) { |
| 86 | + debug!("Local used elsewhere in this block: {statement:?}"); |
| 87 | + continue; |
| 88 | + } |
| 89 | + if !tcx.consider_optimizing(|| format!("Sinking const assignment to {local:?}")) |
| 90 | + { |
| 91 | + debug!("optimization fuel exhausted"); |
| 92 | + break 'outer; |
| 93 | + } |
| 94 | + debug!("Sinking const assignment to {local:?}"); |
| 95 | + let blocks = body.basic_blocks.as_mut_preserves_cfg(); |
| 96 | + let statement = blocks[block].statements[statement_idx].replace_nop(); |
| 97 | + |
| 98 | + for succ in &successors { |
| 99 | + let mut successor_uses = SuccessorUsesVisitor::new(local); |
| 100 | + successor_uses.visit_basic_block_data(*succ, &blocks[*succ]); |
| 101 | + |
| 102 | + if let Some(used) = successor_uses.first_use() { |
| 103 | + if used == blocks[*succ].statements.len() { |
| 104 | + blocks[*succ].statements.push(statement.clone()); |
| 105 | + continue; |
| 106 | + } |
| 107 | + // If the first use of our local in this block is another const |
| 108 | + // assignment to it, do not paste a new assignment right before it |
| 109 | + // because that would just create dead code. |
| 110 | + if let StatementKind::Assign(box ( |
| 111 | + place, |
| 112 | + Rvalue::Use(Operand::Constant(_)), |
| 113 | + )) = &blocks[*succ].statements[used].kind |
| 114 | + { |
| 115 | + if place.local == local && place.projection.is_empty() { |
| 116 | + continue; |
| 117 | + } |
| 118 | + } |
| 119 | + blocks[*succ].statements.insert(used, statement.clone()); |
| 120 | + } else { |
| 121 | + blocks[*succ].statements.push(statement.clone()); |
| 122 | + } |
| 123 | + } |
| 124 | + } |
| 125 | + } |
| 126 | + } |
| 127 | + } |
| 128 | +} |
| 129 | + |
| 130 | +fn branched_locals(body: &Body<'_>) -> BitSet<Local> { |
| 131 | + let mut visitor = BranchedLocals { |
| 132 | + branched: BitSet::new_empty(body.local_decls.len()), |
| 133 | + gets_const_assign: BitSet::new_empty(body.local_decls.len()), |
| 134 | + }; |
| 135 | + visitor.visit_body(body); |
| 136 | + visitor.branched.intersect(&visitor.gets_const_assign); |
| 137 | + visitor.branched |
| 138 | +} |
| 139 | + |
| 140 | +struct BranchedLocals { |
| 141 | + branched: BitSet<Local>, |
| 142 | + gets_const_assign: BitSet<Local>, |
| 143 | +} |
| 144 | + |
| 145 | +impl Visitor<'_> for BranchedLocals { |
| 146 | + fn visit_terminator(&mut self, terminator: &Terminator<'_>, _location: Location) { |
| 147 | + let TerminatorKind::SwitchInt { discr, .. } = &terminator.kind else { return; }; |
| 148 | + if let Some(place) = discr.place() { |
| 149 | + self.branched.insert(place.local); |
| 150 | + } |
| 151 | + } |
| 152 | + |
| 153 | + fn visit_statement(&mut self, statement: &Statement<'_>, _location: Location) { |
| 154 | + if let StatementKind::Assign(box (place, Rvalue::Use(Operand::Constant(_)))) = |
| 155 | + &statement.kind |
| 156 | + { |
| 157 | + if place.projection.is_empty() { |
| 158 | + self.gets_const_assign.insert(place.local); |
| 159 | + } |
| 160 | + } |
| 161 | + } |
| 162 | +} |
| 163 | + |
| 164 | +fn is_empty_resume<'tcx>(block: &BasicBlockData<'tcx>) -> bool { |
| 165 | + block.statements.iter().all(|s| matches!(s.kind, StatementKind::Nop)) |
| 166 | + && block.terminator.as_ref().map(|t| &t.kind) == Some(&TerminatorKind::Resume) |
| 167 | +} |
| 168 | + |
| 169 | +struct CountUsesVisitor { |
| 170 | + counts: FxHashMap<Local, usize>, |
| 171 | + dead: FxHashSet<Local>, |
| 172 | +} |
| 173 | + |
| 174 | +impl CountUsesVisitor { |
| 175 | + fn new() -> Self { |
| 176 | + Self { counts: FxHashMap::default(), dead: FxHashSet::default() } |
| 177 | + } |
| 178 | + |
| 179 | + fn is_used_once(&self, local: Local) -> bool { |
| 180 | + self.counts.get(&local) == Some(&1) |
| 181 | + } |
| 182 | +} |
| 183 | + |
| 184 | +impl Visitor<'_> for CountUsesVisitor { |
| 185 | + fn visit_local(&mut self, local: Local, context: PlaceContext, _location: Location) { |
| 186 | + match context { |
| 187 | + PlaceContext::NonUse(NonUseContext::StorageDead) => { |
| 188 | + self.dead.insert(local); |
| 189 | + } |
| 190 | + PlaceContext::NonUse(_) => {} |
| 191 | + PlaceContext::MutatingUse(_) | PlaceContext::NonMutatingUse(_) => { |
| 192 | + *self.counts.entry(local).or_default() += 1; |
| 193 | + } |
| 194 | + }; |
| 195 | + } |
| 196 | +} |
| 197 | + |
| 198 | +struct SuccessorUsesVisitor { |
| 199 | + local: Local, |
| 200 | + first_use: Option<usize>, |
| 201 | +} |
| 202 | + |
| 203 | +impl SuccessorUsesVisitor { |
| 204 | + fn new(local: Local) -> Self { |
| 205 | + Self { local, first_use: None } |
| 206 | + } |
| 207 | + |
| 208 | + fn first_use(&self) -> Option<usize> { |
| 209 | + self.first_use |
| 210 | + } |
| 211 | +} |
| 212 | + |
| 213 | +impl Visitor<'_> for SuccessorUsesVisitor { |
| 214 | + fn visit_local(&mut self, local: Local, _context: PlaceContext, location: Location) { |
| 215 | + if local == self.local { |
| 216 | + match self.first_use { |
| 217 | + None => self.first_use = Some(location.statement_index), |
| 218 | + Some(first) => self.first_use = Some(first.min(location.statement_index)), |
| 219 | + } |
| 220 | + } |
| 221 | + } |
| 222 | +} |
0 commit comments