compiler/rustc_mir_transform/src/ssa.rs

   1 use either::Either;
   2 use rustc_data_structures::graph::dominators::Dominators;
   3 use rustc_index::bit_set::BitSet;
   4 use rustc_index::vec::IndexVec;
   5 use rustc_middle::middle::resolve_lifetime::Set1;
   6 use rustc_middle::mir::visit::*;
   7 use rustc_middle::mir::*;
   8 use rustc_middle::ty::{ParamEnv, TyCtxt};
   9
  10 #[derive(Debug)]
  11 pub struct SsaLocals {
  12     /// Assignments to each local. This defines whether the local is SSA.
  13     assignments: IndexVec<Local, Set1<LocationExtended>>,
  14     /// We visit the body in reverse postorder, to ensure each local is assigned before it is used.
  15     /// We remember the order in which we saw the assignments to compute the SSA values in a single
  16     /// pass.
  17     assignment_order: Vec<Local>,
  18     /// Copy equivalence classes between locals. See `copy_classes` for documentation.
  19     copy_classes: IndexVec<Local, Local>,
  20 }
  21
  22 impl SsaLocals {
  23     pub fn new<'tcx>(
  24         tcx: TyCtxt<'tcx>,
  25         param_env: ParamEnv<'tcx>,
  26         body: &Body<'tcx>,
  27         borrowed_locals: &BitSet<Local>,
  28     ) -> SsaLocals {
  29         let assignment_order = Vec::new();
  30
  31         let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls);
  32         let dominators = body.basic_blocks.dominators();
  33         let mut visitor = SsaVisitor { assignments, assignment_order, dominators };
  34
  35         for (local, decl) in body.local_decls.iter_enumerated() {
  36             if matches!(body.local_kind(local), LocalKind::Arg) {
  37                 visitor.assignments[local] = Set1::One(LocationExtended::Arg);
  38             }
  39             if borrowed_locals.contains(local) && !decl.ty.is_freeze(tcx, param_env) {
  40                 visitor.assignments[local] = Set1::Many;
  41             }
  42         }
  43
  44         for (bb, data) in traversal::reverse_postorder(body) {
  45             visitor.visit_basic_block_data(bb, data);
  46         }
  47
  48         for var_debug_info in &body.var_debug_info {
  49             visitor.visit_var_debug_info(var_debug_info);
  50         }
  51
  52         debug!(?visitor.assignments);
  53
  54         visitor
  55             .assignment_order
  56             .retain(|&local| matches!(visitor.assignments[local], Set1::One(_)));
  57         debug!(?visitor.assignment_order);
  58
  59         let copy_classes = compute_copy_classes(&visitor, body);
  60
  61         SsaLocals {
  62             assignments: visitor.assignments,
  63             assignment_order: visitor.assignment_order,
  64             copy_classes,
  65         }
  66     }
  67
  68     pub fn is_ssa(&self, local: Local) -> bool {
  69         matches!(self.assignments[local], Set1::One(_))
  70     }
  71
  72     pub fn assignments<'a, 'tcx>(
  73         &'a self,
  74         body: &'a Body<'tcx>,
  75     ) -> impl Iterator<Item = (Local, &'a Rvalue<'tcx>)> + 'a {
  76         self.assignment_order.iter().filter_map(|&local| {
  77             if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local] {
  78                 // `loc` must point to a direct assignment to `local`.
  79                 let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
  80                 let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() };
  81                 assert_eq!(target.as_local(), Some(local));
  82                 Some((local, rvalue))
  83             } else {
  84                 None
  85             }
  86         })
  87     }
  88
  89     /// Compute the equivalence classes for locals, based on copy statements.
  90     ///
  91     /// The returned vector maps each local to the one it copies. In the following case:
  92     ///   _a = &mut _0
  93     ///   _b = move? _a
  94     ///   _c = move? _a
  95     ///   _d = move? _c
  96     /// We return the mapping
  97     ///   _a => _a // not a copy so, represented by itself
  98     ///   _b => _a
  99     ///   _c => _a
 100     ///   _d => _a // transitively through _c
 101     ///
 102     /// Exception: we do not see through the return place, as it cannot be substituted.
 103     pub fn copy_classes(&self) -> &IndexVec<Local, Local> {
 104         &self.copy_classes
 105     }
 106
 107     /// Make a property uniform on a copy equivalence class by removing elements.
 108     pub fn meet_copy_equivalence(&self, property: &mut BitSet<Local>) {
 109         // Consolidate to have a local iff all its copies are.
 110         //
 111         // `copy_classes` defines equivalence classes between locals. The `local`s that recursively
 112         // move/copy the same local all have the same `head`.
 113         for (local, &head) in self.copy_classes.iter_enumerated() {
 114             // If any copy does not have `property`, then the head is not.
 115             if !property.contains(local) {
 116                 property.remove(head);
 117             }
 118         }
 119         for (local, &head) in self.copy_classes.iter_enumerated() {
 120             // If any copy does not have `property`, then the head doesn't either,
 121             // then no copy has `property`.
 122             if !property.contains(head) {
 123                 property.remove(local);
 124             }
 125         }
 126
 127         // Verify that we correctly computed equivalence classes.
 128         #[cfg(debug_assertions)]
 129         for (local, &head) in self.copy_classes.iter_enumerated() {
 130             assert_eq!(property.contains(local), property.contains(head));
 131         }
 132     }
 133 }
 134
 135 #[derive(Copy, Clone, Debug, PartialEq, Eq)]
 136 enum LocationExtended {
 137     Plain(Location),
 138     Arg,
 139 }
 140
 141 struct SsaVisitor {
 142     dominators: Dominators<BasicBlock>,
 143     assignments: IndexVec<Local, Set1<LocationExtended>>,
 144     assignment_order: Vec<Local>,
 145 }
 146
 147 impl<'tcx> Visitor<'tcx> for SsaVisitor {
 148     fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) {
 149         match ctxt {
 150             PlaceContext::MutatingUse(MutatingUseContext::Store) => {
 151                 self.assignments[local].insert(LocationExtended::Plain(loc));
 152                 self.assignment_order.push(local);
 153             }
 154             // Anything can happen with raw pointers, so remove them.
 155             PlaceContext::NonMutatingUse(NonMutatingUseContext::AddressOf)
 156             | PlaceContext::MutatingUse(_) => self.assignments[local] = Set1::Many,
 157             // Immutable borrows are taken into account in `SsaLocals::new` by
 158             // removing non-freeze locals.
 159             PlaceContext::NonMutatingUse(_) => {
 160                 let set = &mut self.assignments[local];
 161                 let assign_dominates = match *set {
 162                     Set1::Empty | Set1::Many => false,
 163                     Set1::One(LocationExtended::Arg) => true,
 164                     Set1::One(LocationExtended::Plain(assign)) => {
 165                         assign.successor_within_block().dominates(loc, &self.dominators)
 166                     }
 167                 };
 168                 // We are visiting a use that is not dominated by an assignment.
 169                 // Either there is a cycle involved, or we are reading for uninitialized local.
 170                 // Bail out.
 171                 if !assign_dominates {
 172                     *set = Set1::Many;
 173                 }
 174             }
 175             PlaceContext::NonUse(_) => {}
 176         }
 177     }
 178 }
 179
 180 #[instrument(level = "trace", skip(ssa, body))]
 181 fn compute_copy_classes(ssa: &SsaVisitor, body: &Body<'_>) -> IndexVec<Local, Local> {
 182     let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len());
 183
 184     for &local in &ssa.assignment_order {
 185         debug!(?local);
 186
 187         if local == RETURN_PLACE {
 188             // `_0` is special, we cannot rename it.
 189             continue;
 190         }
 191
 192         // This is not SSA: mark that we don't know the value.
 193         debug!(assignments = ?ssa.assignments[local]);
 194         let Set1::One(LocationExtended::Plain(loc)) = ssa.assignments[local] else { continue };
 195
 196         // `loc` must point to a direct assignment to `local`.
 197         let Either::Left(stmt) = body.stmt_at(loc) else { bug!() };
 198         let Some((_target, rvalue)) = stmt.kind.as_assign() else { bug!() };
 199         assert_eq!(_target.as_local(), Some(local));
 200
 201         let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place))
 202             = rvalue
 203         else { continue };
 204
 205         let Some(rhs) = place.as_local() else { continue };
 206         let Set1::One(_) = ssa.assignments[rhs] else { continue };
 207
 208         // We visit in `assignment_order`, ie. reverse post-order, so `rhs` has been
 209         // visited before `local`, and we just have to copy the representing local.
 210         copies[local] = copies[rhs];
 211     }
 212
 213     debug!(?copies);
 214
 215     // Invariant: `copies` must point to the head of an equivalence class.
 216     #[cfg(debug_assertions)]
 217     for &head in copies.iter() {
 218         assert_eq!(copies[head], head);
 219     }
 220
 221     copies
 222 }