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