]>
Commit | Line | Data |
---|---|---|
29967ef6 XL |
1 | //! See the docs for [`RenameReturnPlace`]. |
2 | ||
f9f354fc XL |
3 | use rustc_hir::Mutability; |
4 | use rustc_index::bit_set::HybridBitSet; | |
1b1a35ee | 5 | use rustc_middle::mir::visit::{MutVisitor, NonUseContext, PlaceContext, Visitor}; |
f9f354fc XL |
6 | use rustc_middle::mir::{self, BasicBlock, Local, Location}; |
7 | use rustc_middle::ty::TyCtxt; | |
8 | ||
29967ef6 | 9 | use crate::transform::MirPass; |
f9f354fc XL |
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 { | |
29967ef6 | 36 | fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut mir::Body<'tcx>) { |
f9f354fc XL |
37 | if tcx.sess.opts.debugging_opts.mir_opt_level == 0 { |
38 | return; | |
39 | } | |
40 | ||
fc512014 | 41 | let def_id = body.source.def_id(); |
f9f354fc XL |
42 | let returned_local = match local_eligible_for_nrvo(body) { |
43 | Some(l) => l, | |
44 | None => { | |
fc512014 | 45 | debug!("`{:?}` was ineligible for NRVO", def_id); |
f9f354fc XL |
46 | return; |
47 | } | |
48 | }; | |
49 | ||
fc512014 XL |
50 | if !tcx.consider_optimizing(|| format!("RenameReturnPlace {:?}", def_id)) { |
51 | return; | |
52 | } | |
53 | ||
f9f354fc XL |
54 | debug!( |
55 | "`{:?}` was eligible for NRVO, making {:?} the return place", | |
fc512014 | 56 | def_id, returned_local |
f9f354fc XL |
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 | ||
f035d41b | 120 | copied_to_return_place |
f9f354fc XL |
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 | ||
f035d41b | 145 | None |
f9f354fc XL |
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 | ||
1b1a35ee XL |
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 { | |
f9f354fc XL |
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 | } |