From a9194f30eb5b80bd659f3f8e8357fd19b771063e Mon Sep 17 00:00:00 2001 From: Liu Dingming Date: Thu, 4 Jul 2024 04:52:01 +0800 Subject: [PATCH] Use `IndexVec` for coroutine local mapping --- compiler/rustc_index/src/vec.rs | 5 +++ compiler/rustc_mir_transform/src/coroutine.rs | 35 +++++++++---------- 2 files changed, 21 insertions(+), 19 deletions(-) diff --git a/compiler/rustc_index/src/vec.rs b/compiler/rustc_index/src/vec.rs index 88298150a79..b866c8b8433 100644 --- a/compiler/rustc_index/src/vec.rs +++ b/compiler/rustc_index/src/vec.rs @@ -208,6 +208,11 @@ impl IndexVec> { pub fn remove(&mut self, index: I) -> Option { self.get_mut(index)?.take() } + + #[inline] + pub fn contains(&self, index: I) -> bool { + self.get(index).and_then(Option::as_ref).is_some() + } } impl fmt::Debug for IndexVec { diff --git a/compiler/rustc_mir_transform/src/coroutine.rs b/compiler/rustc_mir_transform/src/coroutine.rs index 05674792426..41c9c55cef5 100644 --- a/compiler/rustc_mir_transform/src/coroutine.rs +++ b/compiler/rustc_mir_transform/src/coroutine.rs @@ -58,7 +58,7 @@ use crate::deref_separator::deref_finder; use crate::errors; use crate::pass_manager as pm; use crate::simplify; -use rustc_data_structures::fx::{FxHashMap, FxHashSet}; +use rustc_data_structures::fx::FxHashSet; use rustc_errors::pluralize; use rustc_hir as hir; use rustc_hir::lang_items::LangItem; @@ -236,8 +236,7 @@ struct TransformVisitor<'tcx> { discr_ty: Ty<'tcx>, // Mapping from Local to (type of local, coroutine struct index) - // FIXME(eddyb) This should use `IndexVec>`. - remap: FxHashMap, VariantIdx, FieldIdx)>, + remap: IndexVec, VariantIdx, FieldIdx)>>, // A map from a suspension point in a block to the locals which have live storage at that point storage_liveness: IndexVec>>, @@ -485,7 +484,7 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> { } fn visit_local(&mut self, local: &mut Local, _: PlaceContext, _: Location) { - assert_eq!(self.remap.get(local), None); + assert!(!self.remap.contains(*local)); } fn visit_place( @@ -495,7 +494,7 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> { _location: Location, ) { // Replace an Local in the remap with a coroutine struct access - if let Some(&(ty, variant_index, idx)) = self.remap.get(&place.local) { + if let Some(&Some((ty, variant_index, idx))) = self.remap.get(place.local) { replace_base(place, self.make_field(variant_index, idx, ty), self.tcx); } } @@ -504,7 +503,7 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> { // Remove StorageLive and StorageDead statements for remapped locals data.retain_statements(|s| match s.kind { StatementKind::StorageLive(l) | StatementKind::StorageDead(l) => { - !self.remap.contains_key(&l) + !self.remap.contains(l) } _ => true, }); @@ -529,13 +528,9 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> { // The resume arg target location might itself be remapped if its base local is // live across a yield. - let resume_arg = - if let Some(&(ty, variant, idx)) = self.remap.get(&resume_arg.local) { - replace_base(&mut resume_arg, self.make_field(variant, idx, ty), self.tcx); - resume_arg - } else { - resume_arg - }; + if let Some(&Some((ty, variant, idx))) = self.remap.get(resume_arg.local) { + replace_base(&mut resume_arg, self.make_field(variant, idx, ty), self.tcx); + } let storage_liveness: GrowableBitSet = self.storage_liveness[block].clone().unwrap().into(); @@ -543,7 +538,7 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> { for i in 0..self.always_live_locals.domain_size() { let l = Local::new(i); let needs_storage_dead = storage_liveness.contains(l) - && !self.remap.contains_key(&l) + && !self.remap.contains(l) && !self.always_live_locals.contains(l); if needs_storage_dead { data.statements @@ -1037,7 +1032,7 @@ fn compute_layout<'tcx>( liveness: LivenessInfo, body: &Body<'tcx>, ) -> ( - FxHashMap, VariantIdx, FieldIdx)>, + IndexVec, VariantIdx, FieldIdx)>>, CoroutineLayout<'tcx>, IndexVec>>, ) { @@ -1098,7 +1093,7 @@ fn compute_layout<'tcx>( // Create a map from local indices to coroutine struct indices. let mut variant_fields: IndexVec> = iter::repeat(IndexVec::new()).take(RESERVED_VARIANTS).collect(); - let mut remap = FxHashMap::default(); + let mut remap = IndexVec::from_elem_n(None, saved_locals.domain_size()); for (suspension_point_idx, live_locals) in live_locals_at_suspension_points.iter().enumerate() { let variant_index = VariantIdx::from(RESERVED_VARIANTS + suspension_point_idx); let mut fields = IndexVec::new(); @@ -1109,7 +1104,7 @@ fn compute_layout<'tcx>( // around inside coroutines, so it doesn't matter which variant // index we access them by. let idx = FieldIdx::from_usize(idx); - remap.entry(locals[saved_local]).or_insert((tys[saved_local].ty, variant_index, idx)); + remap[locals[saved_local]] = Some((tys[saved_local].ty, variant_index, idx)); } variant_fields.push(fields); variant_source_info.push(source_info_at_suspension_points[suspension_point_idx]); @@ -1121,7 +1116,9 @@ fn compute_layout<'tcx>( for var in &body.var_debug_info { let VarDebugInfoContents::Place(place) = &var.value else { continue }; let Some(local) = place.as_local() else { continue }; - let Some(&(_, variant, field)) = remap.get(&local) else { continue }; + let Some(&Some((_, variant, field))) = remap.get(local) else { + continue; + }; let saved_local = variant_fields[variant][field]; field_names.get_or_insert_with(saved_local, || var.name); @@ -1524,7 +1521,7 @@ fn create_cases<'tcx>( for i in 0..(body.local_decls.len()) { let l = Local::new(i); let needs_storage_live = point.storage_liveness.contains(l) - && !transform.remap.contains_key(&l) + && !transform.remap.contains(l) && !transform.always_live_locals.contains(l); if needs_storage_live { statements