|
| 1 | +use either::Either; |
| 2 | +use rustc_index::bit_set::BitSet; |
| 3 | +use rustc_index::vec::IndexVec; |
| 4 | +use rustc_middle::middle::resolve_lifetime::Set1; |
| 5 | +use rustc_middle::mir::visit::*; |
| 6 | +use rustc_middle::mir::*; |
| 7 | +use rustc_middle::ty::{ParamEnv, TyCtxt}; |
| 8 | +use rustc_mir_dataflow::impls::borrowed_locals; |
| 9 | + |
| 10 | +use crate::MirPass; |
| 11 | + |
| 12 | +/// Unify locals that copy each other. |
| 13 | +/// |
| 14 | +/// We consider patterns of the form |
| 15 | +/// _a = rvalue |
| 16 | +/// _b = move? _a |
| 17 | +/// _c = move? _a |
| 18 | +/// _d = move? _c |
| 19 | +/// where each of the locals is only assigned once. |
| 20 | +/// |
| 21 | +/// We want to replace all those locals by `_a`, either copied or moved. |
| 22 | +pub struct CopyProp; |
| 23 | + |
| 24 | +impl<'tcx> MirPass<'tcx> for CopyProp { |
| 25 | + fn is_enabled(&self, sess: &rustc_session::Session) -> bool { |
| 26 | + sess.mir_opt_level() >= 4 |
| 27 | + } |
| 28 | + |
| 29 | + #[instrument(level = "trace", skip(self, tcx, body))] |
| 30 | + fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 31 | + debug!(def_id = ?body.source.def_id()); |
| 32 | + propagate_ssa(tcx, body); |
| 33 | + } |
| 34 | +} |
| 35 | + |
| 36 | +fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
| 37 | + let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); |
| 38 | + let ssa = SsaLocals::new(tcx, param_env, body); |
| 39 | + |
| 40 | + let (copy_classes, fully_moved) = compute_copy_classes(&ssa, body); |
| 41 | + debug!(?copy_classes); |
| 42 | + |
| 43 | + let mut storage_to_remove = BitSet::new_empty(fully_moved.domain_size()); |
| 44 | + for (local, &head) in copy_classes.iter_enumerated() { |
| 45 | + if local != head { |
| 46 | + storage_to_remove.insert(head); |
| 47 | + storage_to_remove.insert(local); |
| 48 | + } |
| 49 | + } |
| 50 | + |
| 51 | + let any_replacement = copy_classes.iter_enumerated().any(|(l, &h)| l != h); |
| 52 | + |
| 53 | + Replacer { tcx, copy_classes, fully_moved, storage_to_remove }.visit_body_preserves_cfg(body); |
| 54 | + |
| 55 | + if any_replacement { |
| 56 | + crate::simplify::remove_unused_definitions(body); |
| 57 | + } |
| 58 | +} |
| 59 | + |
| 60 | +#[derive(Copy, Clone, Debug, PartialEq, Eq)] |
| 61 | +enum LocationExtended { |
| 62 | + Plain(Location), |
| 63 | + Arg, |
| 64 | +} |
| 65 | + |
| 66 | +#[derive(Debug)] |
| 67 | +struct SsaLocals { |
| 68 | + /// Assignments to each local. This defines whether the local is SSA. |
| 69 | + assignments: IndexVec<Local, Set1<LocationExtended>>, |
| 70 | + /// We visit the body in reverse postorder, to ensure each local is assigned before it is used. |
| 71 | + /// We remember the order in which we saw the assignments to compute the SSA values in a single |
| 72 | + /// pass. |
| 73 | + assignment_order: Vec<Local>, |
| 74 | +} |
| 75 | + |
| 76 | +impl SsaLocals { |
| 77 | + fn new<'tcx>(tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>, body: &Body<'tcx>) -> SsaLocals { |
| 78 | + let assignment_order = Vec::new(); |
| 79 | + |
| 80 | + let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls); |
| 81 | + let mut this = SsaLocals { assignments, assignment_order }; |
| 82 | + |
| 83 | + let borrowed = borrowed_locals(body); |
| 84 | + for (local, decl) in body.local_decls.iter_enumerated() { |
| 85 | + if matches!(body.local_kind(local), LocalKind::Arg) { |
| 86 | + this.assignments[local] = Set1::One(LocationExtended::Arg); |
| 87 | + } |
| 88 | + if borrowed.contains(local) && !decl.ty.is_freeze(tcx, param_env) { |
| 89 | + this.assignments[local] = Set1::Many; |
| 90 | + } |
| 91 | + } |
| 92 | + |
| 93 | + for (bb, data) in traversal::reverse_postorder(body) { |
| 94 | + this.visit_basic_block_data(bb, data); |
| 95 | + } |
| 96 | + |
| 97 | + for var_debug_info in &body.var_debug_info { |
| 98 | + this.visit_var_debug_info(var_debug_info); |
| 99 | + } |
| 100 | + |
| 101 | + debug!(?this.assignments); |
| 102 | + |
| 103 | + this.assignment_order.retain(|&local| matches!(this.assignments[local], Set1::One(_))); |
| 104 | + debug!(?this.assignment_order); |
| 105 | + |
| 106 | + this |
| 107 | + } |
| 108 | +} |
| 109 | + |
| 110 | +impl<'tcx> Visitor<'tcx> for SsaLocals { |
| 111 | + fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) { |
| 112 | + match ctxt { |
| 113 | + PlaceContext::MutatingUse(MutatingUseContext::Store) => { |
| 114 | + self.assignments[local].insert(LocationExtended::Plain(loc)); |
| 115 | + self.assignment_order.push(local); |
| 116 | + } |
| 117 | + PlaceContext::MutatingUse(_) => self.assignments[local] = Set1::Many, |
| 118 | + // Immutable borrows and AddressOf are taken into account in `SsaLocals::new` by |
| 119 | + // removing non-freeze locals. |
| 120 | + PlaceContext::NonMutatingUse(_) | PlaceContext::NonUse(_) => {} |
| 121 | + } |
| 122 | + } |
| 123 | +} |
| 124 | + |
| 125 | +/// Compute the equivalence classes for locals, based on copy statements. |
| 126 | +/// |
| 127 | +/// The returned vector maps each local to the one it copies. In the following case: |
| 128 | +/// _a = &mut _0 |
| 129 | +/// _b = move? _a |
| 130 | +/// _c = move? _a |
| 131 | +/// _d = move? _c |
| 132 | +/// We return the mapping |
| 133 | +/// _a => _a // not a copy so, represented by itself |
| 134 | +/// _b => _a |
| 135 | +/// _c => _a |
| 136 | +/// _d => _a // transitively through _c |
| 137 | +/// |
| 138 | +/// This function also returns whether all the `move?` in the pattern are `move` and not copies. |
| 139 | +/// A local which is in the bitset can be replaced by `move _a`. Otherwise, it must be |
| 140 | +/// replaced by `copy _a`, as we cannot move multiple times from `_a`. |
| 141 | +/// |
| 142 | +/// If an operand copies `_c`, it must happen before the assignment `_d = _c`, otherwise it is UB. |
| 143 | +/// This means that replacing it by a copy of `_a` if ok, since this copy happens before `_c` is |
| 144 | +/// moved, and therefore that `_d` is moved. |
| 145 | +#[instrument(level = "trace", skip(ssa, body))] |
| 146 | +fn compute_copy_classes( |
| 147 | + ssa: &SsaLocals, |
| 148 | + body: &Body<'_>, |
| 149 | +) -> (IndexVec<Local, Local>, BitSet<Local>) { |
| 150 | + let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len()); |
| 151 | + let mut fully_moved = BitSet::new_filled(copies.len()); |
| 152 | + |
| 153 | + for &local in &ssa.assignment_order { |
| 154 | + debug!(?local); |
| 155 | + |
| 156 | + if local == RETURN_PLACE { |
| 157 | + // `_0` is special, we cannot rename it. |
| 158 | + continue; |
| 159 | + } |
| 160 | + |
| 161 | + // This is not SSA: mark that we don't know the value. |
| 162 | + debug!(assignments = ?ssa.assignments[local]); |
| 163 | + let Set1::One(LocationExtended::Plain(loc)) = ssa.assignments[local] else { continue }; |
| 164 | + |
| 165 | + // `loc` must point to a direct assignment to `local`. |
| 166 | + let Either::Left(stmt) = body.stmt_at(loc) else { bug!() }; |
| 167 | + let Some((_target, rvalue)) = stmt.kind.as_assign() else { bug!() }; |
| 168 | + assert_eq!(_target.as_local(), Some(local)); |
| 169 | + |
| 170 | + let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place)) |
| 171 | + = rvalue |
| 172 | + else { continue }; |
| 173 | + |
| 174 | + let Some(rhs) = place.as_local() else { continue }; |
| 175 | + let Set1::One(_) = ssa.assignments[rhs] else { continue }; |
| 176 | + |
| 177 | + // We visit in `assignment_order`, ie. reverse post-order, so `rhs` has been |
| 178 | + // visited before `local`, and we just have to copy the representing local. |
| 179 | + copies[local] = copies[rhs]; |
| 180 | + |
| 181 | + if let Rvalue::Use(Operand::Copy(_)) | Rvalue::CopyForDeref(_) = rvalue { |
| 182 | + fully_moved.remove(rhs); |
| 183 | + } |
| 184 | + } |
| 185 | + |
| 186 | + debug!(?copies); |
| 187 | + |
| 188 | + // Invariant: `copies` must point to the head of an equivalence class. |
| 189 | + #[cfg(debug_assertions)] |
| 190 | + for &head in copies.iter() { |
| 191 | + assert_eq!(copies[head], head); |
| 192 | + } |
| 193 | + |
| 194 | + meet_copy_equivalence(&copies, &mut fully_moved); |
| 195 | + |
| 196 | + (copies, fully_moved) |
| 197 | +} |
| 198 | + |
| 199 | +/// Make a property uniform on a copy equivalence class by removing elements. |
| 200 | +fn meet_copy_equivalence(copies: &IndexVec<Local, Local>, property: &mut BitSet<Local>) { |
| 201 | + // Consolidate to have a local iff all its copies are. |
| 202 | + // |
| 203 | + // `copies` defines equivalence classes between locals. The `local`s that recursively |
| 204 | + // move/copy the same local all have the same `head`. |
| 205 | + for (local, &head) in copies.iter_enumerated() { |
| 206 | + // If any copy does not have `property`, then the head is not. |
| 207 | + if !property.contains(local) { |
| 208 | + property.remove(head); |
| 209 | + } |
| 210 | + } |
| 211 | + for (local, &head) in copies.iter_enumerated() { |
| 212 | + // If any copy does not have `property`, then the head doesn't either, |
| 213 | + // then no copy has `property`. |
| 214 | + if !property.contains(head) { |
| 215 | + property.remove(local); |
| 216 | + } |
| 217 | + } |
| 218 | + |
| 219 | + // Verify that we correctly computed equivalence classes. |
| 220 | + #[cfg(debug_assertions)] |
| 221 | + for (local, &head) in copies.iter_enumerated() { |
| 222 | + assert_eq!(property.contains(local), property.contains(head)); |
| 223 | + } |
| 224 | +} |
| 225 | + |
| 226 | +/// Utility to help performing subtitution of `*pattern` by `target`. |
| 227 | +struct Replacer<'tcx> { |
| 228 | + tcx: TyCtxt<'tcx>, |
| 229 | + fully_moved: BitSet<Local>, |
| 230 | + storage_to_remove: BitSet<Local>, |
| 231 | + copy_classes: IndexVec<Local, Local>, |
| 232 | +} |
| 233 | + |
| 234 | +impl<'tcx> MutVisitor<'tcx> for Replacer<'tcx> { |
| 235 | + fn tcx(&self) -> TyCtxt<'tcx> { |
| 236 | + self.tcx |
| 237 | + } |
| 238 | + |
| 239 | + fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { |
| 240 | + *local = self.copy_classes[*local]; |
| 241 | + } |
| 242 | + |
| 243 | + fn visit_operand(&mut self, operand: &mut Operand<'tcx>, loc: Location) { |
| 244 | + if let Operand::Move(place) = *operand |
| 245 | + && let Some(local) = place.as_local() |
| 246 | + && !self.fully_moved.contains(local) |
| 247 | + { |
| 248 | + *operand = Operand::Copy(place); |
| 249 | + } |
| 250 | + self.super_operand(operand, loc); |
| 251 | + } |
| 252 | + |
| 253 | + fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) { |
| 254 | + if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = stmt.kind |
| 255 | + && self.storage_to_remove.contains(l) |
| 256 | + { |
| 257 | + stmt.make_nop(); |
| 258 | + } |
| 259 | + if let StatementKind::Assign(box (ref place, _)) = stmt.kind |
| 260 | + && let Some(l) = place.as_local() |
| 261 | + && self.copy_classes[l] != l |
| 262 | + { |
| 263 | + stmt.make_nop(); |
| 264 | + } |
| 265 | + self.super_statement(stmt, loc); |
| 266 | + } |
| 267 | +} |
0 commit comments