compiler/rustc_mir_transform/src/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::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 is_enabled(&self, sess: &rustc_session::Session) -> bool {
  37         sess.mir_opt_level() > 0
  38     }
  39
  40     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) {
  41         let def_id = body.source.def_id();
  42         let Some(returned_local) = local_eligible_for_nrvo(body) else {
  43             debug!("`{:?}` was ineligible for NRVO", def_id);
  44             return;
  45         };
  46
  47         if !tcx.consider_optimizing(|| format!("RenameReturnPlace {:?}", def_id)) {
  48             return;
  49         }
  50
  51         debug!(
  52             "`{:?}` was eligible for NRVO, making {:?} the return place",
  53             def_id, returned_local
  54         );
  55
  56         RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body_preserves_cfg(body);
  57
  58         // Clean up the `NOP`s we inserted for statements made useless by our renaming.
  59         for block_data in body.basic_blocks.as_mut_preserves_cfg() {
  60             block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
  61         }
  62
  63         // Overwrite the debuginfo of `_0` with that of the renamed local.
  64         let (renamed_decl, ret_decl) =
  65             body.local_decls.pick2_mut(returned_local, mir::RETURN_PLACE);
  66
  67         // Sometimes, the return place is assigned a local of a different but coercible type, for
  68         // example `&mut T` instead of `&T`. Overwriting the `LocalInfo` for the return place means
  69         // its type may no longer match the return type of its function. This doesn't cause a
  70         // problem in codegen because these two types are layout-compatible, but may be unexpected.
  71         debug!("_0: {:?} = {:?}: {:?}", ret_decl.ty, returned_local, renamed_decl.ty);
  72         ret_decl.clone_from(renamed_decl);
  73
  74         // The return place is always mutable.
  75         ret_decl.mutability = Mutability::Mut;
  76     }
  77 }
  78
  79 /// MIR that is eligible for the NRVO must fulfill two conditions:
  80 ///   1. The return place must not be read prior to the `Return` terminator.
  81 ///   2. A simple assignment of a whole local to the return place (e.g., `_0 = _1`) must be the
  82 ///      only definition of the return place reaching the `Return` terminator.
  83 ///
  84 /// If the MIR fulfills both these conditions, this function returns the `Local` that is assigned
  85 /// to the return place along all possible paths through the control-flow graph.
  86 fn local_eligible_for_nrvo(body: &mut mir::Body<'_>) -> Option<Local> {
  87     if IsReturnPlaceRead::run(body) {
  88         return None;
  89     }
  90
  91     let mut copied_to_return_place = None;
  92     for block in body.basic_blocks.indices() {
  93         // Look for blocks with a `Return` terminator.
  94         if !matches!(body[block].terminator().kind, mir::TerminatorKind::Return) {
  95             continue;
  96         }
  97
  98         // Look for an assignment of a single local to the return place prior to the `Return`.
  99         let returned_local = find_local_assigned_to_return_place(block, body)?;
 100         match body.local_kind(returned_local) {
 101             // FIXME: Can we do this for arguments as well?
 102             mir::LocalKind::Arg => return None,
 103
 104             mir::LocalKind::ReturnPointer => bug!("Return place was assigned to itself?"),
 105             mir::LocalKind::Var | mir::LocalKind::Temp => {}
 106         }
 107
 108         // If multiple different locals are copied to the return place. We can't pick a
 109         // single one to rename.
 110         if copied_to_return_place.map_or(false, |old| old != returned_local) {
 111             return None;
 112         }
 113
 114         copied_to_return_place = Some(returned_local);
 115     }
 116
 117     copied_to_return_place
 118 }
 119
 120 fn find_local_assigned_to_return_place(
 121     start: BasicBlock,
 122     body: &mut mir::Body<'_>,
 123 ) -> Option<Local> {
 124     let mut block = start;
 125     let mut seen = HybridBitSet::new_empty(body.basic_blocks.len());
 126
 127     // Iterate as long as `block` has exactly one predecessor that we have not yet visited.
 128     while seen.insert(block) {
 129         trace!("Looking for assignments to `_0` in {:?}", block);
 130
 131         let local = body[block].statements.iter().rev().find_map(as_local_assigned_to_return_place);
 132         if local.is_some() {
 133             return local;
 134         }
 135
 136         match body.basic_blocks.predecessors()[block].as_slice() {
 137             &[pred] => block = pred,
 138             _ => return None,
 139         }
 140     }
 141
 142     None
 143 }
 144
 145 // If this statement is an assignment of an unprojected local to the return place,
 146 // return that local.
 147 fn as_local_assigned_to_return_place(stmt: &mir::Statement<'_>) -> Option<Local> {
 148     if let mir::StatementKind::Assign(box (lhs, rhs)) = &stmt.kind {
 149         if lhs.as_local() == Some(mir::RETURN_PLACE) {
 150             if let mir::Rvalue::Use(mir::Operand::Copy(rhs) | mir::Operand::Move(rhs)) = rhs {
 151                 return rhs.as_local();
 152             }
 153         }
 154     }
 155
 156     None
 157 }
 158
 159 struct RenameToReturnPlace<'tcx> {
 160     to_rename: Local,
 161     tcx: TyCtxt<'tcx>,
 162 }
 163
 164 /// Replaces all uses of `self.to_rename` with `_0`.
 165 impl<'tcx> MutVisitor<'tcx> for RenameToReturnPlace<'tcx> {
 166     fn tcx(&self) -> TyCtxt<'tcx> {
 167         self.tcx
 168     }
 169
 170     fn visit_statement(&mut self, stmt: &mut mir::Statement<'tcx>, loc: Location) {
 171         // Remove assignments of the local being replaced to the return place, since it is now the
 172         // return place:
 173         //     _0 = _1
 174         if as_local_assigned_to_return_place(stmt) == Some(self.to_rename) {
 175             stmt.kind = mir::StatementKind::Nop;
 176             return;
 177         }
 178
 179         // Remove storage annotations for the local being replaced:
 180         //     StorageLive(_1)
 181         if let mir::StatementKind::StorageLive(local) | mir::StatementKind::StorageDead(local) =
 182             stmt.kind
 183         {
 184             if local == self.to_rename {
 185                 stmt.kind = mir::StatementKind::Nop;
 186                 return;
 187             }
 188         }
 189
 190         self.super_statement(stmt, loc)
 191     }
 192
 193     fn visit_terminator(&mut self, terminator: &mut mir::Terminator<'tcx>, loc: Location) {
 194         // Ignore the implicit "use" of the return place in a `Return` statement.
 195         if let mir::TerminatorKind::Return = terminator.kind {
 196             return;
 197         }
 198
 199         self.super_terminator(terminator, loc);
 200     }
 201
 202     fn visit_local(&mut self, l: &mut Local, ctxt: PlaceContext, _: Location) {
 203         if *l == mir::RETURN_PLACE {
 204             assert_eq!(ctxt, PlaceContext::NonUse(NonUseContext::VarDebugInfo));
 205         } else if *l == self.to_rename {
 206             *l = mir::RETURN_PLACE;
 207         }
 208     }
 209 }
 210
 211 struct IsReturnPlaceRead(bool);
 212
 213 impl IsReturnPlaceRead {
 214     fn run(body: &mir::Body<'_>) -> bool {
 215         let mut vis = IsReturnPlaceRead(false);
 216         vis.visit_body(body);
 217         vis.0
 218     }
 219 }
 220
 221 impl<'tcx> Visitor<'tcx> for IsReturnPlaceRead {
 222     fn visit_local(&mut self, l: Local, ctxt: PlaceContext, _: Location) {
 223         if l == mir::RETURN_PLACE && ctxt.is_use() && !ctxt.is_place_assignment() {
 224             self.0 = true;
 225         }
 226     }
 227
 228     fn visit_terminator(&mut self, terminator: &mir::Terminator<'tcx>, loc: Location) {
 229         // Ignore the implicit "use" of the return place in a `Return` statement.
 230         if let mir::TerminatorKind::Return = terminator.kind {
 231             return;
 232         }
 233
 234         self.super_terminator(terminator, loc);
 235     }
 236 }