]>
Commit | Line | Data |
---|---|---|
ea8adc8c XL |
1 | //! This is the implementation of the pass which transforms generators into state machines. |
2 | //! | |
3 | //! MIR generation for generators creates a function which has a self argument which | |
4 | //! passes by value. This argument is effectively a generator type which only contains upvars and | |
5 | //! is only used for this argument inside the MIR for the generator. | |
6 | //! It is passed by value to enable upvars to be moved out of it. Drop elaboration runs on that | |
7 | //! MIR before this pass and creates drop flags for MIR locals. | |
8 | //! It will also drop the generator argument (which only consists of upvars) if any of the upvars | |
9 | //! are moved out of. This pass elaborates the drops of upvars / generator argument in the case | |
10 | //! that none of the upvars were moved out of. This is because we cannot have any drops of this | |
11 | //! generator in the MIR, since it is used to create the drop glue for the generator. We'd get | |
12 | //! infinite recursion otherwise. | |
13 | //! | |
14 | //! This pass creates the implementation for the Generator::resume function and the drop shim | |
15 | //! for the generator based on the MIR input. It converts the generator argument from Self to | |
16 | //! &mut Self adding derefs in the MIR as needed. It computes the final layout of the generator | |
17 | //! struct which looks like this: | |
18 | //! First upvars are stored | |
19 | //! It is followed by the generator state field. | |
20 | //! Then finally the MIR locals which are live across a suspension point are stored. | |
21 | //! | |
22 | //! struct Generator { | |
23 | //! upvars..., | |
24 | //! state: u32, | |
25 | //! mir_locals..., | |
26 | //! } | |
27 | //! | |
28 | //! This pass computes the meaning of the state field and the MIR locals which are live | |
532ac7d7 | 29 | //! across a suspension point. There are however three hardcoded generator states: |
ea8adc8c XL |
30 | //! 0 - Generator have not been resumed yet |
31 | //! 1 - Generator has returned / is completed | |
32 | //! 2 - Generator has been poisoned | |
33 | //! | |
34 | //! It also rewrites `return x` and `yield y` as setting a new generator state and returning | |
35 | //! GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively. | |
36 | //! MIR locals which are live across a suspension point are moved to the generator struct | |
37 | //! with references to them being updated with references to the generator struct. | |
38 | //! | |
39 | //! The pass creates two functions which have a switch on the generator state giving | |
40 | //! the action to take. | |
41 | //! | |
42 | //! One of them is the implementation of Generator::resume. | |
43 | //! For generators with state 0 (unresumed) it starts the execution of the generator. | |
44 | //! For generators with state 1 (returned) and state 2 (poisoned) it panics. | |
45 | //! Otherwise it continues the execution from the last suspension point. | |
46 | //! | |
47 | //! The other function is the drop glue for the generator. | |
48 | //! For generators with state 0 (unresumed) it drops the upvars of the generator. | |
49 | //! For generators with state 1 (returned) and state 2 (poisoned) it does nothing. | |
50 | //! Otherwise it drops all the values in scope at the last suspension point. | |
51 | ||
c295e0f8 | 52 | use crate::simplify; |
3dfed10e | 53 | use crate::util::expand_aggregate; |
c295e0f8 | 54 | use crate::MirPass; |
b7449926 | 55 | use rustc_data_structures::fx::FxHashMap; |
dfeec247 | 56 | use rustc_hir as hir; |
3dfed10e | 57 | use rustc_hir::lang_items::LangItem; |
dfeec247 | 58 | use rustc_index::bit_set::{BitMatrix, BitSet}; |
e74abb32 | 59 | use rustc_index::vec::{Idx, IndexVec}; |
c295e0f8 | 60 | use rustc_middle::mir::dump_mir; |
f035d41b | 61 | use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor}; |
ba9703b0 | 62 | use rustc_middle::mir::*; |
3dfed10e | 63 | use rustc_middle::ty::subst::{Subst, SubstsRef}; |
ba9703b0 XL |
64 | use rustc_middle::ty::GeneratorSubsts; |
65 | use rustc_middle::ty::{self, AdtDef, Ty, TyCtxt}; | |
c295e0f8 XL |
66 | use rustc_mir_dataflow::impls::{ |
67 | MaybeBorrowedLocals, MaybeLiveLocals, MaybeRequiresStorage, MaybeStorageLive, | |
68 | }; | |
69 | use rustc_mir_dataflow::storage; | |
70 | use rustc_mir_dataflow::{self, Analysis}; | |
ba9703b0 | 71 | use rustc_target::abi::VariantIdx; |
f9f354fc | 72 | use rustc_target::spec::PanicStrategy; |
f035d41b | 73 | use std::{iter, ops}; |
ea8adc8c XL |
74 | |
75 | pub struct StateTransform; | |
76 | ||
e74abb32 | 77 | struct RenameLocalVisitor<'tcx> { |
ea8adc8c XL |
78 | from: Local, |
79 | to: Local, | |
e74abb32 | 80 | tcx: TyCtxt<'tcx>, |
ea8adc8c XL |
81 | } |
82 | ||
e74abb32 XL |
83 | impl<'tcx> MutVisitor<'tcx> for RenameLocalVisitor<'tcx> { |
84 | fn tcx(&self) -> TyCtxt<'tcx> { | |
85 | self.tcx | |
86 | } | |
87 | ||
dfeec247 | 88 | fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { |
ea8adc8c XL |
89 | if *local == self.from { |
90 | *local = self.to; | |
91 | } | |
92 | } | |
f9f354fc | 93 | |
f035d41b XL |
94 | fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, location: Location) { |
95 | match terminator.kind { | |
f9f354fc XL |
96 | TerminatorKind::Return => { |
97 | // Do not replace the implicit `_0` access here, as that's not possible. The | |
98 | // transform already handles `return` correctly. | |
99 | } | |
f035d41b | 100 | _ => self.super_terminator(terminator, location), |
f9f354fc XL |
101 | } |
102 | } | |
ea8adc8c XL |
103 | } |
104 | ||
e74abb32 XL |
105 | struct DerefArgVisitor<'tcx> { |
106 | tcx: TyCtxt<'tcx>, | |
107 | } | |
108 | ||
109 | impl<'tcx> MutVisitor<'tcx> for DerefArgVisitor<'tcx> { | |
110 | fn tcx(&self) -> TyCtxt<'tcx> { | |
111 | self.tcx | |
112 | } | |
ea8adc8c | 113 | |
dfeec247 | 114 | fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { |
ba9703b0 | 115 | assert_ne!(*local, SELF_ARG); |
ea8adc8c XL |
116 | } |
117 | ||
dfeec247 | 118 | fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) { |
ba9703b0 | 119 | if place.local == SELF_ARG { |
dfeec247 XL |
120 | replace_base( |
121 | place, | |
122 | Place { | |
ba9703b0 | 123 | local: SELF_ARG, |
dfeec247 XL |
124 | projection: self.tcx().intern_place_elems(&[ProjectionElem::Deref]), |
125 | }, | |
126 | self.tcx, | |
127 | ); | |
ea8adc8c | 128 | } else { |
f9f354fc | 129 | self.visit_local(&mut place.local, context, location); |
e74abb32 XL |
130 | |
131 | for elem in place.projection.iter() { | |
132 | if let PlaceElem::Index(local) = elem { | |
f9f354fc | 133 | assert_ne!(local, SELF_ARG); |
e74abb32 XL |
134 | } |
135 | } | |
ea8adc8c XL |
136 | } |
137 | } | |
138 | } | |
139 | ||
9fa01778 XL |
140 | struct PinArgVisitor<'tcx> { |
141 | ref_gen_ty: Ty<'tcx>, | |
e74abb32 | 142 | tcx: TyCtxt<'tcx>, |
9fa01778 XL |
143 | } |
144 | ||
145 | impl<'tcx> MutVisitor<'tcx> for PinArgVisitor<'tcx> { | |
e74abb32 XL |
146 | fn tcx(&self) -> TyCtxt<'tcx> { |
147 | self.tcx | |
148 | } | |
149 | ||
dfeec247 | 150 | fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { |
ba9703b0 | 151 | assert_ne!(*local, SELF_ARG); |
9fa01778 XL |
152 | } |
153 | ||
e74abb32 | 154 | fn visit_place(&mut self, place: &mut Place<'tcx>, context: PlaceContext, location: Location) { |
ba9703b0 | 155 | if place.local == SELF_ARG { |
e74abb32 XL |
156 | replace_base( |
157 | place, | |
158 | Place { | |
ba9703b0 | 159 | local: SELF_ARG, |
dfeec247 XL |
160 | projection: self.tcx().intern_place_elems(&[ProjectionElem::Field( |
161 | Field::new(0), | |
162 | self.ref_gen_ty, | |
e74abb32 XL |
163 | )]), |
164 | }, | |
165 | self.tcx, | |
166 | ); | |
9fa01778 | 167 | } else { |
f9f354fc | 168 | self.visit_local(&mut place.local, context, location); |
e74abb32 XL |
169 | |
170 | for elem in place.projection.iter() { | |
171 | if let PlaceElem::Index(local) = elem { | |
f9f354fc | 172 | assert_ne!(local, SELF_ARG); |
e74abb32 XL |
173 | } |
174 | } | |
9fa01778 XL |
175 | } |
176 | } | |
177 | } | |
178 | ||
e74abb32 | 179 | fn replace_base<'tcx>(place: &mut Place<'tcx>, new_base: Place<'tcx>, tcx: TyCtxt<'tcx>) { |
dfeec247 | 180 | place.local = new_base.local; |
e1599b0c XL |
181 | |
182 | let mut new_projection = new_base.projection.to_vec(); | |
183 | new_projection.append(&mut place.projection.to_vec()); | |
184 | ||
e74abb32 | 185 | place.projection = tcx.intern_place_elems(&new_projection); |
dc9dc135 XL |
186 | } |
187 | ||
ba9703b0 | 188 | const SELF_ARG: Local = Local::from_u32(1); |
ea8adc8c | 189 | |
74b04a01 | 190 | /// Generator has not been resumed yet. |
48663c56 | 191 | const UNRESUMED: usize = GeneratorSubsts::UNRESUMED; |
74b04a01 | 192 | /// Generator has returned / is completed. |
48663c56 | 193 | const RETURNED: usize = GeneratorSubsts::RETURNED; |
74b04a01 | 194 | /// Generator has panicked and is poisoned. |
48663c56 | 195 | const POISONED: usize = GeneratorSubsts::POISONED; |
532ac7d7 | 196 | |
74b04a01 XL |
197 | /// A `yield` point in the generator. |
198 | struct SuspensionPoint<'tcx> { | |
199 | /// State discriminant used when suspending or resuming at this point. | |
48663c56 | 200 | state: usize, |
74b04a01 | 201 | /// The block to jump to after resumption. |
ea8adc8c | 202 | resume: BasicBlock, |
74b04a01 XL |
203 | /// Where to move the resume argument after resumption. |
204 | resume_arg: Place<'tcx>, | |
205 | /// Which block to jump to if the generator is dropped in this state. | |
ea8adc8c | 206 | drop: Option<BasicBlock>, |
74b04a01 | 207 | /// Set of locals that have live storage while at this suspension point. |
f9f354fc | 208 | storage_liveness: BitSet<Local>, |
ea8adc8c XL |
209 | } |
210 | ||
dc9dc135 XL |
211 | struct TransformVisitor<'tcx> { |
212 | tcx: TyCtxt<'tcx>, | |
5e7ed085 | 213 | state_adt_ref: AdtDef<'tcx>, |
532ac7d7 | 214 | state_substs: SubstsRef<'tcx>, |
ea8adc8c | 215 | |
48663c56 XL |
216 | // The type of the discriminant in the generator struct |
217 | discr_ty: Ty<'tcx>, | |
ea8adc8c XL |
218 | |
219 | // Mapping from Local to (type of local, generator struct index) | |
b7449926 | 220 | // FIXME(eddyb) This should use `IndexVec<Local, Option<_>>`. |
48663c56 | 221 | remap: FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>, |
ea8adc8c XL |
222 | |
223 | // A map from a suspension point in a block to the locals which have live storage at that point | |
f9f354fc | 224 | storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>, |
ea8adc8c XL |
225 | |
226 | // A list of suspension points, generated during the transform | |
74b04a01 | 227 | suspension_points: Vec<SuspensionPoint<'tcx>>, |
ea8adc8c | 228 | |
ba9703b0 XL |
229 | // The set of locals that have no `StorageLive`/`StorageDead` annotations. |
230 | always_live_locals: storage::AlwaysLiveLocals, | |
231 | ||
ff7c6d11 | 232 | // The original RETURN_PLACE local |
ea8adc8c XL |
233 | new_ret_local: Local, |
234 | } | |
235 | ||
a2a8927a | 236 | impl<'tcx> TransformVisitor<'tcx> { |
3dfed10e XL |
237 | // Make a GeneratorState variant assignment. `core::ops::GeneratorState` only has single |
238 | // element tuple variants, so we can just write to the downcasted first field and then set the | |
239 | // discriminant to the appropriate variant. | |
240 | fn make_state( | |
241 | &self, | |
242 | idx: VariantIdx, | |
243 | val: Operand<'tcx>, | |
244 | source_info: SourceInfo, | |
245 | ) -> impl Iterator<Item = Statement<'tcx>> { | |
5e7ed085 FG |
246 | let kind = AggregateKind::Adt(self.state_adt_ref.did(), idx, self.state_substs, None, None); |
247 | assert_eq!(self.state_adt_ref.variant(idx).fields.len(), 1); | |
3dfed10e XL |
248 | let ty = self |
249 | .tcx | |
5e7ed085 | 250 | .type_of(self.state_adt_ref.variant(idx).fields[0].did) |
3dfed10e XL |
251 | .subst(self.tcx, self.state_substs); |
252 | expand_aggregate( | |
253 | Place::return_place(), | |
254 | std::iter::once((val, ty)), | |
255 | kind, | |
256 | source_info, | |
257 | self.tcx, | |
258 | ) | |
ea8adc8c XL |
259 | } |
260 | ||
ff7c6d11 | 261 | // Create a Place referencing a generator struct field |
48663c56 | 262 | fn make_field(&self, variant_index: VariantIdx, idx: usize, ty: Ty<'tcx>) -> Place<'tcx> { |
ba9703b0 | 263 | let self_place = Place::from(SELF_ARG); |
e74abb32 | 264 | let base = self.tcx.mk_place_downcast_unnamed(self_place, variant_index); |
e1599b0c XL |
265 | let mut projection = base.projection.to_vec(); |
266 | projection.push(ProjectionElem::Field(Field::new(idx), ty)); | |
267 | ||
dfeec247 | 268 | Place { local: base.local, projection: self.tcx.intern_place_elems(&projection) } |
ea8adc8c XL |
269 | } |
270 | ||
48663c56 XL |
271 | // Create a statement which changes the discriminant |
272 | fn set_discr(&self, state_disc: VariantIdx, source_info: SourceInfo) -> Statement<'tcx> { | |
ba9703b0 | 273 | let self_place = Place::from(SELF_ARG); |
ea8adc8c XL |
274 | Statement { |
275 | source_info, | |
e1599b0c | 276 | kind: StatementKind::SetDiscriminant { |
94222f64 | 277 | place: Box::new(self_place), |
e1599b0c XL |
278 | variant_index: state_disc, |
279 | }, | |
ea8adc8c XL |
280 | } |
281 | } | |
48663c56 XL |
282 | |
283 | // Create a statement which reads the discriminant into a temporary | |
dc9dc135 | 284 | fn get_discr(&self, body: &mut Body<'tcx>) -> (Statement<'tcx>, Place<'tcx>) { |
f9f354fc | 285 | let temp_decl = LocalDecl::new(self.discr_ty, body.span).internal(); |
dc9dc135 XL |
286 | let local_decls_len = body.local_decls.push(temp_decl); |
287 | let temp = Place::from(local_decls_len); | |
48663c56 | 288 | |
ba9703b0 | 289 | let self_place = Place::from(SELF_ARG); |
48663c56 | 290 | let assign = Statement { |
f9f354fc | 291 | source_info: SourceInfo::outermost(body.span), |
94222f64 | 292 | kind: StatementKind::Assign(Box::new((temp, Rvalue::Discriminant(self_place)))), |
48663c56 XL |
293 | }; |
294 | (assign, temp) | |
295 | } | |
ea8adc8c XL |
296 | } |
297 | ||
a2a8927a | 298 | impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> { |
e74abb32 XL |
299 | fn tcx(&self) -> TyCtxt<'tcx> { |
300 | self.tcx | |
301 | } | |
302 | ||
dfeec247 | 303 | fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { |
ea8adc8c XL |
304 | assert_eq!(self.remap.get(local), None); |
305 | } | |
306 | ||
e74abb32 XL |
307 | fn visit_place( |
308 | &mut self, | |
309 | place: &mut Place<'tcx>, | |
dfeec247 XL |
310 | _context: PlaceContext, |
311 | _location: Location, | |
e74abb32 | 312 | ) { |
dfeec247 XL |
313 | // Replace an Local in the remap with a generator struct access |
314 | if let Some(&(ty, variant_index, idx)) = self.remap.get(&place.local) { | |
315 | replace_base(place, self.make_field(variant_index, idx, ty), self.tcx); | |
ea8adc8c XL |
316 | } |
317 | } | |
318 | ||
dfeec247 | 319 | fn visit_basic_block_data(&mut self, block: BasicBlock, data: &mut BasicBlockData<'tcx>) { |
ea8adc8c | 320 | // Remove StorageLive and StorageDead statements for remapped locals |
dfeec247 XL |
321 | data.retain_statements(|s| match s.kind { |
322 | StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => { | |
323 | !self.remap.contains_key(&l) | |
ea8adc8c | 324 | } |
dfeec247 | 325 | _ => true, |
ea8adc8c XL |
326 | }); |
327 | ||
328 | let ret_val = match data.terminator().kind { | |
dfeec247 XL |
329 | TerminatorKind::Return => Some(( |
330 | VariantIdx::new(1), | |
ea8adc8c | 331 | None, |
dc9dc135 | 332 | Operand::Move(Place::from(self.new_ret_local)), |
dfeec247 XL |
333 | None, |
334 | )), | |
74b04a01 XL |
335 | TerminatorKind::Yield { ref value, resume, resume_arg, drop } => { |
336 | Some((VariantIdx::new(0), Some((resume, resume_arg)), value.clone(), drop)) | |
dfeec247 XL |
337 | } |
338 | _ => None, | |
ea8adc8c XL |
339 | }; |
340 | ||
341 | if let Some((state_idx, resume, v, drop)) = ret_val { | |
342 | let source_info = data.terminator().source_info; | |
343 | // We must assign the value first in case it gets declared dead below | |
3dfed10e | 344 | data.statements.extend(self.make_state(state_idx, v, source_info)); |
c295e0f8 | 345 | let state = if let Some((resume, mut resume_arg)) = resume { |
dfeec247 | 346 | // Yield |
48663c56 | 347 | let state = 3 + self.suspension_points.len(); |
ea8adc8c | 348 | |
74b04a01 XL |
349 | // The resume arg target location might itself be remapped if its base local is |
350 | // live across a yield. | |
351 | let resume_arg = | |
352 | if let Some(&(ty, variant, idx)) = self.remap.get(&resume_arg.local) { | |
c295e0f8 XL |
353 | replace_base(&mut resume_arg, self.make_field(variant, idx, ty), self.tcx); |
354 | resume_arg | |
74b04a01 XL |
355 | } else { |
356 | resume_arg | |
357 | }; | |
358 | ||
ea8adc8c XL |
359 | self.suspension_points.push(SuspensionPoint { |
360 | state, | |
361 | resume, | |
74b04a01 | 362 | resume_arg, |
ea8adc8c | 363 | drop, |
f9f354fc | 364 | storage_liveness: self.storage_liveness[block].clone().unwrap(), |
ea8adc8c XL |
365 | }); |
366 | ||
48663c56 | 367 | VariantIdx::new(state) |
dfeec247 XL |
368 | } else { |
369 | // Return | |
48663c56 | 370 | VariantIdx::new(RETURNED) // state for returned |
ea8adc8c | 371 | }; |
48663c56 | 372 | data.statements.push(self.set_discr(state, source_info)); |
60c5eb7d | 373 | data.terminator_mut().kind = TerminatorKind::Return; |
ea8adc8c XL |
374 | } |
375 | ||
376 | self.super_basic_block_data(block, data); | |
377 | } | |
378 | } | |
379 | ||
f9f354fc | 380 | fn make_generator_state_argument_indirect<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
dc9dc135 | 381 | let gen_ty = body.local_decls.raw[1].ty; |
ea8adc8c | 382 | |
ba9703b0 XL |
383 | let ref_gen_ty = |
384 | tcx.mk_ref(tcx.lifetimes.re_erased, ty::TypeAndMut { ty: gen_ty, mutbl: Mutability::Mut }); | |
ea8adc8c XL |
385 | |
386 | // Replace the by value generator argument | |
dc9dc135 | 387 | body.local_decls.raw[1].ty = ref_gen_ty; |
ea8adc8c XL |
388 | |
389 | // Add a deref to accesses of the generator state | |
e74abb32 | 390 | DerefArgVisitor { tcx }.visit_body(body); |
ea8adc8c XL |
391 | } |
392 | ||
f9f354fc | 393 | fn make_generator_state_argument_pinned<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
dc9dc135 | 394 | let ref_gen_ty = body.local_decls.raw[1].ty; |
9fa01778 | 395 | |
3dfed10e | 396 | let pin_did = tcx.require_lang_item(LangItem::Pin, Some(body.span)); |
9fa01778 XL |
397 | let pin_adt_ref = tcx.adt_def(pin_did); |
398 | let substs = tcx.intern_substs(&[ref_gen_ty.into()]); | |
399 | let pin_ref_gen_ty = tcx.mk_adt(pin_adt_ref, substs); | |
400 | ||
401 | // Replace the by ref generator argument | |
dc9dc135 | 402 | body.local_decls.raw[1].ty = pin_ref_gen_ty; |
9fa01778 XL |
403 | |
404 | // Add the Pin field access to accesses of the generator state | |
e74abb32 | 405 | PinArgVisitor { ref_gen_ty, tcx }.visit_body(body); |
9fa01778 XL |
406 | } |
407 | ||
74b04a01 XL |
408 | /// Allocates a new local and replaces all references of `local` with it. Returns the new local. |
409 | /// | |
410 | /// `local` will be changed to a new local decl with type `ty`. | |
411 | /// | |
412 | /// Note that the new local will be uninitialized. It is the caller's responsibility to assign some | |
413 | /// valid value to it before its first use. | |
414 | fn replace_local<'tcx>( | |
415 | local: Local, | |
416 | ty: Ty<'tcx>, | |
f9f354fc | 417 | body: &mut Body<'tcx>, |
e74abb32 | 418 | tcx: TyCtxt<'tcx>, |
8faf50e0 | 419 | ) -> Local { |
f9f354fc XL |
420 | let new_decl = LocalDecl::new(ty, body.span); |
421 | let new_local = body.local_decls.push(new_decl); | |
74b04a01 | 422 | body.local_decls.swap(local, new_local); |
ea8adc8c | 423 | |
74b04a01 | 424 | RenameLocalVisitor { from: local, to: new_local, tcx }.visit_body(body); |
ea8adc8c | 425 | |
74b04a01 | 426 | new_local |
ea8adc8c XL |
427 | } |
428 | ||
dc9dc135 XL |
429 | struct LivenessInfo { |
430 | /// Which locals are live across any suspension point. | |
f035d41b | 431 | saved_locals: GeneratorSavedLocals, |
dc9dc135 XL |
432 | |
433 | /// The set of saved locals live at each suspension point. | |
434 | live_locals_at_suspension_points: Vec<BitSet<GeneratorSavedLocal>>, | |
435 | ||
f035d41b XL |
436 | /// Parallel vec to the above with SourceInfo for each yield terminator. |
437 | source_info_at_suspension_points: Vec<SourceInfo>, | |
438 | ||
dc9dc135 XL |
439 | /// For every saved local, the set of other saved locals that are |
440 | /// storage-live at the same time as this local. We cannot overlap locals in | |
441 | /// the layout which have conflicting storage. | |
442 | storage_conflicts: BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>, | |
443 | ||
444 | /// For every suspending block, the locals which are storage-live across | |
445 | /// that suspension point. | |
f9f354fc | 446 | storage_liveness: IndexVec<BasicBlock, Option<BitSet<Local>>>, |
dc9dc135 XL |
447 | } |
448 | ||
a2a8927a | 449 | fn locals_live_across_suspend_points<'tcx>( |
dc9dc135 | 450 | tcx: TyCtxt<'tcx>, |
f9f354fc | 451 | body: &Body<'tcx>, |
ba9703b0 | 452 | always_live_locals: &storage::AlwaysLiveLocals, |
b7449926 | 453 | movable: bool, |
dc9dc135 | 454 | ) -> LivenessInfo { |
60c5eb7d | 455 | let body_ref: &Body<'_> = &body; |
2c00a5a8 XL |
456 | |
457 | // Calculate when MIR locals have live storage. This gives us an upper bound of their | |
458 | // lifetimes. | |
ba9703b0 | 459 | let mut storage_live = MaybeStorageLive::new(always_live_locals.clone()) |
29967ef6 | 460 | .into_engine(tcx, body_ref) |
74b04a01 XL |
461 | .iterate_to_fixpoint() |
462 | .into_results_cursor(body_ref); | |
ea8adc8c | 463 | |
2c00a5a8 XL |
464 | // Calculate the MIR locals which have been previously |
465 | // borrowed (even if they are still active). | |
5e7ed085 FG |
466 | let borrowed_locals_results = |
467 | MaybeBorrowedLocals.into_engine(tcx, body_ref).pass_name("generator").iterate_to_fixpoint(); | |
74b04a01 XL |
468 | |
469 | let mut borrowed_locals_cursor = | |
c295e0f8 | 470 | rustc_mir_dataflow::ResultsCursor::new(body_ref, &borrowed_locals_results); |
dc9dc135 XL |
471 | |
472 | // Calculate the MIR locals that we actually need to keep storage around | |
473 | // for. | |
74b04a01 | 474 | let requires_storage_results = MaybeRequiresStorage::new(body, &borrowed_locals_results) |
29967ef6 | 475 | .into_engine(tcx, body_ref) |
74b04a01 | 476 | .iterate_to_fixpoint(); |
dfeec247 | 477 | let mut requires_storage_cursor = |
c295e0f8 | 478 | rustc_mir_dataflow::ResultsCursor::new(body_ref, &requires_storage_results); |
dfeec247 XL |
479 | |
480 | // Calculate the liveness of MIR locals ignoring borrows. | |
f9f354fc | 481 | let mut liveness = MaybeLiveLocals |
29967ef6 | 482 | .into_engine(tcx, body_ref) |
1b1a35ee | 483 | .pass_name("generator") |
f9f354fc XL |
484 | .iterate_to_fixpoint() |
485 | .into_results_cursor(body_ref); | |
ea8adc8c | 486 | |
f9f354fc | 487 | let mut storage_liveness_map = IndexVec::from_elem(None, body.basic_blocks()); |
dc9dc135 | 488 | let mut live_locals_at_suspension_points = Vec::new(); |
f035d41b | 489 | let mut source_info_at_suspension_points = Vec::new(); |
f9f354fc | 490 | let mut live_locals_at_any_suspension_point = BitSet::new_empty(body.local_decls.len()); |
ea8adc8c | 491 | |
dc9dc135 | 492 | for (block, data) in body.basic_blocks().iter_enumerated() { |
ea8adc8c | 493 | if let TerminatorKind::Yield { .. } = data.terminator().kind { |
74b04a01 | 494 | let loc = Location { block, statement_index: data.statements.len() }; |
ea8adc8c | 495 | |
f9f354fc XL |
496 | liveness.seek_to_block_end(block); |
497 | let mut live_locals = liveness.get().clone(); | |
498 | ||
dc9dc135 | 499 | if !movable { |
2c00a5a8 XL |
500 | // The `liveness` variable contains the liveness of MIR locals ignoring borrows. |
501 | // This is correct for movable generators since borrows cannot live across | |
502 | // suspension points. However for immovable generators we need to account for | |
5e7ed085 | 503 | // borrows, so we conservatively assume that all borrowed locals are live until |
b7449926 | 504 | // we find a StorageDead statement referencing the locals. |
2c00a5a8 XL |
505 | // To do this we just union our `liveness` result with `borrowed_locals`, which |
506 | // contains all the locals which has been borrowed before this suspension point. | |
507 | // If a borrow is converted to a raw reference, we must also assume that it lives | |
508 | // forever. Note that the final liveness is still bounded by the storage liveness | |
509 | // of the local, which happens using the `intersect` operation below. | |
f9f354fc XL |
510 | borrowed_locals_cursor.seek_before_primary_effect(loc); |
511 | live_locals.union(borrowed_locals_cursor.get()); | |
2c00a5a8 XL |
512 | } |
513 | ||
2c00a5a8 XL |
514 | // Store the storage liveness for later use so we can restore the state |
515 | // after a suspension point | |
f9f354fc XL |
516 | storage_live.seek_before_primary_effect(loc); |
517 | storage_liveness_map[block] = Some(storage_live.get().clone()); | |
2c00a5a8 XL |
518 | |
519 | // Locals live are live at this point only if they are used across | |
520 | // suspension points (the `liveness` variable) | |
dc9dc135 | 521 | // and their storage is required (the `storage_required` variable) |
f9f354fc XL |
522 | requires_storage_cursor.seek_before_primary_effect(loc); |
523 | live_locals.intersect(requires_storage_cursor.get()); | |
dc9dc135 | 524 | |
ba9703b0 | 525 | // The generator argument is ignored. |
f9f354fc | 526 | live_locals.remove(SELF_ARG); |
ea8adc8c | 527 | |
f9f354fc | 528 | debug!("loc = {:?}, live_locals = {:?}", loc, live_locals); |
ea8adc8c | 529 | |
dc9dc135 | 530 | // Add the locals live at this suspension point to the set of locals which live across |
ea8adc8c | 531 | // any suspension points |
f9f354fc | 532 | live_locals_at_any_suspension_point.union(&live_locals); |
dc9dc135 | 533 | |
f9f354fc | 534 | live_locals_at_suspension_points.push(live_locals); |
f035d41b | 535 | source_info_at_suspension_points.push(data.terminator().source_info); |
ea8adc8c XL |
536 | } |
537 | } | |
f035d41b | 538 | |
f9f354fc | 539 | debug!("live_locals_anywhere = {:?}", live_locals_at_any_suspension_point); |
f035d41b | 540 | let saved_locals = GeneratorSavedLocals(live_locals_at_any_suspension_point); |
dc9dc135 XL |
541 | |
542 | // Renumber our liveness_map bitsets to include only the locals we are | |
543 | // saving. | |
544 | let live_locals_at_suspension_points = live_locals_at_suspension_points | |
545 | .iter() | |
f035d41b | 546 | .map(|live_here| saved_locals.renumber_bitset(&live_here)) |
dc9dc135 XL |
547 | .collect(); |
548 | ||
ba9703b0 XL |
549 | let storage_conflicts = compute_storage_conflicts( |
550 | body_ref, | |
f035d41b | 551 | &saved_locals, |
ba9703b0 XL |
552 | always_live_locals.clone(), |
553 | requires_storage_results, | |
554 | ); | |
dc9dc135 XL |
555 | |
556 | LivenessInfo { | |
f035d41b | 557 | saved_locals, |
dc9dc135 | 558 | live_locals_at_suspension_points, |
f035d41b | 559 | source_info_at_suspension_points, |
dc9dc135 XL |
560 | storage_conflicts, |
561 | storage_liveness: storage_liveness_map, | |
562 | } | |
563 | } | |
564 | ||
f035d41b | 565 | /// The set of `Local`s that must be saved across yield points. |
dc9dc135 | 566 | /// |
f035d41b XL |
567 | /// `GeneratorSavedLocal` is indexed in terms of the elements in this set; |
568 | /// i.e. `GeneratorSavedLocal::new(1)` corresponds to the second local | |
569 | /// included in this set. | |
570 | struct GeneratorSavedLocals(BitSet<Local>); | |
571 | ||
572 | impl GeneratorSavedLocals { | |
573 | /// Returns an iterator over each `GeneratorSavedLocal` along with the `Local` it corresponds | |
574 | /// to. | |
575 | fn iter_enumerated(&self) -> impl '_ + Iterator<Item = (GeneratorSavedLocal, Local)> { | |
576 | self.iter().enumerate().map(|(i, l)| (GeneratorSavedLocal::from(i), l)) | |
577 | } | |
578 | ||
579 | /// Transforms a `BitSet<Local>` that contains only locals saved across yield points to the | |
580 | /// equivalent `BitSet<GeneratorSavedLocal>`. | |
581 | fn renumber_bitset(&self, input: &BitSet<Local>) -> BitSet<GeneratorSavedLocal> { | |
582 | assert!(self.superset(&input), "{:?} not a superset of {:?}", self.0, input); | |
583 | let mut out = BitSet::new_empty(self.count()); | |
584 | for (saved_local, local) in self.iter_enumerated() { | |
585 | if input.contains(local) { | |
586 | out.insert(saved_local); | |
587 | } | |
dc9dc135 | 588 | } |
f035d41b XL |
589 | out |
590 | } | |
591 | ||
592 | fn get(&self, local: Local) -> Option<GeneratorSavedLocal> { | |
593 | if !self.contains(local) { | |
594 | return None; | |
595 | } | |
596 | ||
597 | let idx = self.iter().take_while(|&l| l < local).count(); | |
598 | Some(GeneratorSavedLocal::new(idx)) | |
599 | } | |
600 | } | |
601 | ||
602 | impl ops::Deref for GeneratorSavedLocals { | |
603 | type Target = BitSet<Local>; | |
604 | ||
605 | fn deref(&self) -> &Self::Target { | |
606 | &self.0 | |
dc9dc135 | 607 | } |
dc9dc135 | 608 | } |
ea8adc8c | 609 | |
dc9dc135 XL |
610 | /// For every saved local, looks for which locals are StorageLive at the same |
611 | /// time. Generates a bitset for every local of all the other locals that may be | |
612 | /// StorageLive simultaneously with that local. This is used in the layout | |
613 | /// computation; see `GeneratorLayout` for more. | |
a2a8927a | 614 | fn compute_storage_conflicts<'mir, 'tcx>( |
dc9dc135 | 615 | body: &'mir Body<'tcx>, |
f035d41b | 616 | saved_locals: &GeneratorSavedLocals, |
ba9703b0 | 617 | always_live_locals: storage::AlwaysLiveLocals, |
c295e0f8 | 618 | requires_storage: rustc_mir_dataflow::Results<'tcx, MaybeRequiresStorage<'mir, 'tcx>>, |
dc9dc135 | 619 | ) -> BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal> { |
f035d41b | 620 | assert_eq!(body.local_decls.len(), saved_locals.domain_size()); |
ba9703b0 | 621 | |
dc9dc135 | 622 | debug!("compute_storage_conflicts({:?})", body.span); |
ba9703b0 | 623 | debug!("always_live = {:?}", always_live_locals); |
dc9dc135 | 624 | |
ba9703b0 XL |
625 | // Locals that are always live or ones that need to be stored across |
626 | // suspension points are not eligible for overlap. | |
627 | let mut ineligible_locals = always_live_locals.into_inner(); | |
94222f64 | 628 | ineligible_locals.intersect(&**saved_locals); |
dc9dc135 XL |
629 | |
630 | // Compute the storage conflicts for all eligible locals. | |
631 | let mut visitor = StorageConflictVisitor { | |
632 | body, | |
f035d41b | 633 | saved_locals: &saved_locals, |
dc9dc135 XL |
634 | local_conflicts: BitMatrix::from_row_n(&ineligible_locals, body.local_decls.len()), |
635 | }; | |
74b04a01 | 636 | |
3dfed10e | 637 | requires_storage.visit_reachable_with(body, &mut visitor); |
74b04a01 | 638 | |
dc9dc135 XL |
639 | let local_conflicts = visitor.local_conflicts; |
640 | ||
641 | // Compress the matrix using only stored locals (Local -> GeneratorSavedLocal). | |
642 | // | |
643 | // NOTE: Today we store a full conflict bitset for every local. Technically | |
644 | // this is twice as many bits as we need, since the relation is symmetric. | |
645 | // However, in practice these bitsets are not usually large. The layout code | |
646 | // also needs to keep track of how many conflicts each local has, so it's | |
647 | // simpler to keep it this way for now. | |
f035d41b XL |
648 | let mut storage_conflicts = BitMatrix::new(saved_locals.count(), saved_locals.count()); |
649 | for (saved_local_a, local_a) in saved_locals.iter_enumerated() { | |
dc9dc135 XL |
650 | if ineligible_locals.contains(local_a) { |
651 | // Conflicts with everything. | |
652 | storage_conflicts.insert_all_into_row(saved_local_a); | |
653 | } else { | |
654 | // Keep overlap information only for stored locals. | |
f035d41b | 655 | for (saved_local_b, local_b) in saved_locals.iter_enumerated() { |
dc9dc135 XL |
656 | if local_conflicts.contains(local_a, local_b) { |
657 | storage_conflicts.insert(saved_local_a, saved_local_b); | |
658 | } | |
659 | } | |
660 | } | |
661 | } | |
662 | storage_conflicts | |
663 | } | |
ea8adc8c | 664 | |
74b04a01 XL |
665 | struct StorageConflictVisitor<'mir, 'tcx, 's> { |
666 | body: &'mir Body<'tcx>, | |
f035d41b | 667 | saved_locals: &'s GeneratorSavedLocals, |
dc9dc135 XL |
668 | // FIXME(tmandry): Consider using sparse bitsets here once we have good |
669 | // benchmarks for generators. | |
670 | local_conflicts: BitMatrix<Local, Local>, | |
ea8adc8c XL |
671 | } |
672 | ||
a2a8927a XL |
673 | impl<'mir, 'tcx> rustc_mir_dataflow::ResultsVisitor<'mir, 'tcx> |
674 | for StorageConflictVisitor<'mir, 'tcx, '_> | |
675 | { | |
74b04a01 | 676 | type FlowState = BitSet<Local>; |
dc9dc135 | 677 | |
f9f354fc | 678 | fn visit_statement_before_primary_effect( |
dfeec247 | 679 | &mut self, |
74b04a01 XL |
680 | state: &Self::FlowState, |
681 | _statement: &'mir Statement<'tcx>, | |
dfeec247 | 682 | loc: Location, |
dfeec247 | 683 | ) { |
74b04a01 | 684 | self.apply_state(state, loc); |
dc9dc135 XL |
685 | } |
686 | ||
f9f354fc | 687 | fn visit_terminator_before_primary_effect( |
dfeec247 | 688 | &mut self, |
74b04a01 XL |
689 | state: &Self::FlowState, |
690 | _terminator: &'mir Terminator<'tcx>, | |
dfeec247 | 691 | loc: Location, |
dfeec247 | 692 | ) { |
74b04a01 | 693 | self.apply_state(state, loc); |
dc9dc135 XL |
694 | } |
695 | } | |
696 | ||
a2a8927a | 697 | impl StorageConflictVisitor<'_, '_, '_> { |
74b04a01 | 698 | fn apply_state(&mut self, flow_state: &BitSet<Local>, loc: Location) { |
dc9dc135 | 699 | // Ignore unreachable blocks. |
ba9703b0 XL |
700 | if self.body.basic_blocks()[loc.block].terminator().kind == TerminatorKind::Unreachable { |
701 | return; | |
702 | } | |
dc9dc135 | 703 | |
74b04a01 | 704 | let mut eligible_storage_live = flow_state.clone(); |
94222f64 | 705 | eligible_storage_live.intersect(&**self.saved_locals); |
dc9dc135 XL |
706 | |
707 | for local in eligible_storage_live.iter() { | |
708 | self.local_conflicts.union_row_with(&eligible_storage_live, local); | |
709 | } | |
710 | ||
711 | if eligible_storage_live.count() > 1 { | |
712 | trace!("at {:?}, eligible_storage_live={:?}", loc, eligible_storage_live); | |
713 | } | |
714 | } | |
715 | } | |
716 | ||
f035d41b XL |
717 | /// Validates the typeck view of the generator against the actual set of types saved between |
718 | /// yield points. | |
719 | fn sanitize_witness<'tcx>( | |
dc9dc135 | 720 | tcx: TyCtxt<'tcx>, |
f035d41b | 721 | body: &Body<'tcx>, |
f035d41b | 722 | witness: Ty<'tcx>, |
fc512014 | 723 | upvars: Vec<Ty<'tcx>>, |
f035d41b | 724 | saved_locals: &GeneratorSavedLocals, |
dc9dc135 | 725 | ) { |
29967ef6 | 726 | let did = body.source.def_id(); |
5099ac24 FG |
727 | let param_env = tcx.param_env(did); |
728 | ||
729 | let allowed_upvars = tcx.normalize_erasing_regions(param_env, upvars); | |
1b1a35ee | 730 | let allowed = match witness.kind() { |
5099ac24 FG |
731 | &ty::GeneratorWitness(interior_tys) => { |
732 | tcx.normalize_erasing_late_bound_regions(param_env, interior_tys) | |
733 | } | |
f035d41b XL |
734 | _ => { |
735 | tcx.sess.delay_span_bug( | |
736 | body.span, | |
1b1a35ee | 737 | &format!("unexpected generator witness type {:?}", witness.kind()), |
f035d41b XL |
738 | ); |
739 | return; | |
740 | } | |
2c00a5a8 | 741 | }; |
ea8adc8c | 742 | |
dc9dc135 | 743 | for (local, decl) in body.local_decls.iter_enumerated() { |
f035d41b XL |
744 | // Ignore locals which are internal or not saved between yields. |
745 | if !saved_locals.contains(local) || decl.internal { | |
ea8adc8c XL |
746 | continue; |
747 | } | |
ba9703b0 | 748 | let decl_ty = tcx.normalize_erasing_regions(param_env, decl.ty); |
ea8adc8c XL |
749 | |
750 | // Sanity check that typeck knows about the type of locals which are | |
751 | // live across a suspension point | |
ba9703b0 | 752 | if !allowed.contains(&decl_ty) && !allowed_upvars.contains(&decl_ty) { |
dfeec247 XL |
753 | span_bug!( |
754 | body.span, | |
755 | "Broken MIR: generator contains type {} in MIR, \ | |
cdc7bbd5 XL |
756 | but typeck only knows about {} and {:?}", |
757 | decl_ty, | |
758 | allowed, | |
759 | allowed_upvars | |
dfeec247 | 760 | ); |
ea8adc8c XL |
761 | } |
762 | } | |
f035d41b XL |
763 | } |
764 | ||
765 | fn compute_layout<'tcx>( | |
766 | liveness: LivenessInfo, | |
767 | body: &mut Body<'tcx>, | |
768 | ) -> ( | |
769 | FxHashMap<Local, (Ty<'tcx>, VariantIdx, usize)>, | |
770 | GeneratorLayout<'tcx>, | |
771 | IndexVec<BasicBlock, Option<BitSet<Local>>>, | |
772 | ) { | |
773 | let LivenessInfo { | |
774 | saved_locals, | |
775 | live_locals_at_suspension_points, | |
776 | source_info_at_suspension_points, | |
777 | storage_conflicts, | |
778 | storage_liveness, | |
779 | } = liveness; | |
ea8adc8c | 780 | |
60c5eb7d | 781 | // Gather live local types and their indices. |
dc9dc135 XL |
782 | let mut locals = IndexVec::<GeneratorSavedLocal, _>::new(); |
783 | let mut tys = IndexVec::<GeneratorSavedLocal, _>::new(); | |
f035d41b | 784 | for (saved_local, local) in saved_locals.iter_enumerated() { |
dc9dc135 | 785 | locals.push(local); |
60c5eb7d | 786 | tys.push(body.local_decls[local].ty); |
f035d41b | 787 | debug!("generator saved local {:?} => {:?}", saved_local, local); |
dc9dc135 | 788 | } |
ea8adc8c | 789 | |
dc9dc135 | 790 | // Leave empty variants for the UNRESUMED, RETURNED, and POISONED states. |
f035d41b XL |
791 | // In debuginfo, these will correspond to the beginning (UNRESUMED) or end |
792 | // (RETURNED, POISONED) of the function. | |
dc9dc135 | 793 | const RESERVED_VARIANTS: usize = 3; |
f035d41b XL |
794 | let body_span = body.source_scopes[OUTERMOST_SOURCE_SCOPE].span; |
795 | let mut variant_source_info: IndexVec<VariantIdx, SourceInfo> = [ | |
796 | SourceInfo::outermost(body_span.shrink_to_lo()), | |
797 | SourceInfo::outermost(body_span.shrink_to_hi()), | |
798 | SourceInfo::outermost(body_span.shrink_to_hi()), | |
799 | ] | |
800 | .iter() | |
801 | .copied() | |
802 | .collect(); | |
48663c56 | 803 | |
dc9dc135 | 804 | // Build the generator variant field list. |
ea8adc8c | 805 | // Create a map from local indices to generator struct indices. |
dc9dc135 XL |
806 | let mut variant_fields: IndexVec<VariantIdx, IndexVec<Field, GeneratorSavedLocal>> = |
807 | iter::repeat(IndexVec::new()).take(RESERVED_VARIANTS).collect(); | |
48663c56 | 808 | let mut remap = FxHashMap::default(); |
dc9dc135 XL |
809 | for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() { |
810 | let variant_index = VariantIdx::from(RESERVED_VARIANTS + suspension_point_idx); | |
811 | let mut fields = IndexVec::new(); | |
812 | for (idx, saved_local) in live_locals.iter().enumerate() { | |
813 | fields.push(saved_local); | |
814 | // Note that if a field is included in multiple variants, we will | |
815 | // just use the first one here. That's fine; fields do not move | |
816 | // around inside generators, so it doesn't matter which variant | |
817 | // index we access them by. | |
818 | remap.entry(locals[saved_local]).or_insert((tys[saved_local], variant_index, idx)); | |
819 | } | |
820 | variant_fields.push(fields); | |
f035d41b | 821 | variant_source_info.push(source_info_at_suspension_points[suspension_point_idx]); |
48663c56 | 822 | } |
dc9dc135 XL |
823 | debug!("generator variant_fields = {:?}", variant_fields); |
824 | debug!("generator storage_conflicts = {:#?}", storage_conflicts); | |
ea8adc8c | 825 | |
f035d41b XL |
826 | let layout = |
827 | GeneratorLayout { field_tys: tys, variant_fields, variant_source_info, storage_conflicts }; | |
ea8adc8c XL |
828 | |
829 | (remap, layout, storage_liveness) | |
830 | } | |
831 | ||
74b04a01 XL |
832 | /// Replaces the entry point of `body` with a block that switches on the generator discriminant and |
833 | /// dispatches to blocks according to `cases`. | |
834 | /// | |
835 | /// After this function, the former entry point of the function will be bb1. | |
dc9dc135 | 836 | fn insert_switch<'tcx>( |
f9f354fc | 837 | body: &mut Body<'tcx>, |
dc9dc135 XL |
838 | cases: Vec<(usize, BasicBlock)>, |
839 | transform: &TransformVisitor<'tcx>, | |
840 | default: TerminatorKind<'tcx>, | |
841 | ) { | |
842 | let default_block = insert_term_block(body, default); | |
843 | let (assign, discr) = transform.get_discr(body); | |
29967ef6 XL |
844 | let switch_targets = |
845 | SwitchTargets::new(cases.iter().map(|(i, bb)| ((*i) as u128, *bb)), default_block); | |
ea8adc8c | 846 | let switch = TerminatorKind::SwitchInt { |
48663c56 XL |
847 | discr: Operand::Move(discr), |
848 | switch_ty: transform.discr_ty, | |
29967ef6 | 849 | targets: switch_targets, |
ea8adc8c XL |
850 | }; |
851 | ||
f9f354fc | 852 | let source_info = SourceInfo::outermost(body.span); |
dfeec247 XL |
853 | body.basic_blocks_mut().raw.insert( |
854 | 0, | |
855 | BasicBlockData { | |
856 | statements: vec![assign], | |
857 | terminator: Some(Terminator { source_info, kind: switch }), | |
858 | is_cleanup: false, | |
859 | }, | |
860 | ); | |
ea8adc8c | 861 | |
dc9dc135 | 862 | let blocks = body.basic_blocks_mut().iter_mut(); |
ea8adc8c XL |
863 | |
864 | for target in blocks.flat_map(|b| b.terminator_mut().successors_mut()) { | |
865 | *target = BasicBlock::new(target.index() + 1); | |
866 | } | |
867 | } | |
868 | ||
29967ef6 | 869 | fn elaborate_generator_drops<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
dfeec247 | 870 | use crate::shim::DropShimElaborator; |
c295e0f8 XL |
871 | use rustc_middle::mir::patch::MirPatch; |
872 | use rustc_mir_dataflow::elaborate_drops::{elaborate_drop, Unwind}; | |
ea8adc8c XL |
873 | |
874 | // Note that `elaborate_drops` only drops the upvars of a generator, and | |
875 | // this is ok because `open_drop` can only be reached within that own | |
876 | // generator's resume function. | |
877 | ||
29967ef6 | 878 | let def_id = body.source.def_id(); |
ea8adc8c | 879 | let param_env = tcx.param_env(def_id); |
ea8adc8c | 880 | |
dfeec247 | 881 | let mut elaborator = DropShimElaborator { body, patch: MirPatch::new(body), tcx, param_env }; |
532ac7d7 | 882 | |
dc9dc135 | 883 | for (block, block_data) in body.basic_blocks().iter_enumerated() { |
532ac7d7 | 884 | let (target, unwind, source_info) = match block_data.terminator() { |
f035d41b XL |
885 | Terminator { source_info, kind: TerminatorKind::Drop { place, target, unwind } } => { |
886 | if let Some(local) = place.as_local() { | |
ba9703b0 | 887 | if local == SELF_ARG { |
e74abb32 XL |
888 | (target, unwind, source_info) |
889 | } else { | |
890 | continue; | |
891 | } | |
892 | } else { | |
893 | continue; | |
894 | } | |
895 | } | |
ea8adc8c XL |
896 | _ => continue, |
897 | }; | |
532ac7d7 | 898 | let unwind = if block_data.is_cleanup { |
ea8adc8c | 899 | Unwind::InCleanup |
532ac7d7 XL |
900 | } else { |
901 | Unwind::To(unwind.unwrap_or_else(|| elaborator.patch.resume_block())) | |
ea8adc8c | 902 | }; |
532ac7d7 XL |
903 | elaborate_drop( |
904 | &mut elaborator, | |
e74abb32 | 905 | *source_info, |
ba9703b0 | 906 | Place::from(SELF_ARG), |
532ac7d7 | 907 | (), |
e74abb32 | 908 | *target, |
532ac7d7 XL |
909 | unwind, |
910 | block, | |
911 | ); | |
ea8adc8c | 912 | } |
dc9dc135 | 913 | elaborator.patch.apply(body); |
ea8adc8c XL |
914 | } |
915 | ||
dc9dc135 XL |
916 | fn create_generator_drop_shim<'tcx>( |
917 | tcx: TyCtxt<'tcx>, | |
918 | transform: &TransformVisitor<'tcx>, | |
dc9dc135 | 919 | gen_ty: Ty<'tcx>, |
f9f354fc | 920 | body: &mut Body<'tcx>, |
dc9dc135 | 921 | drop_clean: BasicBlock, |
f9f354fc | 922 | ) -> Body<'tcx> { |
dc9dc135 | 923 | let mut body = body.clone(); |
74b04a01 | 924 | body.arg_count = 1; // make sure the resume argument is not included here |
ea8adc8c | 925 | |
f9f354fc | 926 | let source_info = SourceInfo::outermost(body.span); |
ea8adc8c | 927 | |
74b04a01 | 928 | let mut cases = create_cases(&mut body, transform, Operation::Drop); |
ea8adc8c | 929 | |
532ac7d7 | 930 | cases.insert(0, (UNRESUMED, drop_clean)); |
ea8adc8c | 931 | |
532ac7d7 XL |
932 | // The returned state and the poisoned state fall through to the default |
933 | // case which is just to return | |
ea8adc8c | 934 | |
dc9dc135 | 935 | insert_switch(&mut body, cases, &transform, TerminatorKind::Return); |
ea8adc8c | 936 | |
dc9dc135 | 937 | for block in body.basic_blocks_mut() { |
ea8adc8c XL |
938 | let kind = &mut block.terminator_mut().kind; |
939 | if let TerminatorKind::GeneratorDrop = *kind { | |
940 | *kind = TerminatorKind::Return; | |
941 | } | |
942 | } | |
943 | ||
944 | // Replace the return variable | |
f9f354fc | 945 | body.local_decls[RETURN_PLACE] = LocalDecl::with_source_info(tcx.mk_unit(), source_info); |
ea8adc8c | 946 | |
ba9703b0 | 947 | make_generator_state_argument_indirect(tcx, &mut body); |
ea8adc8c XL |
948 | |
949 | // Change the generator argument from &mut to *mut | |
f9f354fc XL |
950 | body.local_decls[SELF_ARG] = LocalDecl::with_source_info( |
951 | tcx.mk_ptr(ty::TypeAndMut { ty: gen_ty, mutbl: hir::Mutability::Mut }), | |
ea8adc8c | 952 | source_info, |
f9f354fc | 953 | ); |
a1dfa0c6 XL |
954 | if tcx.sess.opts.debugging_opts.mir_emit_retag { |
955 | // Alias tracking must know we changed the type | |
dfeec247 XL |
956 | body.basic_blocks_mut()[START_BLOCK].statements.insert( |
957 | 0, | |
958 | Statement { | |
959 | source_info, | |
94222f64 | 960 | kind: StatementKind::Retag(RetagKind::Raw, Box::new(Place::from(SELF_ARG))), |
dfeec247 XL |
961 | }, |
962 | ) | |
a1dfa0c6 | 963 | } |
ea8adc8c | 964 | |
ea8adc8c XL |
965 | // Make sure we remove dead blocks to remove |
966 | // unrelated code from the resume part of the function | |
17df50a5 | 967 | simplify::remove_dead_blocks(tcx, &mut body); |
ea8adc8c | 968 | |
29967ef6 | 969 | dump_mir(tcx, None, "generator_drop", &0, &body, |_, _| Ok(())); |
ea8adc8c | 970 | |
dc9dc135 | 971 | body |
ea8adc8c XL |
972 | } |
973 | ||
f9f354fc XL |
974 | fn insert_term_block<'tcx>(body: &mut Body<'tcx>, kind: TerminatorKind<'tcx>) -> BasicBlock { |
975 | let source_info = SourceInfo::outermost(body.span); | |
dc9dc135 | 976 | body.basic_blocks_mut().push(BasicBlockData { |
ea8adc8c | 977 | statements: Vec::new(), |
dfeec247 | 978 | terminator: Some(Terminator { source_info, kind }), |
ea8adc8c | 979 | is_cleanup: false, |
f9f354fc | 980 | }) |
ea8adc8c XL |
981 | } |
982 | ||
dc9dc135 XL |
983 | fn insert_panic_block<'tcx>( |
984 | tcx: TyCtxt<'tcx>, | |
f9f354fc | 985 | body: &mut Body<'tcx>, |
dc9dc135 XL |
986 | message: AssertMessage<'tcx>, |
987 | ) -> BasicBlock { | |
988 | let assert_block = BasicBlock::new(body.basic_blocks().len()); | |
ea8adc8c | 989 | let term = TerminatorKind::Assert { |
94222f64 | 990 | cond: Operand::Constant(Box::new(Constant { |
dc9dc135 | 991 | span: body.span, |
b7449926 | 992 | user_ty: None, |
6a06907d | 993 | literal: ty::Const::from_bool(tcx, false).into(), |
94222f64 | 994 | })), |
ea8adc8c XL |
995 | expected: true, |
996 | msg: message, | |
997 | target: assert_block, | |
998 | cleanup: None, | |
999 | }; | |
1000 | ||
f9f354fc | 1001 | let source_info = SourceInfo::outermost(body.span); |
dc9dc135 | 1002 | body.basic_blocks_mut().push(BasicBlockData { |
ea8adc8c | 1003 | statements: Vec::new(), |
dfeec247 | 1004 | terminator: Some(Terminator { source_info, kind: term }), |
ea8adc8c XL |
1005 | is_cleanup: false, |
1006 | }); | |
1007 | ||
1008 | assert_block | |
1009 | } | |
1010 | ||
6a06907d | 1011 | fn can_return<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>, param_env: ty::ParamEnv<'tcx>) -> bool { |
ba9703b0 | 1012 | // Returning from a function with an uninhabited return type is undefined behavior. |
6a06907d | 1013 | if tcx.conservative_is_privately_uninhabited(param_env.and(body.return_ty())) { |
ba9703b0 XL |
1014 | return false; |
1015 | } | |
1016 | ||
1017 | // If there's a return terminator the function may return. | |
1018 | for block in body.basic_blocks() { | |
1019 | if let TerminatorKind::Return = block.terminator().kind { | |
1020 | return true; | |
1021 | } | |
1022 | } | |
1023 | ||
1024 | // Otherwise the function can't return. | |
1025 | false | |
1026 | } | |
1027 | ||
1028 | fn can_unwind<'tcx>(tcx: TyCtxt<'tcx>, body: &Body<'tcx>) -> bool { | |
1029 | // Nothing can unwind when landing pads are off. | |
f9f354fc | 1030 | if tcx.sess.panic_strategy() == PanicStrategy::Abort { |
ba9703b0 XL |
1031 | return false; |
1032 | } | |
1033 | ||
1034 | // Unwinds can only start at certain terminators. | |
1035 | for block in body.basic_blocks() { | |
1036 | match block.terminator().kind { | |
1037 | // These never unwind. | |
1038 | TerminatorKind::Goto { .. } | |
1039 | | TerminatorKind::SwitchInt { .. } | |
1040 | | TerminatorKind::Abort | |
1041 | | TerminatorKind::Return | |
1042 | | TerminatorKind::Unreachable | |
1043 | | TerminatorKind::GeneratorDrop | |
f035d41b | 1044 | | TerminatorKind::FalseEdge { .. } |
f9f354fc XL |
1045 | | TerminatorKind::FalseUnwind { .. } |
1046 | | TerminatorKind::InlineAsm { .. } => {} | |
ba9703b0 XL |
1047 | |
1048 | // Resume will *continue* unwinding, but if there's no other unwinding terminator it | |
1049 | // will never be reached. | |
1050 | TerminatorKind::Resume => {} | |
1051 | ||
1052 | TerminatorKind::Yield { .. } => { | |
1053 | unreachable!("`can_unwind` called before generator transform") | |
1054 | } | |
1055 | ||
1056 | // These may unwind. | |
1057 | TerminatorKind::Drop { .. } | |
1058 | | TerminatorKind::DropAndReplace { .. } | |
1059 | | TerminatorKind::Call { .. } | |
1060 | | TerminatorKind::Assert { .. } => return true, | |
1061 | } | |
1062 | } | |
1063 | ||
1064 | // If we didn't find an unwinding terminator, the function cannot unwind. | |
1065 | false | |
1066 | } | |
1067 | ||
dc9dc135 XL |
1068 | fn create_generator_resume_function<'tcx>( |
1069 | tcx: TyCtxt<'tcx>, | |
1070 | transform: TransformVisitor<'tcx>, | |
f9f354fc | 1071 | body: &mut Body<'tcx>, |
ba9703b0 | 1072 | can_return: bool, |
dc9dc135 | 1073 | ) { |
ba9703b0 XL |
1074 | let can_unwind = can_unwind(tcx, body); |
1075 | ||
ea8adc8c | 1076 | // Poison the generator when it unwinds |
ba9703b0 | 1077 | if can_unwind { |
f9f354fc XL |
1078 | let source_info = SourceInfo::outermost(body.span); |
1079 | let poison_block = body.basic_blocks_mut().push(BasicBlockData { | |
ba9703b0 XL |
1080 | statements: vec![transform.set_discr(VariantIdx::new(POISONED), source_info)], |
1081 | terminator: Some(Terminator { source_info, kind: TerminatorKind::Resume }), | |
1082 | is_cleanup: true, | |
1083 | }); | |
1084 | ||
1085 | for (idx, block) in body.basic_blocks_mut().iter_enumerated_mut() { | |
1086 | let source_info = block.terminator().source_info; | |
1087 | ||
1088 | if let TerminatorKind::Resume = block.terminator().kind { | |
1089 | // An existing `Resume` terminator is redirected to jump to our dedicated | |
1090 | // "poisoning block" above. | |
1091 | if idx != poison_block { | |
1092 | *block.terminator_mut() = Terminator { | |
1093 | source_info, | |
1094 | kind: TerminatorKind::Goto { target: poison_block }, | |
1095 | }; | |
1096 | } | |
1097 | } else if !block.is_cleanup { | |
1098 | // Any terminators that *can* unwind but don't have an unwind target set are also | |
1099 | // pointed at our poisoning block (unless they're part of the cleanup path). | |
1100 | if let Some(unwind @ None) = block.terminator_mut().unwind_mut() { | |
1101 | *unwind = Some(poison_block); | |
1102 | } | |
1103 | } | |
ea8adc8c XL |
1104 | } |
1105 | } | |
1106 | ||
74b04a01 | 1107 | let mut cases = create_cases(body, &transform, Operation::Resume); |
ea8adc8c | 1108 | |
ba9703b0 | 1109 | use rustc_middle::mir::AssertKind::{ResumedAfterPanic, ResumedAfterReturn}; |
83c7162d | 1110 | |
532ac7d7 XL |
1111 | // Jump to the entry point on the unresumed |
1112 | cases.insert(0, (UNRESUMED, BasicBlock::new(0))); | |
60c5eb7d XL |
1113 | |
1114 | // Panic when resumed on the returned or poisoned state | |
6a06907d | 1115 | let generator_kind = body.generator_kind().unwrap(); |
ba9703b0 XL |
1116 | |
1117 | if can_unwind { | |
1118 | cases.insert( | |
1119 | 1, | |
1120 | (POISONED, insert_panic_block(tcx, body, ResumedAfterPanic(generator_kind))), | |
1121 | ); | |
1122 | } | |
1123 | ||
1124 | if can_return { | |
1125 | cases.insert( | |
1126 | 1, | |
1127 | (RETURNED, insert_panic_block(tcx, body, ResumedAfterReturn(generator_kind))), | |
1128 | ); | |
1129 | } | |
ea8adc8c | 1130 | |
dc9dc135 | 1131 | insert_switch(body, cases, &transform, TerminatorKind::Unreachable); |
ea8adc8c | 1132 | |
ba9703b0 | 1133 | make_generator_state_argument_indirect(tcx, body); |
dc9dc135 | 1134 | make_generator_state_argument_pinned(tcx, body); |
ea8adc8c | 1135 | |
ea8adc8c XL |
1136 | // Make sure we remove dead blocks to remove |
1137 | // unrelated code from the drop part of the function | |
17df50a5 | 1138 | simplify::remove_dead_blocks(tcx, body); |
ea8adc8c | 1139 | |
29967ef6 | 1140 | dump_mir(tcx, None, "generator_resume", &0, body, |_, _| Ok(())); |
ea8adc8c XL |
1141 | } |
1142 | ||
f9f354fc | 1143 | fn insert_clean_drop(body: &mut Body<'_>) -> BasicBlock { |
dc9dc135 | 1144 | let return_block = insert_term_block(body, TerminatorKind::Return); |
ea8adc8c | 1145 | |
f035d41b XL |
1146 | let term = |
1147 | TerminatorKind::Drop { place: Place::from(SELF_ARG), target: return_block, unwind: None }; | |
f9f354fc XL |
1148 | let source_info = SourceInfo::outermost(body.span); |
1149 | ||
1150 | // Create a block to destroy an unresumed generators. This can only destroy upvars. | |
dc9dc135 | 1151 | body.basic_blocks_mut().push(BasicBlockData { |
ea8adc8c | 1152 | statements: Vec::new(), |
dfeec247 | 1153 | terminator: Some(Terminator { source_info, kind: term }), |
ea8adc8c | 1154 | is_cleanup: false, |
f9f354fc | 1155 | }) |
ea8adc8c XL |
1156 | } |
1157 | ||
74b04a01 XL |
1158 | /// An operation that can be performed on a generator. |
1159 | #[derive(PartialEq, Copy, Clone)] | |
1160 | enum Operation { | |
1161 | Resume, | |
1162 | Drop, | |
1163 | } | |
1164 | ||
1165 | impl Operation { | |
1166 | fn target_block(self, point: &SuspensionPoint<'_>) -> Option<BasicBlock> { | |
1167 | match self { | |
1168 | Operation::Resume => Some(point.resume), | |
1169 | Operation::Drop => point.drop, | |
1170 | } | |
1171 | } | |
1172 | } | |
1173 | ||
1174 | fn create_cases<'tcx>( | |
f9f354fc | 1175 | body: &mut Body<'tcx>, |
dc9dc135 | 1176 | transform: &TransformVisitor<'tcx>, |
74b04a01 XL |
1177 | operation: Operation, |
1178 | ) -> Vec<(usize, BasicBlock)> { | |
f9f354fc | 1179 | let source_info = SourceInfo::outermost(body.span); |
ea8adc8c | 1180 | |
dfeec247 XL |
1181 | transform |
1182 | .suspension_points | |
1183 | .iter() | |
1184 | .filter_map(|point| { | |
1185 | // Find the target for this suspension point, if applicable | |
74b04a01 | 1186 | operation.target_block(point).map(|target| { |
dfeec247 XL |
1187 | let mut statements = Vec::new(); |
1188 | ||
1189 | // Create StorageLive instructions for locals with live storage | |
1190 | for i in 0..(body.local_decls.len()) { | |
74b04a01 XL |
1191 | if i == 2 { |
1192 | // The resume argument is live on function entry. Don't insert a | |
1193 | // `StorageLive`, or the following `Assign` will read from uninitialized | |
1194 | // memory. | |
1195 | continue; | |
1196 | } | |
1197 | ||
dfeec247 | 1198 | let l = Local::new(i); |
ba9703b0 XL |
1199 | let needs_storage_live = point.storage_liveness.contains(l) |
1200 | && !transform.remap.contains_key(&l) | |
1201 | && !transform.always_live_locals.contains(l); | |
1202 | if needs_storage_live { | |
dfeec247 XL |
1203 | statements |
1204 | .push(Statement { source_info, kind: StatementKind::StorageLive(l) }); | |
1205 | } | |
ea8adc8c | 1206 | } |
ea8adc8c | 1207 | |
74b04a01 XL |
1208 | if operation == Operation::Resume { |
1209 | // Move the resume argument to the destination place of the `Yield` terminator | |
1210 | let resume_arg = Local::new(2); // 0 = return, 1 = self | |
1211 | statements.push(Statement { | |
1212 | source_info, | |
94222f64 | 1213 | kind: StatementKind::Assign(Box::new(( |
74b04a01 XL |
1214 | point.resume_arg, |
1215 | Rvalue::Use(Operand::Move(resume_arg.into())), | |
94222f64 | 1216 | ))), |
74b04a01 XL |
1217 | }); |
1218 | } | |
1219 | ||
dfeec247 | 1220 | // Then jump to the real target |
f9f354fc | 1221 | let block = body.basic_blocks_mut().push(BasicBlockData { |
dfeec247 XL |
1222 | statements, |
1223 | terminator: Some(Terminator { | |
1224 | source_info, | |
1225 | kind: TerminatorKind::Goto { target }, | |
1226 | }), | |
1227 | is_cleanup: false, | |
1228 | }); | |
ea8adc8c | 1229 | |
dfeec247 XL |
1230 | (point.state, block) |
1231 | }) | |
ea8adc8c | 1232 | }) |
dfeec247 | 1233 | .collect() |
ea8adc8c XL |
1234 | } |
1235 | ||
e1599b0c | 1236 | impl<'tcx> MirPass<'tcx> for StateTransform { |
a2a8927a | 1237 | fn phase_change(&self) -> Option<MirPhase> { |
5e7ed085 | 1238 | Some(MirPhase::GeneratorsLowered) |
a2a8927a XL |
1239 | } |
1240 | ||
29967ef6 | 1241 | fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) { |
5099ac24 | 1242 | let Some(yield_ty) = body.yield_ty() else { |
ea8adc8c | 1243 | // This only applies to generators |
dfeec247 | 1244 | return; |
ea8adc8c XL |
1245 | }; |
1246 | ||
6a06907d | 1247 | assert!(body.generator_drop().is_none()); |
ea8adc8c | 1248 | |
ea8adc8c | 1249 | // The first argument is the generator type passed by value |
dc9dc135 | 1250 | let gen_ty = body.local_decls.raw[1].ty; |
ea8adc8c | 1251 | |
2c00a5a8 | 1252 | // Get the interior types and substs which typeck computed |
1b1a35ee | 1253 | let (upvars, interior, discr_ty, movable) = match *gen_ty.kind() { |
b7449926 | 1254 | ty::Generator(_, substs, movability) => { |
e74abb32 | 1255 | let substs = substs.as_generator(); |
dfeec247 | 1256 | ( |
ba9703b0 XL |
1257 | substs.upvar_tys().collect(), |
1258 | substs.witness(), | |
dfeec247 XL |
1259 | substs.discr_ty(tcx), |
1260 | movability == hir::Movability::Movable, | |
1261 | ) | |
2c00a5a8 | 1262 | } |
f035d41b XL |
1263 | _ => { |
1264 | tcx.sess | |
1265 | .delay_span_bug(body.span, &format!("unexpected generator type {}", gen_ty)); | |
1266 | return; | |
1267 | } | |
2c00a5a8 XL |
1268 | }; |
1269 | ||
ea8adc8c | 1270 | // Compute GeneratorState<yield_ty, return_ty> |
3dfed10e | 1271 | let state_did = tcx.require_lang_item(LangItem::GeneratorState, None); |
ea8adc8c | 1272 | let state_adt_ref = tcx.adt_def(state_did); |
dfeec247 | 1273 | let state_substs = tcx.intern_substs(&[yield_ty.into(), body.return_ty().into()]); |
ea8adc8c XL |
1274 | let ret_ty = tcx.mk_adt(state_adt_ref, state_substs); |
1275 | ||
ff7c6d11 XL |
1276 | // We rename RETURN_PLACE which has type mir.return_ty to new_ret_local |
1277 | // RETURN_PLACE then is a fresh unused local with type ret_ty. | |
74b04a01 XL |
1278 | let new_ret_local = replace_local(RETURN_PLACE, ret_ty, body, tcx); |
1279 | ||
1280 | // We also replace the resume argument and insert an `Assign`. | |
1281 | // This is needed because the resume argument `_2` might be live across a `yield`, in which | |
1282 | // case there is no `Assign` to it that the transform can turn into a store to the generator | |
1283 | // state. After the yield the slot in the generator state would then be uninitialized. | |
1284 | let resume_local = Local::new(2); | |
1285 | let new_resume_local = | |
1286 | replace_local(resume_local, body.local_decls[resume_local].ty, body, tcx); | |
1287 | ||
1288 | // When first entering the generator, move the resume argument into its new local. | |
f9f354fc | 1289 | let source_info = SourceInfo::outermost(body.span); |
74b04a01 XL |
1290 | let stmts = &mut body.basic_blocks_mut()[BasicBlock::new(0)].statements; |
1291 | stmts.insert( | |
1292 | 0, | |
1293 | Statement { | |
1294 | source_info, | |
94222f64 | 1295 | kind: StatementKind::Assign(Box::new(( |
74b04a01 XL |
1296 | new_resume_local.into(), |
1297 | Rvalue::Use(Operand::Move(resume_local.into())), | |
94222f64 | 1298 | ))), |
74b04a01 XL |
1299 | }, |
1300 | ); | |
ea8adc8c | 1301 | |
ba9703b0 XL |
1302 | let always_live_locals = storage::AlwaysLiveLocals::new(&body); |
1303 | ||
f035d41b | 1304 | let liveness_info = |
29967ef6 | 1305 | locals_live_across_suspend_points(tcx, body, &always_live_locals, movable); |
f035d41b | 1306 | |
fc512014 | 1307 | sanitize_witness(tcx, body, interior, upvars, &liveness_info.saved_locals); |
f035d41b XL |
1308 | |
1309 | if tcx.sess.opts.debugging_opts.validate_mir { | |
1310 | let mut vis = EnsureGeneratorFieldAssignmentsNeverAlias { | |
1311 | assigned_local: None, | |
1312 | saved_locals: &liveness_info.saved_locals, | |
1313 | storage_conflicts: &liveness_info.storage_conflicts, | |
1314 | }; | |
1315 | ||
1316 | vis.visit_body(body); | |
1317 | } | |
1318 | ||
ea8adc8c XL |
1319 | // Extract locals which are live across suspension point into `layout` |
1320 | // `remap` gives a mapping from local indices onto generator struct indices | |
1321 | // `storage_liveness` tells us which locals have live storage at suspension points | |
f035d41b | 1322 | let (remap, layout, storage_liveness) = compute_layout(liveness_info, body); |
ba9703b0 | 1323 | |
6a06907d | 1324 | let can_return = can_return(tcx, body, tcx.param_env(body.source.def_id())); |
ea8adc8c | 1325 | |
ff7c6d11 | 1326 | // Run the transformation which converts Places from Local to generator struct |
ea8adc8c XL |
1327 | // accesses for locals in `remap`. |
1328 | // It also rewrites `return x` and `yield y` as writing a new generator state and returning | |
1329 | // GeneratorState::Complete(x) and GeneratorState::Yielded(y) respectively. | |
1330 | let mut transform = TransformVisitor { | |
1331 | tcx, | |
1332 | state_adt_ref, | |
1333 | state_substs, | |
1334 | remap, | |
1335 | storage_liveness, | |
ba9703b0 | 1336 | always_live_locals, |
ea8adc8c XL |
1337 | suspension_points: Vec::new(), |
1338 | new_ret_local, | |
48663c56 | 1339 | discr_ty, |
ea8adc8c | 1340 | }; |
dc9dc135 | 1341 | transform.visit_body(body); |
ea8adc8c | 1342 | |
74b04a01 | 1343 | // Update our MIR struct to reflect the changes we've made |
74b04a01 | 1344 | body.arg_count = 2; // self, resume arg |
dc9dc135 | 1345 | body.spread_arg = None; |
6a06907d XL |
1346 | |
1347 | body.generator.as_mut().unwrap().yield_ty = None; | |
1348 | body.generator.as_mut().unwrap().generator_layout = Some(layout); | |
ea8adc8c XL |
1349 | |
1350 | // Insert `drop(generator_struct)` which is used to drop upvars for generators in | |
532ac7d7 | 1351 | // the unresumed state. |
ea8adc8c | 1352 | // This is expanded to a drop ladder in `elaborate_generator_drops`. |
dc9dc135 | 1353 | let drop_clean = insert_clean_drop(body); |
ea8adc8c | 1354 | |
29967ef6 | 1355 | dump_mir(tcx, None, "generator_pre-elab", &0, body, |_, _| Ok(())); |
ea8adc8c XL |
1356 | |
1357 | // Expand `drop(generator_struct)` to a drop ladder which destroys upvars. | |
1358 | // If any upvars are moved out of, drop elaboration will handle upvar destruction. | |
1359 | // However we need to also elaborate the code generated by `insert_clean_drop`. | |
29967ef6 | 1360 | elaborate_generator_drops(tcx, body); |
ea8adc8c | 1361 | |
29967ef6 | 1362 | dump_mir(tcx, None, "generator_post-transform", &0, body, |_, _| Ok(())); |
ea8adc8c XL |
1363 | |
1364 | // Create a copy of our MIR and use it to create the drop shim for the generator | |
29967ef6 | 1365 | let drop_shim = create_generator_drop_shim(tcx, &transform, gen_ty, body, drop_clean); |
ea8adc8c | 1366 | |
6a06907d | 1367 | body.generator.as_mut().unwrap().generator_drop = Some(drop_shim); |
ea8adc8c XL |
1368 | |
1369 | // Create the Generator::resume function | |
29967ef6 | 1370 | create_generator_resume_function(tcx, transform, body, can_return); |
ea8adc8c XL |
1371 | } |
1372 | } | |
f035d41b XL |
1373 | |
1374 | /// Looks for any assignments between locals (e.g., `_4 = _5`) that will both be converted to fields | |
1375 | /// in the generator state machine but whose storage is not marked as conflicting | |
1376 | /// | |
1377 | /// Validation needs to happen immediately *before* `TransformVisitor` is invoked, not after. | |
1378 | /// | |
1379 | /// This condition would arise when the assignment is the last use of `_5` but the initial | |
1380 | /// definition of `_4` if we weren't extra careful to mark all locals used inside a statement as | |
1381 | /// conflicting. Non-conflicting generator saved locals may be stored at the same location within | |
1382 | /// the generator state machine, which would result in ill-formed MIR: the left-hand and right-hand | |
1383 | /// sides of an assignment may not alias. This caused a miscompilation in [#73137]. | |
1384 | /// | |
1385 | /// [#73137]: https://github.com/rust-lang/rust/issues/73137 | |
1386 | struct EnsureGeneratorFieldAssignmentsNeverAlias<'a> { | |
1387 | saved_locals: &'a GeneratorSavedLocals, | |
1388 | storage_conflicts: &'a BitMatrix<GeneratorSavedLocal, GeneratorSavedLocal>, | |
1389 | assigned_local: Option<GeneratorSavedLocal>, | |
1390 | } | |
1391 | ||
1392 | impl EnsureGeneratorFieldAssignmentsNeverAlias<'_> { | |
1393 | fn saved_local_for_direct_place(&self, place: Place<'_>) -> Option<GeneratorSavedLocal> { | |
1394 | if place.is_indirect() { | |
1395 | return None; | |
1396 | } | |
1397 | ||
1398 | self.saved_locals.get(place.local) | |
1399 | } | |
1400 | ||
a2a8927a | 1401 | fn check_assigned_place(&mut self, place: Place<'_>, f: impl FnOnce(&mut Self)) { |
f035d41b XL |
1402 | if let Some(assigned_local) = self.saved_local_for_direct_place(place) { |
1403 | assert!(self.assigned_local.is_none(), "`check_assigned_place` must not recurse"); | |
1404 | ||
1405 | self.assigned_local = Some(assigned_local); | |
1406 | f(self); | |
1407 | self.assigned_local = None; | |
1408 | } | |
1409 | } | |
1410 | } | |
1411 | ||
a2a8927a | 1412 | impl<'tcx> Visitor<'tcx> for EnsureGeneratorFieldAssignmentsNeverAlias<'_> { |
f035d41b | 1413 | fn visit_place(&mut self, place: &Place<'tcx>, context: PlaceContext, location: Location) { |
5e7ed085 FG |
1414 | let Some(lhs) = self.assigned_local else { |
1415 | // This visitor only invokes `visit_place` for the right-hand side of an assignment | |
1416 | // and only after setting `self.assigned_local`. However, the default impl of | |
1417 | // `Visitor::super_body` may call `visit_place` with a `NonUseContext` for places | |
1418 | // with debuginfo. Ignore them here. | |
1419 | assert!(!context.is_use()); | |
1420 | return; | |
f035d41b XL |
1421 | }; |
1422 | ||
5e7ed085 | 1423 | let Some(rhs) = self.saved_local_for_direct_place(*place) else { return }; |
f035d41b XL |
1424 | |
1425 | if !self.storage_conflicts.contains(lhs, rhs) { | |
1426 | bug!( | |
1427 | "Assignment between generator saved locals whose storage is not \ | |
1428 | marked as conflicting: {:?}: {:?} = {:?}", | |
1429 | location, | |
1430 | lhs, | |
1431 | rhs, | |
1432 | ); | |
1433 | } | |
1434 | } | |
1435 | ||
1436 | fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) { | |
1437 | match &statement.kind { | |
1438 | StatementKind::Assign(box (lhs, rhs)) => { | |
1439 | self.check_assigned_place(*lhs, |this| this.visit_rvalue(rhs, location)); | |
1440 | } | |
1441 | ||
f035d41b XL |
1442 | StatementKind::FakeRead(..) |
1443 | | StatementKind::SetDiscriminant { .. } | |
1444 | | StatementKind::StorageLive(_) | |
1445 | | StatementKind::StorageDead(_) | |
1446 | | StatementKind::Retag(..) | |
1447 | | StatementKind::AscribeUserType(..) | |
3dfed10e | 1448 | | StatementKind::Coverage(..) |
6a06907d | 1449 | | StatementKind::CopyNonOverlapping(..) |
f035d41b XL |
1450 | | StatementKind::Nop => {} |
1451 | } | |
1452 | } | |
1453 | ||
1454 | fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) { | |
1455 | // Checking for aliasing in terminators is probably overkill, but until we have actual | |
1456 | // semantics, we should be conservative here. | |
1457 | match &terminator.kind { | |
1458 | TerminatorKind::Call { | |
1459 | func, | |
1460 | args, | |
1461 | destination: Some((dest, _)), | |
1462 | cleanup: _, | |
1463 | from_hir_call: _, | |
1464 | fn_span: _, | |
1465 | } => { | |
1466 | self.check_assigned_place(*dest, |this| { | |
1467 | this.visit_operand(func, location); | |
1468 | for arg in args { | |
1469 | this.visit_operand(arg, location); | |
1470 | } | |
1471 | }); | |
1472 | } | |
1473 | ||
1474 | TerminatorKind::Yield { value, resume: _, resume_arg, drop: _ } => { | |
1475 | self.check_assigned_place(*resume_arg, |this| this.visit_operand(value, location)); | |
1476 | } | |
1477 | ||
1478 | // FIXME: Does `asm!` have any aliasing requirements? | |
1479 | TerminatorKind::InlineAsm { .. } => {} | |
1480 | ||
1481 | TerminatorKind::Call { .. } | |
1482 | | TerminatorKind::Goto { .. } | |
1483 | | TerminatorKind::SwitchInt { .. } | |
1484 | | TerminatorKind::Resume | |
1485 | | TerminatorKind::Abort | |
1486 | | TerminatorKind::Return | |
1487 | | TerminatorKind::Unreachable | |
1488 | | TerminatorKind::Drop { .. } | |
1489 | | TerminatorKind::DropAndReplace { .. } | |
1490 | | TerminatorKind::Assert { .. } | |
1491 | | TerminatorKind::GeneratorDrop | |
1492 | | TerminatorKind::FalseEdge { .. } | |
1493 | | TerminatorKind::FalseUnwind { .. } => {} | |
1494 | } | |
1495 | } | |
1496 | } |