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