rustc_mir_transform/
nrvo.rs

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

rustc_mir_transform/
nrvo.rs
