optimization round 2

- moved work from `find_local` to `gather_statement`
- created custom iterator for `iter_projections`
- reverted change from `IndexVec` to `FxIndexMap`

3 files changed, 96 insertions, 34 deletions

diff --git a/compiler/rustc_mir_dataflow/src/move_paths/builder.rs b/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
index 537df1588a3..5052de99184 100644
--- a/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
+++ b/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
@@ -39,18 +39,18 @@ 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(),
@@ -285,8 +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!(self.builder.body.local_decls[place.local].is_deref_temp());
-                self.builder.data.rev_lookup.un_derefer.insert(place.as_local().unwrap(), *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);
diff --git a/compiler/rustc_mir_dataflow/src/move_paths/mod.rs b/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
index 0f7da4f46e3..0c7aa6676ec 100644
--- a/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
+++ b/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
@@ -1,6 +1,6 @@
 use crate::move_paths::builder::MoveDat;
 use crate::un_derefer::UnDerefer;
-use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
+use rustc_data_structures::fx::FxHashMap;
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::mir::*;
 use rustc_middle::ty::{ParamEnv, Ty, TyCtxt};
@@ -291,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
@@ -331,17 +331,9 @@ impl<'tcx> MovePathLookup<'tcx> {
         LookupResult::Exact(result)
     }
 
+    #[inline]
     pub fn find_local(&self, local: Local) -> MovePathIndex {
-        let deref_chain = self.un_derefer.deref_chain(local);
-
-        let local = match deref_chain.first() {
-            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
@@ -349,7 +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))
+        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
index fdd2eb0233d..874d50ffd0e 100644
--- a/compiler/rustc_mir_dataflow/src/un_derefer.rs
+++ b/compiler/rustc_mir_dataflow/src/un_derefer.rs
@@ -1,36 +1,100 @@
 use rustc_data_structures::fx::FxHashMap;
 use rustc_middle::mir::*;
 
-use std::iter::once;
-
 /// Used for reverting changes made by `DerefSeparator`
 #[derive(Default, Debug)]
 pub struct UnDerefer<'tcx> {
-    deref_chains: FxHashMap<Local, Vec<Place<'tcx>>>,
+    deref_chains: FxHashMap<Local, Vec<PlaceRef<'tcx>>>,
 }
 
 impl<'tcx> UnDerefer<'tcx> {
-    pub fn insert(&mut self, local: Local, reffed: Place<'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
-    pub fn deref_chain(&self, local: Local) -> &[Place<'tcx>] {
+    #[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>)> + DoubleEndedIterator + '_ {
-        let deref_chain = self.deref_chain(place.local);
-
-        deref_chain
-            .iter()
-            .map(Place::as_ref)
-            .chain(once(place))
-            .flat_map(|place| place.iter_projections())
+    ) -> 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,
+        }
     }
 }