|
71 | 71 |
|
72 | 72 | use rustc_data_structures::unord::UnordMap;
|
73 | 73 | use rustc_hir as hir;
|
74 |
| -use rustc_middle::hir::place::{PlaceBase, Projection, ProjectionKind}; |
| 74 | +use rustc_middle::hir::place::{Projection, ProjectionKind}; |
75 | 75 | use rustc_middle::mir::visit::MutVisitor;
|
76 | 76 | use rustc_middle::mir::{self, dump_mir, MirPass};
|
77 | 77 | use rustc_middle::ty::{self, InstanceDef, Ty, TyCtxt, TypeVisitableExt};
|
@@ -124,62 +124,10 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
|
124 | 124 | .tuple_fields()
|
125 | 125 | .len();
|
126 | 126 |
|
127 |
| - let mut field_remapping = UnordMap::default(); |
128 |
| - |
129 |
| - // One parent capture may correspond to several child captures if we end up |
130 |
| - // refining the set of captures via edition-2021 precise captures. We want to |
131 |
| - // match up any number of child captures with one parent capture, so we keep |
132 |
| - // peeking off this `Peekable` until the child doesn't match anymore. |
133 |
| - let mut parent_captures = |
134 |
| - tcx.closure_captures(parent_def_id).iter().copied().enumerate().peekable(); |
135 |
| - // Make sure we use every field at least once, b/c why are we capturing something |
136 |
| - // if it's not used in the inner coroutine. |
137 |
| - let mut field_used_at_least_once = false; |
138 |
| - |
139 |
| - for (child_field_idx, child_capture) in tcx |
140 |
| - .closure_captures(coroutine_def_id) |
141 |
| - .iter() |
142 |
| - .copied() |
143 |
| - // By construction we capture all the args first. |
144 |
| - .skip(num_args) |
145 |
| - .enumerate() |
146 |
| - { |
147 |
| - loop { |
148 |
| - let Some(&(parent_field_idx, parent_capture)) = parent_captures.peek() else { |
149 |
| - bug!("we ran out of parent captures!") |
150 |
| - }; |
151 |
| - |
152 |
| - let PlaceBase::Upvar(parent_base) = parent_capture.place.base else { |
153 |
| - bug!("expected capture to be an upvar"); |
154 |
| - }; |
155 |
| - let PlaceBase::Upvar(child_base) = child_capture.place.base else { |
156 |
| - bug!("expected capture to be an upvar"); |
157 |
| - }; |
158 |
| - |
159 |
| - assert!( |
160 |
| - child_capture.place.projections.len() >= parent_capture.place.projections.len() |
161 |
| - ); |
162 |
| - // A parent matches a child they share the same prefix of projections. |
163 |
| - // The child may have more, if it is capturing sub-fields out of |
164 |
| - // something that is captured by-move in the parent closure. |
165 |
| - if parent_base.var_path.hir_id != child_base.var_path.hir_id |
166 |
| - || !std::iter::zip( |
167 |
| - &child_capture.place.projections, |
168 |
| - &parent_capture.place.projections, |
169 |
| - ) |
170 |
| - .all(|(child, parent)| child.kind == parent.kind) |
171 |
| - { |
172 |
| - // Make sure the field was used at least once. |
173 |
| - assert!( |
174 |
| - field_used_at_least_once, |
175 |
| - "we captured {parent_capture:#?} but it was not used in the child coroutine?" |
176 |
| - ); |
177 |
| - field_used_at_least_once = false; |
178 |
| - // Skip this field. |
179 |
| - let _ = parent_captures.next().unwrap(); |
180 |
| - continue; |
181 |
| - } |
182 |
| - |
| 127 | + let field_remapping: UnordMap<_, _> = ty::analyze_coroutine_closure_captures( |
| 128 | + tcx.closure_captures(parent_def_id).iter().copied(), |
| 129 | + tcx.closure_captures(coroutine_def_id).iter().skip(num_args).copied(), |
| 130 | + |(parent_field_idx, parent_capture), (child_field_idx, child_capture)| { |
183 | 131 | // Store this set of additional projections (fields and derefs).
|
184 | 132 | // We need to re-apply them later.
|
185 | 133 | let child_precise_captures =
|
@@ -210,26 +158,18 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
|
210 | 158 | ),
|
211 | 159 | };
|
212 | 160 |
|
213 |
| - field_remapping.insert( |
| 161 | + ( |
214 | 162 | FieldIdx::from_usize(child_field_idx + num_args),
|
215 | 163 | (
|
216 | 164 | FieldIdx::from_usize(parent_field_idx + num_args),
|
217 | 165 | parent_capture_ty,
|
218 | 166 | needs_deref,
|
219 | 167 | child_precise_captures,
|
220 | 168 | ),
|
221 |
| - ); |
222 |
| - |
223 |
| - field_used_at_least_once = true; |
224 |
| - break; |
225 |
| - } |
226 |
| - } |
227 |
| - |
228 |
| - // Pop the last parent capture |
229 |
| - if field_used_at_least_once { |
230 |
| - let _ = parent_captures.next().unwrap(); |
231 |
| - } |
232 |
| - assert_eq!(parent_captures.next(), None, "leftover parent captures?"); |
| 169 | + ) |
| 170 | + }, |
| 171 | + ) |
| 172 | + .collect(); |
233 | 173 |
|
234 | 174 | if coroutine_kind == ty::ClosureKind::FnOnce {
|
235 | 175 | assert_eq!(field_remapping.len(), tcx.closure_captures(parent_def_id).len());
|
|
0 commit comments