From 688e531925577de3d8f57f717976f7d2d69fd45b Mon Sep 17 00:00:00 2001 From: Oli Scherer Date: Tue, 9 Apr 2024 07:31:20 +0000 Subject: [PATCH 1/5] Split out a complex if condition into a named function --- .../src/coroutine/by_move_body.rs | 36 +++++++++---------- 1 file changed, 18 insertions(+), 18 deletions(-) diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs index 320d8fd3977..d88c7bebc73 100644 --- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs +++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs @@ -148,27 +148,10 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody { let Some(&(parent_field_idx, parent_capture)) = parent_captures.peek() else { bug!("we ran out of parent captures!") }; - - let PlaceBase::Upvar(parent_base) = parent_capture.place.base else { - bug!("expected capture to be an upvar"); - }; - let PlaceBase::Upvar(child_base) = child_capture.place.base else { - bug!("expected capture to be an upvar"); - }; - - assert!( - child_capture.place.projections.len() >= parent_capture.place.projections.len() - ); // A parent matches a child they share the same prefix of projections. // The child may have more, if it is capturing sub-fields out of // something that is captured by-move in the parent closure. - if parent_base.var_path.hir_id != child_base.var_path.hir_id - || !std::iter::zip( - &child_capture.place.projections, - &parent_capture.place.projections, - ) - .all(|(child, parent)| child.kind == parent.kind) - { + if !child_prefix_matches_parent_projections(parent_capture, child_capture) { // Make sure the field was used at least once. assert!( field_used_at_least_once, @@ -258,6 +241,23 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody { } } +fn child_prefix_matches_parent_projections( + parent_capture: &ty::CapturedPlace<'_>, + child_capture: &ty::CapturedPlace<'_>, +) -> bool { + let PlaceBase::Upvar(parent_base) = parent_capture.place.base else { + bug!("expected capture to be an upvar"); + }; + let PlaceBase::Upvar(child_base) = child_capture.place.base else { + bug!("expected capture to be an upvar"); + }; + + assert!(child_capture.place.projections.len() >= parent_capture.place.projections.len()); + parent_base.var_path.hir_id == child_base.var_path.hir_id + && std::iter::zip(&child_capture.place.projections, &parent_capture.place.projections) + .all(|(child, parent)| child.kind == parent.kind) +} + struct MakeByMoveBody<'tcx> { tcx: TyCtxt<'tcx>, field_remapping: UnordMap, bool, &'tcx [Projection<'tcx>])>, From 626fd4b258f71ba5f044c31e8e4d50d730d14a1f Mon Sep 17 00:00:00 2001 From: Oli Scherer Date: Tue, 9 Apr 2024 07:48:55 +0000 Subject: [PATCH 2/5] prefer `expect` over `let else bug!` --- compiler/rustc_mir_transform/src/coroutine/by_move_body.rs | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs index d88c7bebc73..f0a00ebb586 100644 --- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs +++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs @@ -145,9 +145,8 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody { .enumerate() { loop { - let Some(&(parent_field_idx, parent_capture)) = parent_captures.peek() else { - bug!("we ran out of parent captures!") - }; + let &(parent_field_idx, parent_capture) = + parent_captures.peek().expect("we ran out of parent captures!"); // A parent matches a child they share the same prefix of projections. // The child may have more, if it is capturing sub-fields out of // something that is captured by-move in the parent closure. From f014e91a38d64caa3e1231bb86cfb7f1dbda279f Mon Sep 17 00:00:00 2001 From: Oli Scherer Date: Tue, 9 Apr 2024 07:53:12 +0000 Subject: [PATCH 3/5] Shrink a loop to its looping part and move out the part that runs after the loop --- .../src/coroutine/by_move_body.rs | 108 +++++++++--------- 1 file changed, 54 insertions(+), 54 deletions(-) diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs index f0a00ebb586..7979b85225c 100644 --- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs +++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs @@ -144,67 +144,67 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody { .skip(num_args) .enumerate() { + let (mut parent_field_idx, mut parent_capture); loop { - let &(parent_field_idx, parent_capture) = - parent_captures.peek().expect("we ran out of parent captures!"); + (parent_field_idx, parent_capture) = + *parent_captures.peek().expect("we ran out of parent captures!"); // A parent matches a child they share the same prefix of projections. // The child may have more, if it is capturing sub-fields out of // something that is captured by-move in the parent closure. - if !child_prefix_matches_parent_projections(parent_capture, child_capture) { - // Make sure the field was used at least once. - assert!( - field_used_at_least_once, - "we captured {parent_capture:#?} but it was not used in the child coroutine?" - ); - field_used_at_least_once = false; - // Skip this field. - let _ = parent_captures.next().unwrap(); - continue; + if child_prefix_matches_parent_projections(parent_capture, child_capture) { + break; } - - // Store this set of additional projections (fields and derefs). - // We need to re-apply them later. - let child_precise_captures = - &child_capture.place.projections[parent_capture.place.projections.len()..]; - - // If the parent captures by-move, and the child captures by-ref, then we - // need to peel an additional `deref` off of the body of the child. - let needs_deref = child_capture.is_by_ref() && !parent_capture.is_by_ref(); - if needs_deref { - assert_ne!( - coroutine_kind, - ty::ClosureKind::FnOnce, - "`FnOnce` coroutine-closures return coroutines that capture from \ - their body; it will always result in a borrowck error!" - ); - } - - // Finally, store the type of the parent's captured place. We need - // this when building the field projection in the MIR body later on. - let mut parent_capture_ty = parent_capture.place.ty(); - parent_capture_ty = match parent_capture.info.capture_kind { - ty::UpvarCapture::ByValue => parent_capture_ty, - ty::UpvarCapture::ByRef(kind) => Ty::new_ref( - tcx, - tcx.lifetimes.re_erased, - parent_capture_ty, - kind.to_mutbl_lossy(), - ), - }; - - field_remapping.insert( - FieldIdx::from_usize(child_field_idx + num_args), - ( - FieldIdx::from_usize(parent_field_idx + num_args), - parent_capture_ty, - needs_deref, - child_precise_captures, - ), + // Make sure the field was used at least once. + assert!( + field_used_at_least_once, + "we captured {parent_capture:#?} but it was not used in the child coroutine?" ); - - field_used_at_least_once = true; - break; + field_used_at_least_once = false; + // Skip this field. + let _ = parent_captures.next().unwrap(); } + + // Store this set of additional projections (fields and derefs). + // We need to re-apply them later. + let child_precise_captures = + &child_capture.place.projections[parent_capture.place.projections.len()..]; + + // If the parent captures by-move, and the child captures by-ref, then we + // need to peel an additional `deref` off of the body of the child. + let needs_deref = child_capture.is_by_ref() && !parent_capture.is_by_ref(); + if needs_deref { + assert_ne!( + coroutine_kind, + ty::ClosureKind::FnOnce, + "`FnOnce` coroutine-closures return coroutines that capture from \ + their body; it will always result in a borrowck error!" + ); + } + + // Finally, store the type of the parent's captured place. We need + // this when building the field projection in the MIR body later on. + let mut parent_capture_ty = parent_capture.place.ty(); + parent_capture_ty = match parent_capture.info.capture_kind { + ty::UpvarCapture::ByValue => parent_capture_ty, + ty::UpvarCapture::ByRef(kind) => Ty::new_ref( + tcx, + tcx.lifetimes.re_erased, + parent_capture_ty, + kind.to_mutbl_lossy(), + ), + }; + + field_remapping.insert( + FieldIdx::from_usize(child_field_idx + num_args), + ( + FieldIdx::from_usize(parent_field_idx + num_args), + parent_capture_ty, + needs_deref, + child_precise_captures, + ), + ); + + field_used_at_least_once = true; } // Pop the last parent capture From bf497b8347fc4c7df523c88b5bcf23c6eec9ca16 Mon Sep 17 00:00:00 2001 From: Oli Scherer Date: Tue, 9 Apr 2024 08:59:09 +0000 Subject: [PATCH 4/5] Add a FIXME --- compiler/rustc_mir_transform/src/coroutine/by_move_body.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs index 7979b85225c..35a4d3fff06 100644 --- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs +++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs @@ -233,6 +233,7 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody { let mut by_move_body = body.clone(); MakeByMoveBody { tcx, field_remapping, by_move_coroutine_ty }.visit_body(&mut by_move_body); dump_mir(tcx, false, "coroutine_by_move", &0, &by_move_body, |_, _| Ok(())); + // FIXME: use query feeding to generate the body right here and then only store the `DefId` of the new body. by_move_body.source = mir::MirSource::from_instance(InstanceDef::CoroutineKindShim { coroutine_def_id: coroutine_def_id.to_def_id(), }); From e14e7954aee7e82d427dedca94f185d0a539c7c9 Mon Sep 17 00:00:00 2001 From: Oli Scherer Date: Tue, 9 Apr 2024 09:11:41 +0000 Subject: [PATCH 5/5] Iterate over parent captures first, as there is a 1:N mapping of parent captures to child captures --- .../src/coroutine/by_move_body.rs | 134 ++++++++---------- 1 file changed, 63 insertions(+), 71 deletions(-) diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs index 35a4d3fff06..1d022d159a1 100644 --- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs +++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs @@ -126,92 +126,84 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody { let mut field_remapping = UnordMap::default(); - // One parent capture may correspond to several child captures if we end up - // refining the set of captures via edition-2021 precise captures. We want to - // match up any number of child captures with one parent capture, so we keep - // peeking off this `Peekable` until the child doesn't match anymore. - let mut parent_captures = - tcx.closure_captures(parent_def_id).iter().copied().enumerate().peekable(); - // Make sure we use every field at least once, b/c why are we capturing something - // if it's not used in the inner coroutine. - let mut field_used_at_least_once = false; - - for (child_field_idx, child_capture) in tcx + let mut child_captures = tcx .closure_captures(coroutine_def_id) .iter() .copied() // By construction we capture all the args first. .skip(num_args) .enumerate() + .peekable(); + + // One parent capture may correspond to several child captures if we end up + // refining the set of captures via edition-2021 precise captures. We want to + // match up any number of child captures with one parent capture, so we keep + // peeking off this `Peekable` until the child doesn't match anymore. + for (parent_field_idx, parent_capture) in + tcx.closure_captures(parent_def_id).iter().copied().enumerate() { - let (mut parent_field_idx, mut parent_capture); - loop { - (parent_field_idx, parent_capture) = - *parent_captures.peek().expect("we ran out of parent captures!"); - // A parent matches a child they share the same prefix of projections. - // The child may have more, if it is capturing sub-fields out of - // something that is captured by-move in the parent closure. - if child_prefix_matches_parent_projections(parent_capture, child_capture) { - break; - } - // Make sure the field was used at least once. - assert!( - field_used_at_least_once, - "we captured {parent_capture:#?} but it was not used in the child coroutine?" - ); - field_used_at_least_once = false; - // Skip this field. - let _ = parent_captures.next().unwrap(); - } + // Make sure we use every field at least once, b/c why are we capturing something + // if it's not used in the inner coroutine. + let mut field_used_at_least_once = false; - // Store this set of additional projections (fields and derefs). - // We need to re-apply them later. - let child_precise_captures = - &child_capture.place.projections[parent_capture.place.projections.len()..]; + // A parent matches a child if they share the same prefix of projections. + // The child may have more, if it is capturing sub-fields out of + // something that is captured by-move in the parent closure. + while child_captures.peek().map_or(false, |(_, child_capture)| { + child_prefix_matches_parent_projections(parent_capture, child_capture) + }) { + let (child_field_idx, child_capture) = child_captures.next().unwrap(); - // If the parent captures by-move, and the child captures by-ref, then we - // need to peel an additional `deref` off of the body of the child. - let needs_deref = child_capture.is_by_ref() && !parent_capture.is_by_ref(); - if needs_deref { - assert_ne!( - coroutine_kind, - ty::ClosureKind::FnOnce, - "`FnOnce` coroutine-closures return coroutines that capture from \ + // Store this set of additional projections (fields and derefs). + // We need to re-apply them later. + let child_precise_captures = + &child_capture.place.projections[parent_capture.place.projections.len()..]; + + // If the parent captures by-move, and the child captures by-ref, then we + // need to peel an additional `deref` off of the body of the child. + let needs_deref = child_capture.is_by_ref() && !parent_capture.is_by_ref(); + if needs_deref { + assert_ne!( + coroutine_kind, + ty::ClosureKind::FnOnce, + "`FnOnce` coroutine-closures return coroutines that capture from \ their body; it will always result in a borrowck error!" + ); + } + + // Finally, store the type of the parent's captured place. We need + // this when building the field projection in the MIR body later on. + let mut parent_capture_ty = parent_capture.place.ty(); + parent_capture_ty = match parent_capture.info.capture_kind { + ty::UpvarCapture::ByValue => parent_capture_ty, + ty::UpvarCapture::ByRef(kind) => Ty::new_ref( + tcx, + tcx.lifetimes.re_erased, + parent_capture_ty, + kind.to_mutbl_lossy(), + ), + }; + + field_remapping.insert( + FieldIdx::from_usize(child_field_idx + num_args), + ( + FieldIdx::from_usize(parent_field_idx + num_args), + parent_capture_ty, + needs_deref, + child_precise_captures, + ), ); + + field_used_at_least_once = true; } - // Finally, store the type of the parent's captured place. We need - // this when building the field projection in the MIR body later on. - let mut parent_capture_ty = parent_capture.place.ty(); - parent_capture_ty = match parent_capture.info.capture_kind { - ty::UpvarCapture::ByValue => parent_capture_ty, - ty::UpvarCapture::ByRef(kind) => Ty::new_ref( - tcx, - tcx.lifetimes.re_erased, - parent_capture_ty, - kind.to_mutbl_lossy(), - ), - }; - - field_remapping.insert( - FieldIdx::from_usize(child_field_idx + num_args), - ( - FieldIdx::from_usize(parent_field_idx + num_args), - parent_capture_ty, - needs_deref, - child_precise_captures, - ), + // Make sure the field was used at least once. + assert!( + field_used_at_least_once, + "we captured {parent_capture:#?} but it was not used in the child coroutine?" ); - - field_used_at_least_once = true; } - - // Pop the last parent capture - if field_used_at_least_once { - let _ = parent_captures.next().unwrap(); - } - assert_eq!(parent_captures.next(), None, "leftover parent captures?"); + assert_eq!(child_captures.next(), None, "leftover child captures?"); if coroutine_kind == ty::ClosureKind::FnOnce { assert_eq!(field_remapping.len(), tcx.closure_captures(parent_def_id).len());