compiler/rustc_mir/src/transform/nrvo.rs

   1 //! See the docs for [`RenameReturnPlace`].
   2
   3 use rustc_hir::Mutability;
   4 use rustc_index::bit_set::HybridBitSet;
   5 use rustc_middle::mir::visit::{MutVisitor, NonUseContext, PlaceContext, Visitor};
   6 use rustc_middle::mir::{self, BasicBlock, Local, Location};
   7 use rustc_middle::ty::TyCtxt;
   8
   9 use crate::transform::MirPass;
  10
  11 /// This pass looks for MIR that always copies the same local into the return place and eliminates
  12 /// the copy by renaming all uses of that local to `_0`.
  13 ///
  14 /// This allows LLVM to perform an optimization similar to the named return value optimization
  15 /// (NRVO) that is guaranteed in C++. This avoids a stack allocation and `memcpy` for the
  16 /// relatively common pattern of allocating a buffer on the stack, mutating it, and returning it by
  17 /// value like so:
  18 ///
  19 /// ```rust
  20 /// fn foo(init: fn(&mut [u8; 1024])) -> [u8; 1024] {
  21 ///     let mut buf = [0; 1024];
  22 ///     init(&mut buf);
  23 ///     buf
  24 /// }
  25 /// ```
  26 ///
  27 /// For now, this pass is very simple and only capable of eliminating a single copy. A more general
  28 /// version of copy propagation, such as the one based on non-overlapping live ranges in [#47954] and
  29 /// [#71003], could yield even more benefits.
  30 ///
  31 /// [#47954]: https://github.com/rust-lang/rust/pull/47954
  32 /// [#71003]: https://github.com/rust-lang/rust/pull/71003
  33 pub struct RenameReturnPlace;
  34
  35 impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
  36     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
  37         if tcx.sess.opts.debugging_opts.mir_opt_level == 0 {
  38             return;
  39         }
  40
  41         let returned_local = match local_eligible_for_nrvo(body) {
  42             Some(l) => l,
  43             None => {
  44                 debug!("`{:?}` was ineligible for NRVO", body.source.def_id());
  45                 return;
  46             }
  47         };
  48
  49         debug!(
  50             "`{:?}` was eligible for NRVO, making {:?} the return place",
  51             body.source.def_id(),
  52             returned_local
  53         );
  54
  55         RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);
  56
  57         // Clean up the `NOP`s we inserted for statements made useless by our renaming.
  58         for block_data in body.basic_blocks_mut() {
  59             block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
  60         }
  61
  62         // Overwrite the debuginfo of `_0` with that of the renamed local.
  63         let (renamed_decl, ret_decl) =
  64             body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);
  65
  66         // Sometimes, the return place is assigned a local of a different but coercable type, for
  67         // example `&mut T` instead of `&T`. Overwriting the `LocalInfo` for the return place means
  68         // its type may no longer match the return type of its function. This doesn't cause a
  69         // problem in codegen because these two types are layout-compatible, but may be unexpected.
  70         debug!("_0: {:?} = {:?}: {:?}", ret_decl.ty, returned_local, renamed_decl.ty);
  71         ret_decl.clone_from(renamed_decl);
  72
  73         // The return place is always mutable.
  74         ret_decl.mutability = Mutability::Mut;
  75     }
  76 }
  77
  78 /// MIR that is eligible for the NRVO must fulfill two conditions:
  79 ///   1. The return place must not be read prior to the `Return` terminator.
  80 ///   2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the
  81 ///      only definition of the return place reaching the `Return` terminator.
  82 ///
  83 /// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
  84 /// to the return place along all possible paths through the control-flow graph.
  85 fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
  86     if IsReturnPlaceRead::run(body) {
  87         return None;
  88     }
  89
  90     let mut copied_to_return_place = None;
  91     for block in body.basic_blocks().indices() {
  92         // Look for blocks with a `Return` terminator.
  93         if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
  94             continue;
  95         }
  96
  97         // Look for an assignment of a single local to the return place prior to the `Return`.
  98         let returned_local = find_local_assigned_to_return_place(block, body)?;
  99         match body.local_kind(returned_local) {
 100             // FIXME: Can we do this for arguments as well?
 101             mir::LocalKind::Arg => return None,
 102
 103             mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
 104             mir::LocalKind::Var | mir::LocalKind::Temp => {}
 105         }
 106
 107         // If multiple different locals are copied to the return place. We can't pick a
 108         // single one to rename.
 109         if copied_to_return_place.map_or(false, |old| old != returned_local) {
 110             return None;
 111         }
 112
 113         copied_to_return_place = Some(returned_local);
 114     }
 115
 116     copied_to_return_place
 117 }
 118
 119 fn find_local_assigned_to_return_place(
 120     start: BasicBlock,
 121     body: &mut mir::Body<'_>,
 122 ) -> Option<Local> {
 123     let mut block = start;
 124     let mut seen = HybridBitSet::new_empty(body.basic_blocks().len());
 125
 126     // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
 127     while seen.insert(block) {
 128         trace!("Looking for assignments to `_0` in {:?}", block);
 129
 130         let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
 131         if local.is_some() {
 132             return local;
 133         }
 134
 135         match body.predecessors()[block].as_slice() {
 136             &[pred] => block = pred,
 137             _ => return None,
 138         }
 139     }
 140
 141     None
 142 }
 143
 144 // If this statement is an assignment of an unprojected local to the return place,
 145 // return that local.
 146 fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
 147     if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
 148         if lhs.as_local() == Some(mir::RETURN_PLACE) {
 149             if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
 150                 return rhs.as_local();
 151             }
 152         }
 153     }
 154
 155     None
 156 }
 157
 158 struct RenameToReturnPlace<'tcx> {
 159     to_rename: Local,
 160     tcx: TyCtxt<'tcx>,
 161 }
 162
 163 /// Replaces all uses of `self.to_rename` with `_0`.
 164 impl MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
 165     fn tcx(&self) -> TyCtxt<'tcx> {
 166         self.tcx
 167     }
 168
 169     fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
 170         // Remove assignments of the local being replaced to the return place, since it is now the
 171         // return place:
 172         //     _0 = _1
 173         if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
 174             stmt.kind = mir::StatementKind::Nop;
 175             return;
 176         }
 177
 178         // Remove storage annotations for the local being replaced:
 179         //     StorageLive(_1)
 180         if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
 181             stmt.kind
 182         {
 183             if local == self.to_rename {
 184                 stmt.kind = mir::StatementKind::Nop;
 185                 return;
 186             }
 187         }
 188
 189         self.super_statement(stmt, loc)
 190     }
 191
 192     fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
 193         // Ignore the implicit "use" of the return place in a `Return` statement.
 194         if let mir::TerminatorKind::Return = terminator.kind {
 195             return;
 196         }
 197
 198         self.super_terminator(terminator, loc);
 199     }
 200
 201     fn visit_local(&mut self, l: &mut Local, ctxt: PlaceContext, _: Location) {
 202         if *l == mir::RETURN_PLACE {
 203             assert_eq!(ctxt, PlaceContext::NonUse(NonUseContext::VarDebugInfo));
 204         } else if *l == self.to_rename {
 205             *l = mir::RETURN_PLACE;
 206         }
 207     }
 208 }
 209
 210 struct IsReturnPlaceRead(bool);
 211
 212 impl IsReturnPlaceRead {
 213     fn run(body: &mir::Body<'_>) -> bool {
 214         let mut vis = IsReturnPlaceRead(false);
 215         vis.visit_body(body);
 216         vis.0
 217     }
 218 }
 219
 220 impl Visitor<'tcx> for IsReturnPlaceRead {
 221     fn visit_local(&mut self, &l: &Local, ctxt: PlaceContext, _: Location) {
 222         if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
 223             self.0 = true;
 224         }
 225     }
 226
 227     fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
 228         // Ignore the implicit "use" of the return place in a `Return` statement.
 229         if let mir::TerminatorKind::Return = terminator.kind {
 230             return;
 231         }
 232
 233         self.super_terminator(terminator, loc);
 234     }
 235 }