compiler/rustc_mir_transform/src/copy_prop.rs

   1 use rustc_index::bit_set::BitSet;
   2 use rustc_index::vec::IndexVec;
   3 use rustc_middle::mir::visit::*;
   4 use rustc_middle::mir::*;
   5 use rustc_middle::ty::TyCtxt;
   6 use rustc_mir_dataflow::impls::borrowed_locals;
   7
   8 use crate::ssa::SsaLocals;
   9 use crate::MirPass;
  10
  11 /// Unify locals that copy each other.
  12 ///
  13 /// We consider patterns of the form
  14 ///   _a = rvalue
  15 ///   _b = move? _a
  16 ///   _c = move? _a
  17 ///   _d = move? _c
  18 /// where each of the locals is only assigned once.
  19 ///
  20 /// We want to replace all those locals by `_a`, either copied or moved.
  21 pub struct CopyProp;
  22
  23 impl<'tcx> MirPass<'tcx> for CopyProp {
  24     fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
  25         sess.mir_opt_level() >= 4
  26     }
  27
  28     #[instrument(level = "trace", skip(self, tcx, body))]
  29     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  30         debug!(def_id = ?body.source.def_id());
  31         propagate_ssa(tcx, body);
  32     }
  33 }
  34
  35 fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
  36     let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
  37     let borrowed_locals = borrowed_locals(body);
  38     let ssa = SsaLocals::new(tcx, param_env, body, &borrowed_locals);
  39
  40     let fully_moved = fully_moved_locals(&ssa, body);
  41     debug!(?fully_moved);
  42
  43     let mut storage_to_remove = BitSet::new_empty(fully_moved.domain_size());
  44     for (local, &head) in ssa.copy_classes().iter_enumerated() {
  45         if local != head {
  46             storage_to_remove.insert(head);
  47         }
  48     }
  49
  50     let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h);
  51
  52     Replacer {
  53         tcx,
  54         copy_classes: &ssa.copy_classes(),
  55         fully_moved,
  56         borrowed_locals,
  57         storage_to_remove,
  58     }
  59     .visit_body_preserves_cfg(body);
  60
  61     if any_replacement {
  62         crate::simplify::remove_unused_definitions(body);
  63     }
  64 }
  65
  66 /// `SsaLocals` computed equivalence classes between locals considering copy/move assignments.
  67 ///
  68 /// This function also returns whether all the `move?` in the pattern are `move` and not copies.
  69 /// A local which is in the bitset can be replaced by `move _a`. Otherwise, it must be
  70 /// replaced by `copy _a`, as we cannot move multiple times from `_a`.
  71 ///
  72 /// If an operand copies `_c`, it must happen before the assignment `_d = _c`, otherwise it is UB.
  73 /// This means that replacing it by a copy of `_a` if ok, since this copy happens before `_c` is
  74 /// moved, and therefore that `_d` is moved.
  75 #[instrument(level = "trace", skip(ssa, body))]
  76 fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> BitSet<Local> {
  77     let mut fully_moved = BitSet::new_filled(body.local_decls.len());
  78
  79     for (_, rvalue) in ssa.assignments(body) {
  80         let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place))
  81             = rvalue
  82         else { continue };
  83
  84         let Some(rhs) = place.as_local() else { continue };
  85         if !ssa.is_ssa(rhs) {
  86             continue;
  87         }
  88
  89         if let Rvalue::Use(Operand::Copy(_)) | Rvalue::CopyForDeref(_) = rvalue {
  90             fully_moved.remove(rhs);
  91         }
  92     }
  93
  94     ssa.meet_copy_equivalence(&mut fully_moved);
  95
  96     fully_moved
  97 }
  98
  99 /// Utility to help performing subtitution of `*pattern` by `target`.
 100 struct Replacer<'a, 'tcx> {
 101     tcx: TyCtxt<'tcx>,
 102     fully_moved: BitSet<Local>,
 103     storage_to_remove: BitSet<Local>,
 104     borrowed_locals: BitSet<Local>,
 105     copy_classes: &'a IndexVec<Local, Local>,
 106 }
 107
 108 impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
 109     fn tcx(&self) -> TyCtxt<'tcx> {
 110         self.tcx
 111     }
 112
 113     fn visit_local(&mut self, local: &mut Local, ctxt: PlaceContext, _: Location) {
 114         let new_local = self.copy_classes[*local];
 115         match ctxt {
 116             // Do not modify the local in storage statements.
 117             PlaceContext::NonUse(NonUseContext::StorageLive | NonUseContext::StorageDead) => {}
 118             // The local should have been marked as non-SSA.
 119             PlaceContext::MutatingUse(_) => assert_eq!(*local, new_local),
 120             // We access the value.
 121             _ => *local = new_local,
 122         }
 123     }
 124
 125     fn visit_place(&mut self, place: &mut Place<'tcx>, ctxt: PlaceContext, loc: Location) {
 126         if let Some(new_projection) = self.process_projection(&place.projection, loc) {
 127             place.projection = self.tcx().intern_place_elems(&new_projection);
 128         }
 129
 130         let observes_address = match ctxt {
 131             PlaceContext::NonMutatingUse(
 132                 NonMutatingUseContext::SharedBorrow
 133                 | NonMutatingUseContext::ShallowBorrow
 134                 | NonMutatingUseContext::UniqueBorrow
 135                 | NonMutatingUseContext::AddressOf,
 136             ) => true,
 137             // For debuginfo, merging locals is ok.
 138             PlaceContext::NonUse(NonUseContext::VarDebugInfo) => {
 139                 self.borrowed_locals.contains(place.local)
 140             }
 141             _ => false,
 142         };
 143         if observes_address && !place.is_indirect() {
 144             // We observe the address of `place.local`. Do not replace it.
 145         } else {
 146             self.visit_local(
 147                 &mut place.local,
 148                 PlaceContext::NonMutatingUse(NonMutatingUseContext::Copy),
 149                 loc,
 150             )
 151         }
 152     }
 153
 154     fn visit_operand(&mut self, operand: &mut Operand<'tcx>, loc: Location) {
 155         if let Operand::Move(place) = *operand
 156             && let Some(local) = place.as_local()
 157             && !self.fully_moved.contains(local)
 158         {
 159             *operand = Operand::Copy(place);
 160         }
 161         self.super_operand(operand, loc);
 162     }
 163
 164     fn visit_statement(&mut self, stmt: &mut Statement<'tcx>, loc: Location) {
 165         match stmt.kind {
 166             // When removing storage statements, we need to remove both (#107511).
 167             StatementKind::StorageLive(l) | StatementKind::StorageDead(l)
 168                 if self.storage_to_remove.contains(l) =>
 169             {
 170                 stmt.make_nop()
 171             }
 172             StatementKind::Assign(box (ref place, ref mut rvalue))
 173                 if place.as_local().is_some() =>
 174             {
 175                 // Do not replace assignments.
 176                 self.visit_rvalue(rvalue, loc)
 177             }
 178             _ => self.super_statement(stmt, loc),
 179         }
 180     }
 181 }