Auto merge of #113316 - DrMeepster:underefer_perf, r=oli-obk

Rewrite `UnDerefer`, again

This PR is intended to improve the perf regression introduced by #112882.

`UnDerefer` has been separated out again for borrowck reasons. It was a bit overzealous to remove it in the previous PR.

r? `@oli-obk`

6 files changed, 186 insertions, 124 deletions

diff --git a/compiler/rustc_borrowck/src/type_check/liveness/polonius.rs b/compiler/rustc_borrowck/src/type_check/liveness/polonius.rs
index b344ab46adb..012075d7114 100644
--- a/compiler/rustc_borrowck/src/type_check/liveness/polonius.rs
+++ b/compiler/rustc_borrowck/src/type_check/liveness/polonius.rs
@@ -20,7 +20,7 @@ struct UseFactsExtractor<'me, 'tcx> {
 }
 
 // A Visitor to walk through the MIR and extract point-wise facts
-impl UseFactsExtractor<'_, '_> {
+impl<'tcx> UseFactsExtractor<'_, 'tcx> {
     fn location_to_index(&self, location: Location) -> LocationIndex {
         self.location_table.mid_index(location)
     }
@@ -45,7 +45,7 @@ impl UseFactsExtractor<'_, '_> {
         self.path_accessed_at_base.push((path, self.location_to_index(location)));
     }
 
-    fn place_to_mpi(&self, place: &Place<'_>) -> Option<MovePathIndex> {
+    fn place_to_mpi(&self, place: &Place<'tcx>) -> Option<MovePathIndex> {
         match self.move_data.rev_lookup.find(place.as_ref()) {
             LookupResult::Exact(mpi) => Some(mpi),
             LookupResult::Parent(mmpi) => mmpi,
diff --git a/compiler/rustc_mir_dataflow/src/impls/mod.rs b/compiler/rustc_mir_dataflow/src/impls/mod.rs
index 98cec1c6760..cb74bea724a 100644
--- a/compiler/rustc_mir_dataflow/src/impls/mod.rs
+++ b/compiler/rustc_mir_dataflow/src/impls/mod.rs
@@ -731,11 +731,11 @@ fn switch_on_enum_discriminant<'mir, 'tcx>(
 
 struct OnMutBorrow<F>(F);
 
-impl<F> Visitor<'_> for OnMutBorrow<F>
+impl<'tcx, F> Visitor<'tcx> for OnMutBorrow<F>
 where
-    F: FnMut(&mir::Place<'_>),
+    F: FnMut(&mir::Place<'tcx>),
 {
-    fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'_>, location: Location) {
+    fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'tcx>, location: Location) {
         // FIXME: Does `&raw const foo` allow mutation? See #90413.
         match rvalue {
             mir::Rvalue::Ref(_, mir::BorrowKind::Mut { .. }, place)
@@ -756,7 +756,7 @@ where
 fn for_each_mut_borrow<'tcx>(
     mir: &impl MirVisitable<'tcx>,
     location: Location,
-    f: impl FnMut(&mir::Place<'_>),
+    f: impl FnMut(&mir::Place<'tcx>),
 ) {
     let mut vis = OnMutBorrow(f);
 
diff --git a/compiler/rustc_mir_dataflow/src/lib.rs b/compiler/rustc_mir_dataflow/src/lib.rs
index d43446bc5b2..900d438f8d5 100644
--- a/compiler/rustc_mir_dataflow/src/lib.rs
+++ b/compiler/rustc_mir_dataflow/src/lib.rs
@@ -43,6 +43,7 @@ pub mod impls;
 pub mod move_paths;
 pub mod rustc_peek;
 pub mod storage;
+pub mod un_derefer;
 pub mod value_analysis;
 
 fluent_messages! { "../messages.ftl" }
diff --git a/compiler/rustc_mir_dataflow/src/move_paths/builder.rs b/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
index dc7e9ab3ce6..5052de99184 100644
--- a/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
+++ b/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
@@ -4,7 +4,6 @@ use rustc_middle::mir::*;
 use rustc_middle::ty::{self, TyCtxt};
 use smallvec::{smallvec, SmallVec};
 
-use std::iter;
 use std::mem;
 
 use super::abs_domain::Lift;
@@ -40,22 +39,22 @@ impl<'a, 'tcx> MoveDataBuilder<'a, 'tcx> {
                     locals: body
                         .local_decls
                         .iter_enumerated()
-                        .filter(|(_, l)| !l.is_deref_temp())
-                        .map(|(i, _)| {
-                            (
-                                i,
+                        .map(|(i, l)| {
+                            if l.is_deref_temp() {
+                                MovePathIndex::MAX
+                            } else {
                                 Self::new_move_path(
                                     &mut move_paths,
                                     &mut path_map,
                                     &mut init_path_map,
                                     None,
                                     Place::from(i),
-                                ),
-                            )
+                                )
+                            }
                         })
                         .collect(),
                     projections: Default::default(),
-                    derefer_sidetable: Default::default(),
+                    un_derefer: Default::default(),
                 },
                 move_paths,
                 path_map,
@@ -100,11 +99,10 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
     ///
     /// Maybe we should have separate "borrowck" and "moveck" modes.
     fn move_path_for(&mut self, place: Place<'tcx>) -> Result<MovePathIndex, MoveError<'tcx>> {
-        let deref_chain = self.builder.data.rev_lookup.deref_chain(place.as_ref());
+        let data = &mut self.builder.data;
 
         debug!("lookup({:?})", place);
-        let mut base =
-            self.builder.data.rev_lookup.find_local(deref_chain.first().unwrap_or(&place).local);
+        let mut base = data.rev_lookup.find_local(place.local);
 
         // The move path index of the first union that we find. Once this is
         // some we stop creating child move paths, since moves from unions
@@ -113,55 +111,60 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
         // from `*(u.f: &_)` isn't allowed.
         let mut union_path = None;
 
-        for place in deref_chain.into_iter().chain(iter::once(place)) {
-            for (place_ref, elem) in place.as_ref().iter_projections() {
-                let body = self.builder.body;
-                let tcx = self.builder.tcx;
-                let place_ty = place_ref.ty(body, tcx).ty;
-                match place_ty.kind() {
-                    ty::Ref(..) | ty::RawPtr(..) => {
-                        return Err(MoveError::cannot_move_out_of(
-                            self.loc,
-                            BorrowedContent {
-                                target_place: place_ref.project_deeper(&[elem], tcx),
-                            },
-                        ));
-                    }
-                    ty::Adt(adt, _) if adt.has_dtor(tcx) && !adt.is_box() => {
-                        return Err(MoveError::cannot_move_out_of(
-                            self.loc,
-                            InteriorOfTypeWithDestructor { container_ty: place_ty },
-                        ));
-                    }
-                    ty::Adt(adt, _) if adt.is_union() => {
-                        union_path.get_or_insert(base);
-                    }
-                    ty::Slice(_) => {
+        for (place_ref, elem) in data.rev_lookup.un_derefer.iter_projections(place.as_ref()) {
+            let body = self.builder.body;
+            let tcx = self.builder.tcx;
+            let place_ty = place_ref.ty(body, tcx).ty;
+            match place_ty.kind() {
+                ty::Ref(..) | ty::RawPtr(..) => {
+                    return Err(MoveError::cannot_move_out_of(
+                        self.loc,
+                        BorrowedContent { target_place: place_ref.project_deeper(&[elem], tcx) },
+                    ));
+                }
+                ty::Adt(adt, _) if adt.has_dtor(tcx) && !adt.is_box() => {
+                    return Err(MoveError::cannot_move_out_of(
+                        self.loc,
+                        InteriorOfTypeWithDestructor { container_ty: place_ty },
+                    ));
+                }
+                ty::Adt(adt, _) if adt.is_union() => {
+                    union_path.get_or_insert(base);
+                }
+                ty::Slice(_) => {
+                    return Err(MoveError::cannot_move_out_of(
+                        self.loc,
+                        InteriorOfSliceOrArray {
+                            ty: place_ty,
+                            is_index: matches!(elem, ProjectionElem::Index(..)),
+                        },
+                    ));
+                }
+
+                ty::Array(..) => {
+                    if let ProjectionElem::Index(..) = elem {
                         return Err(MoveError::cannot_move_out_of(
                             self.loc,
-                            InteriorOfSliceOrArray {
-                                ty: place_ty,
-                                is_index: matches!(elem, ProjectionElem::Index(..)),
-                            },
+                            InteriorOfSliceOrArray { ty: place_ty, is_index: true },
                         ));
                     }
+                }
 
-                    ty::Array(..) => {
-                        if let ProjectionElem::Index(..) = elem {
-                            return Err(MoveError::cannot_move_out_of(
-                                self.loc,
-                                InteriorOfSliceOrArray { ty: place_ty, is_index: true },
-                            ));
-                        }
-                    }
-
-                    _ => {}
-                };
+                _ => {}
+            };
 
-                if union_path.is_none() {
-                    base = self
-                        .add_move_path(base, elem, |tcx| place_ref.project_deeper(&[elem], tcx));
-                }
+            if union_path.is_none() {
+                // inlined from add_move_path because of a borrowck conflict with the iterator
+                base =
+                    *data.rev_lookup.projections.entry((base, elem.lift())).or_insert_with(|| {
+                        MoveDataBuilder::new_move_path(
+                            &mut data.move_paths,
+                            &mut data.path_map,
+                            &mut data.init_path_map,
+                            Some(base),
+                            place_ref.project_deeper(&[elem], tcx),
+                        )
+                    })
             }
         }
 
@@ -282,10 +285,14 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
     fn gather_statement(&mut self, stmt: &Statement<'tcx>) {
         match &stmt.kind {
             StatementKind::Assign(box (place, Rvalue::CopyForDeref(reffed))) => {
-                assert!(place.projection.is_empty());
-                if self.builder.body.local_decls[place.local].is_deref_temp() {
-                    self.builder.data.rev_lookup.derefer_sidetable.insert(place.local, *reffed);
-                }
+                let local = place.as_local().unwrap();
+                assert!(self.builder.body.local_decls[local].is_deref_temp());
+
+                let rev_lookup = &mut self.builder.data.rev_lookup;
+
+                rev_lookup.un_derefer.insert(local, reffed.as_ref());
+                let base_local = rev_lookup.un_derefer.deref_chain(local).first().unwrap().local;
+                rev_lookup.locals[local] = rev_lookup.locals[base_local];
             }
             StatementKind::Assign(box (place, rval)) => {
                 self.create_move_path(*place);
@@ -306,7 +313,7 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
             StatementKind::StorageLive(_) => {}
             StatementKind::StorageDead(local) => {
                 // DerefTemp locals (results of CopyForDeref) don't actually move anything.
-                if !self.builder.data.rev_lookup.derefer_sidetable.contains_key(&local) {
+                if !self.builder.body.local_decls[*local].is_deref_temp() {
                     self.gather_move(Place::from(*local));
                 }
             }
diff --git a/compiler/rustc_mir_dataflow/src/move_paths/mod.rs b/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
index aa901f66d3f..0c7aa6676ec 100644
--- a/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
+++ b/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
@@ -1,5 +1,6 @@
 use crate::move_paths::builder::MoveDat;
-use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
+use crate::un_derefer::UnDerefer;
+use rustc_data_structures::fx::FxHashMap;
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::mir::*;
 use rustc_middle::ty::{ParamEnv, Ty, TyCtxt};
@@ -290,7 +291,7 @@ impl Init {
 /// Tables mapping from a place to its MovePathIndex.
 #[derive(Debug)]
 pub struct MovePathLookup<'tcx> {
-    locals: FxIndexMap<Local, MovePathIndex>,
+    locals: IndexVec<Local, MovePathIndex>,
 
     /// projections are made from a base-place and a projection
     /// elem. The base-place will have a unique MovePathIndex; we use
@@ -300,8 +301,7 @@ pub struct MovePathLookup<'tcx> {
     /// elem to the associated MovePathIndex.
     projections: FxHashMap<(MovePathIndex, AbstractElem), MovePathIndex>,
 
-    /// Maps `DerefTemp` locals to the `Place`s assigned to them.
-    derefer_sidetable: FxHashMap<Local, Place<'tcx>>,
+    un_derefer: UnDerefer<'tcx>,
 }
 
 mod builder;
@@ -317,54 +317,23 @@ impl<'tcx> MovePathLookup<'tcx> {
     // alternative will *not* create a MovePath on the fly for an
     // unknown place, but will rather return the nearest available
     // parent.
-    pub fn find(&self, place: PlaceRef<'_>) -> LookupResult {
-        let deref_chain = self.deref_chain(place);
+    pub fn find(&self, place: PlaceRef<'tcx>) -> LookupResult {
+        let mut result = self.find_local(place.local);
 
-        let local = match deref_chain.first() {
-            Some(place) => place.local,
-            None => place.local,
-        };
-
-        let mut result = *self.locals.get(&local).unwrap_or_else(|| {
-            bug!("base local ({local:?}) of deref_chain should not be a deref temp")
-        });
-
-        // this needs to be a closure because `place` has a different lifetime than `prefix`'s places
-        let mut subpaths_for_place = |place: PlaceRef<'_>| {
-            for elem in place.projection.iter() {
-                if let Some(&subpath) = self.projections.get(&(result, elem.lift())) {
-                    result = subpath;
-                } else {
-                    return Some(result);
-                }
-            }
-            None
-        };
-
-        for place in deref_chain {
-            if let Some(result) = subpaths_for_place(place.as_ref()) {
+        for (_, elem) in self.un_derefer.iter_projections(place) {
+            if let Some(&subpath) = self.projections.get(&(result, elem.lift())) {
+                result = subpath;
+            } else {
                 return LookupResult::Parent(Some(result));
             }
         }
 
-        if let Some(result) = subpaths_for_place(place) {
-            return LookupResult::Parent(Some(result));
-        }
-
         LookupResult::Exact(result)
     }
 
+    #[inline]
     pub fn find_local(&self, local: Local) -> MovePathIndex {
-        let deref_chain = self.deref_chain(Place::from(local).as_ref());
-
-        let local = match deref_chain.last() {
-            Some(place) => place.local,
-            None => local,
-        };
-
-        *self.locals.get(&local).unwrap_or_else(|| {
-            bug!("base local ({local:?}) of deref_chain should not be a deref temp")
-        })
+        self.locals[local]
     }
 
     /// An enumerated iterator of `local`s and their associated
@@ -372,22 +341,7 @@ impl<'tcx> MovePathLookup<'tcx> {
     pub fn iter_locals_enumerated(
         &self,
     ) -> impl DoubleEndedIterator<Item = (Local, MovePathIndex)> + ExactSizeIterator + '_ {
-        self.locals.iter().map(|(&l, &idx)| (l, idx))
-    }
-
-    /// Returns the chain of places behind `DerefTemp` locals in `place`
-    pub fn deref_chain(&self, place: PlaceRef<'_>) -> Vec<Place<'tcx>> {
-        let mut prefix = Vec::new();
-        let mut local = place.local;
-
-        while let Some(&reffed) = self.derefer_sidetable.get(&local) {
-            prefix.insert(0, reffed);
-            local = reffed.local;
-        }
-
-        debug!("deref_chain({place:?}) = {prefix:?}");
-
-        prefix
+        self.locals.iter_enumerated().map(|(l, &idx)| (l, idx))
     }
 }
 
diff --git a/compiler/rustc_mir_dataflow/src/un_derefer.rs b/compiler/rustc_mir_dataflow/src/un_derefer.rs
new file mode 100644
index 00000000000..874d50ffd0e
--- /dev/null
+++ b/compiler/rustc_mir_dataflow/src/un_derefer.rs
@@ -0,0 +1,100 @@
+use rustc_data_structures::fx::FxHashMap;
+use rustc_middle::mir::*;
+
+/// Used for reverting changes made by `DerefSeparator`
+#[derive(Default, Debug)]
+pub struct UnDerefer<'tcx> {
+    deref_chains: FxHashMap<Local, Vec<PlaceRef<'tcx>>>,
+}
+
+impl<'tcx> UnDerefer<'tcx> {
+    #[inline]
+    pub fn insert(&mut self, local: Local, reffed: PlaceRef<'tcx>) {
+        let mut chain = self.deref_chains.remove(&reffed.local).unwrap_or_default();
+        chain.push(reffed);
+        self.deref_chains.insert(local, chain);
+    }
+
+    /// Returns the chain of places behind `DerefTemp` locals
+    #[inline]
+    pub fn deref_chain(&self, local: Local) -> &[PlaceRef<'tcx>] {
+        self.deref_chains.get(&local).map(Vec::as_slice).unwrap_or_default()
+    }
+
+    /// Iterates over the projections of a place and its deref chain.
+    ///
+    /// See [`PlaceRef::iter_projections`]
+    #[inline]
+    pub fn iter_projections(
+        &self,
+        place: PlaceRef<'tcx>,
+    ) -> impl Iterator<Item = (PlaceRef<'tcx>, PlaceElem<'tcx>)> + '_ {
+        ProjectionIter::new(self.deref_chain(place.local), place)
+    }
+}
+
+/// The iterator returned by [`UnDerefer::iter_projections`].
+struct ProjectionIter<'a, 'tcx> {
+    places: SlicePlusOne<'a, PlaceRef<'tcx>>,
+    proj_idx: usize,
+}
+
+impl<'a, 'tcx> ProjectionIter<'a, 'tcx> {
+    #[inline]
+    fn new(deref_chain: &'a [PlaceRef<'tcx>], place: PlaceRef<'tcx>) -> Self {
+        // just return an empty iterator for a bare local
+        let last = if place.as_local().is_none() {
+            Some(place)
+        } else {
+            debug_assert!(deref_chain.is_empty());
+            None
+        };
+
+        ProjectionIter { places: SlicePlusOne { slice: deref_chain, last }, proj_idx: 0 }
+    }
+}
+
+impl<'tcx> Iterator for ProjectionIter<'_, 'tcx> {
+    type Item = (PlaceRef<'tcx>, PlaceElem<'tcx>);
+
+    #[inline]
+    fn next(&mut self) -> Option<(PlaceRef<'tcx>, PlaceElem<'tcx>)> {
+        let place = self.places.read()?;
+
+        // the projection should never be empty except for a bare local which is handled in new
+        let partial_place =
+            PlaceRef { local: place.local, projection: &place.projection[..self.proj_idx] };
+        let elem = place.projection[self.proj_idx];
+
+        if self.proj_idx == place.projection.len() - 1 {
+            self.proj_idx = 0;
+            self.places.advance();
+        } else {
+            self.proj_idx += 1;
+        }
+
+        Some((partial_place, elem))
+    }
+}
+
+struct SlicePlusOne<'a, T> {
+    slice: &'a [T],
+    last: Option<T>,
+}
+
+impl<T: Copy> SlicePlusOne<'_, T> {
+    #[inline]
+    fn read(&self) -> Option<T> {
+        self.slice.first().copied().or(self.last)
+    }
+
+    #[inline]
+    fn advance(&mut self) {
+        match self.slice {
+            [_, ref remainder @ ..] => {
+                self.slice = remainder;
+            }
+            [] => self.last = None,
+        }
+    }
+}