diff --git a/compiler/rustc_mir_transform/src/copy_prop.rs b/compiler/rustc_mir_transform/src/copy_prop.rs index 8b024721718..58211bc795f 100644 --- a/compiler/rustc_mir_transform/src/copy_prop.rs +++ b/compiler/rustc_mir_transform/src/copy_prop.rs @@ -1,13 +1,10 @@ -use either::Either; -use rustc_data_structures::graph::dominators::Dominators; use rustc_index::bit_set::BitSet; use rustc_index::vec::IndexVec; -use rustc_middle::middle::resolve_lifetime::Set1; use rustc_middle::mir::visit::*; use rustc_middle::mir::*; -use rustc_middle::ty::{ParamEnv, TyCtxt}; -use rustc_mir_dataflow::impls::borrowed_locals; +use rustc_middle::ty::TyCtxt; +use crate::ssa::SsaLocals; use crate::MirPass; /// Unify locals that copy each other. @@ -38,123 +35,28 @@ fn propagate_ssa<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id()); let ssa = SsaLocals::new(tcx, param_env, body); - let (copy_classes, fully_moved) = compute_copy_classes(&ssa, body); - debug!(?copy_classes); + let fully_moved = fully_moved_locals(&ssa, body); + debug!(?fully_moved); let mut storage_to_remove = BitSet::new_empty(fully_moved.domain_size()); - for (local, &head) in copy_classes.iter_enumerated() { + for (local, &head) in ssa.copy_classes().iter_enumerated() { if local != head { storage_to_remove.insert(head); storage_to_remove.insert(local); } } - let any_replacement = copy_classes.iter_enumerated().any(|(l, &h)| l != h); + let any_replacement = ssa.copy_classes().iter_enumerated().any(|(l, &h)| l != h); - Replacer { tcx, copy_classes, fully_moved, storage_to_remove }.visit_body_preserves_cfg(body); + Replacer { tcx, copy_classes: &ssa.copy_classes(), fully_moved, storage_to_remove } + .visit_body_preserves_cfg(body); if any_replacement { crate::simplify::remove_unused_definitions(body); } } -#[derive(Copy, Clone, Debug, PartialEq, Eq)] -enum LocationExtended { - Plain(Location), - Arg, -} - -#[derive(Debug)] -struct SsaLocals { - dominators: Dominators, - /// Assignments to each local. This defines whether the local is SSA. - assignments: IndexVec>, - /// We visit the body in reverse postorder, to ensure each local is assigned before it is used. - /// We remember the order in which we saw the assignments to compute the SSA values in a single - /// pass. - assignment_order: Vec, -} - -impl SsaLocals { - fn new<'tcx>(tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>, body: &Body<'tcx>) -> SsaLocals { - let assignment_order = Vec::new(); - - let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls); - let dominators = body.basic_blocks.dominators(); - let mut this = SsaLocals { assignments, assignment_order, dominators }; - - let borrowed = borrowed_locals(body); - for (local, decl) in body.local_decls.iter_enumerated() { - if matches!(body.local_kind(local), LocalKind::Arg) { - this.assignments[local] = Set1::One(LocationExtended::Arg); - } - if borrowed.contains(local) && !decl.ty.is_freeze(tcx, param_env) { - this.assignments[local] = Set1::Many; - } - } - - for (bb, data) in traversal::reverse_postorder(body) { - this.visit_basic_block_data(bb, data); - } - - for var_debug_info in &body.var_debug_info { - this.visit_var_debug_info(var_debug_info); - } - - debug!(?this.assignments); - - this.assignment_order.retain(|&local| matches!(this.assignments[local], Set1::One(_))); - debug!(?this.assignment_order); - - this - } -} - -impl<'tcx> Visitor<'tcx> for SsaLocals { - fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) { - match ctxt { - PlaceContext::MutatingUse(MutatingUseContext::Store) => { - self.assignments[local].insert(LocationExtended::Plain(loc)); - self.assignment_order.push(local); - } - // Anything can happen with raw pointers, so remove them. - PlaceContext::NonMutatingUse(NonMutatingUseContext::AddressOf) - | PlaceContext::MutatingUse(_) => self.assignments[local] = Set1::Many, - // Immutable borrows are taken into account in `SsaLocals::new` by - // removing non-freeze locals. - PlaceContext::NonMutatingUse(_) => { - let set = &mut self.assignments[local]; - let assign_dominates = match *set { - Set1::Empty | Set1::Many => false, - Set1::One(LocationExtended::Arg) => true, - Set1::One(LocationExtended::Plain(assign)) => { - assign.dominates(loc, &self.dominators) - } - }; - // We are visiting a use that is not dominated by an assignment. - // Either there is a cycle involved, or we are reading for uninitialized local. - // Bail out. - if !assign_dominates { - *set = Set1::Many; - } - } - PlaceContext::NonUse(_) => {} - } - } -} - -/// Compute the equivalence classes for locals, based on copy statements. -/// -/// The returned vector maps each local to the one it copies. In the following case: -/// _a = &mut _0 -/// _b = move? _a -/// _c = move? _a -/// _d = move? _c -/// We return the mapping -/// _a => _a // not a copy so, represented by itself -/// _b => _a -/// _c => _a -/// _d => _a // transitively through _c +/// `SsaLocals` computed equivalence classes between locals considering copy/move assignments. /// /// This function also returns whether all the `move?` in the pattern are `move` and not copies. /// A local which is in the bitset can be replaced by `move _a`. Otherwise, it must be @@ -164,95 +66,38 @@ impl<'tcx> Visitor<'tcx> for SsaLocals { /// This means that replacing it by a copy of `_a` if ok, since this copy happens before `_c` is /// moved, and therefore that `_d` is moved. #[instrument(level = "trace", skip(ssa, body))] -fn compute_copy_classes( - ssa: &SsaLocals, - body: &Body<'_>, -) -> (IndexVec, BitSet) { - let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len()); - let mut fully_moved = BitSet::new_filled(copies.len()); - - for &local in &ssa.assignment_order { - debug!(?local); - - if local == RETURN_PLACE { - // `_0` is special, we cannot rename it. - continue; - } - - // This is not SSA: mark that we don't know the value. - debug!(assignments = ?ssa.assignments[local]); - let Set1::One(LocationExtended::Plain(loc)) = ssa.assignments[local] else { continue }; - - // `loc` must point to a direct assignment to `local`. - let Either::Left(stmt) = body.stmt_at(loc) else { bug!() }; - let Some((_target, rvalue)) = stmt.kind.as_assign() else { bug!() }; - assert_eq!(_target.as_local(), Some(local)); +fn fully_moved_locals(ssa: &SsaLocals, body: &Body<'_>) -> BitSet { + let mut fully_moved = BitSet::new_filled(body.local_decls.len()); + for (_, rvalue) in ssa.assignments(body) { let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place)) = rvalue else { continue }; let Some(rhs) = place.as_local() else { continue }; - let Set1::One(_) = ssa.assignments[rhs] else { continue }; - - // We visit in `assignment_order`, ie. reverse post-order, so `rhs` has been - // visited before `local`, and we just have to copy the representing local. - copies[local] = copies[rhs]; + if !ssa.is_ssa(rhs) { + continue; + } if let Rvalue::Use(Operand::Copy(_)) | Rvalue::CopyForDeref(_) = rvalue { fully_moved.remove(rhs); } } - debug!(?copies); + ssa.meet_copy_equivalence(&mut fully_moved); - // Invariant: `copies` must point to the head of an equivalence class. - #[cfg(debug_assertions)] - for &head in copies.iter() { - assert_eq!(copies[head], head); - } - - meet_copy_equivalence(&copies, &mut fully_moved); - - (copies, fully_moved) -} - -/// Make a property uniform on a copy equivalence class by removing elements. -fn meet_copy_equivalence(copies: &IndexVec, property: &mut BitSet) { - // Consolidate to have a local iff all its copies are. - // - // `copies` defines equivalence classes between locals. The `local`s that recursively - // move/copy the same local all have the same `head`. - for (local, &head) in copies.iter_enumerated() { - // If any copy does not have `property`, then the head is not. - if !property.contains(local) { - property.remove(head); - } - } - for (local, &head) in copies.iter_enumerated() { - // If any copy does not have `property`, then the head doesn't either, - // then no copy has `property`. - if !property.contains(head) { - property.remove(local); - } - } - - // Verify that we correctly computed equivalence classes. - #[cfg(debug_assertions)] - for (local, &head) in copies.iter_enumerated() { - assert_eq!(property.contains(local), property.contains(head)); - } + fully_moved } /// Utility to help performing subtitution of `*pattern` by `target`. -struct Replacer<'tcx> { +struct Replacer<'a, 'tcx> { tcx: TyCtxt<'tcx>, fully_moved: BitSet, storage_to_remove: BitSet, - copy_classes: IndexVec, + copy_classes: &'a IndexVec, } -impl<'tcx> MutVisitor<'tcx> for Replacer<'tcx> { +impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> { fn tcx(&self) -> TyCtxt<'tcx> { self.tcx } diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs index 1330f1e26b8..8a7c14472ac 100644 --- a/compiler/rustc_mir_transform/src/lib.rs +++ b/compiler/rustc_mir_transform/src/lib.rs @@ -87,6 +87,7 @@ mod required_consts; mod reveal_all; mod separate_const_switch; mod shim; +mod ssa; // This pass is public to allow external drivers to perform MIR cleanup pub mod simplify; mod simplify_branches; diff --git a/compiler/rustc_mir_transform/src/ssa.rs b/compiler/rustc_mir_transform/src/ssa.rs new file mode 100644 index 00000000000..a5ae671be81 --- /dev/null +++ b/compiler/rustc_mir_transform/src/ssa.rs @@ -0,0 +1,219 @@ +use either::Either; +use rustc_data_structures::graph::dominators::Dominators; +use rustc_index::bit_set::BitSet; +use rustc_index::vec::IndexVec; +use rustc_middle::middle::resolve_lifetime::Set1; +use rustc_middle::mir::visit::*; +use rustc_middle::mir::*; +use rustc_middle::ty::{ParamEnv, TyCtxt}; +use rustc_mir_dataflow::impls::borrowed_locals; + +#[derive(Debug)] +pub struct SsaLocals { + /// Assignments to each local. This defines whether the local is SSA. + assignments: IndexVec>, + /// We visit the body in reverse postorder, to ensure each local is assigned before it is used. + /// We remember the order in which we saw the assignments to compute the SSA values in a single + /// pass. + assignment_order: Vec, + /// Copy equivalence classes between locals. See `copy_classes` for documentation. + copy_classes: IndexVec, +} + +impl SsaLocals { + pub fn new<'tcx>(tcx: TyCtxt<'tcx>, param_env: ParamEnv<'tcx>, body: &Body<'tcx>) -> SsaLocals { + let assignment_order = Vec::new(); + + let assignments = IndexVec::from_elem(Set1::Empty, &body.local_decls); + let dominators = body.basic_blocks.dominators(); + let mut visitor = SsaVisitor { assignments, assignment_order, dominators }; + + let borrowed = borrowed_locals(body); + for (local, decl) in body.local_decls.iter_enumerated() { + if matches!(body.local_kind(local), LocalKind::Arg) { + visitor.assignments[local] = Set1::One(LocationExtended::Arg); + } + if borrowed.contains(local) && !decl.ty.is_freeze(tcx, param_env) { + visitor.assignments[local] = Set1::Many; + } + } + + for (bb, data) in traversal::reverse_postorder(body) { + visitor.visit_basic_block_data(bb, data); + } + + for var_debug_info in &body.var_debug_info { + visitor.visit_var_debug_info(var_debug_info); + } + + debug!(?visitor.assignments); + + visitor + .assignment_order + .retain(|&local| matches!(visitor.assignments[local], Set1::One(_))); + debug!(?visitor.assignment_order); + + let copy_classes = compute_copy_classes(&visitor, body); + + SsaLocals { + assignments: visitor.assignments, + assignment_order: visitor.assignment_order, + copy_classes, + } + } + + pub fn is_ssa(&self, local: Local) -> bool { + matches!(self.assignments[local], Set1::One(_)) + } + + pub fn assignments<'a, 'tcx>( + &'a self, + body: &'a Body<'tcx>, + ) -> impl Iterator)> + 'a { + self.assignment_order.iter().filter_map(|&local| { + if let Set1::One(LocationExtended::Plain(loc)) = self.assignments[local] { + // `loc` must point to a direct assignment to `local`. + let Either::Left(stmt) = body.stmt_at(loc) else { bug!() }; + let Some((target, rvalue)) = stmt.kind.as_assign() else { bug!() }; + assert_eq!(target.as_local(), Some(local)); + Some((local, rvalue)) + } else { + None + } + }) + } + + /// Compute the equivalence classes for locals, based on copy statements. + /// + /// The returned vector maps each local to the one it copies. In the following case: + /// _a = &mut _0 + /// _b = move? _a + /// _c = move? _a + /// _d = move? _c + /// We return the mapping + /// _a => _a // not a copy so, represented by itself + /// _b => _a + /// _c => _a + /// _d => _a // transitively through _c + /// + /// Exception: we do not see through the return place, as it cannot be substituted. + pub fn copy_classes(&self) -> &IndexVec { + &self.copy_classes + } + + /// Make a property uniform on a copy equivalence class by removing elements. + pub fn meet_copy_equivalence(&self, property: &mut BitSet) { + // Consolidate to have a local iff all its copies are. + // + // `copy_classes` defines equivalence classes between locals. The `local`s that recursively + // move/copy the same local all have the same `head`. + for (local, &head) in self.copy_classes.iter_enumerated() { + // If any copy does not have `property`, then the head is not. + if !property.contains(local) { + property.remove(head); + } + } + for (local, &head) in self.copy_classes.iter_enumerated() { + // If any copy does not have `property`, then the head doesn't either, + // then no copy has `property`. + if !property.contains(head) { + property.remove(local); + } + } + + // Verify that we correctly computed equivalence classes. + #[cfg(debug_assertions)] + for (local, &head) in self.copy_classes.iter_enumerated() { + assert_eq!(property.contains(local), property.contains(head)); + } + } +} + +#[derive(Copy, Clone, Debug, PartialEq, Eq)] +enum LocationExtended { + Plain(Location), + Arg, +} + +struct SsaVisitor { + dominators: Dominators, + assignments: IndexVec>, + assignment_order: Vec, +} + +impl<'tcx> Visitor<'tcx> for SsaVisitor { + fn visit_local(&mut self, local: Local, ctxt: PlaceContext, loc: Location) { + match ctxt { + PlaceContext::MutatingUse(MutatingUseContext::Store) => { + self.assignments[local].insert(LocationExtended::Plain(loc)); + self.assignment_order.push(local); + } + // Anything can happen with raw pointers, so remove them. + PlaceContext::NonMutatingUse(NonMutatingUseContext::AddressOf) + | PlaceContext::MutatingUse(_) => self.assignments[local] = Set1::Many, + // Immutable borrows are taken into account in `SsaLocals::new` by + // removing non-freeze locals. + PlaceContext::NonMutatingUse(_) => { + let set = &mut self.assignments[local]; + let assign_dominates = match *set { + Set1::Empty | Set1::Many => false, + Set1::One(LocationExtended::Arg) => true, + Set1::One(LocationExtended::Plain(assign)) => { + assign.dominates(loc, &self.dominators) + } + }; + // We are visiting a use that is not dominated by an assignment. + // Either there is a cycle involved, or we are reading for uninitialized local. + // Bail out. + if !assign_dominates { + *set = Set1::Many; + } + } + PlaceContext::NonUse(_) => {} + } + } +} + +#[instrument(level = "trace", skip(ssa, body))] +fn compute_copy_classes(ssa: &SsaVisitor, body: &Body<'_>) -> IndexVec { + let mut copies = IndexVec::from_fn_n(|l| l, body.local_decls.len()); + + for &local in &ssa.assignment_order { + debug!(?local); + + if local == RETURN_PLACE { + // `_0` is special, we cannot rename it. + continue; + } + + // This is not SSA: mark that we don't know the value. + debug!(assignments = ?ssa.assignments[local]); + let Set1::One(LocationExtended::Plain(loc)) = ssa.assignments[local] else { continue }; + + // `loc` must point to a direct assignment to `local`. + let Either::Left(stmt) = body.stmt_at(loc) else { bug!() }; + let Some((_target, rvalue)) = stmt.kind.as_assign() else { bug!() }; + assert_eq!(_target.as_local(), Some(local)); + + let (Rvalue::Use(Operand::Copy(place) | Operand::Move(place)) | Rvalue::CopyForDeref(place)) + = rvalue + else { continue }; + + let Some(rhs) = place.as_local() else { continue }; + let Set1::One(_) = ssa.assignments[rhs] else { continue }; + + // We visit in `assignment_order`, ie. reverse post-order, so `rhs` has been + // visited before `local`, and we just have to copy the representing local. + copies[local] = copies[rhs]; + } + + debug!(?copies); + + // Invariant: `copies` must point to the head of an equivalence class. + #[cfg(debug_assertions)] + for &head in copies.iter() { + assert_eq!(copies[head], head); + } + + copies +}