Merge `un_derefer` into `MovePathLookup`

5 files changed, 129 insertions, 115 deletions

diff --git a/compiler/rustc_mir_dataflow/src/lib.rs b/compiler/rustc_mir_dataflow/src/lib.rs
index 900d438f8d5..d43446bc5b2 100644
--- a/compiler/rustc_mir_dataflow/src/lib.rs
+++ b/compiler/rustc_mir_dataflow/src/lib.rs
@@ -43,7 +43,6 @@ 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 2b37b55d278..dc7e9ab3ce6 100644
--- a/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
+++ b/compiler/rustc_mir_dataflow/src/move_paths/builder.rs
@@ -1,11 +1,10 @@
-use crate::move_paths::FxHashMap;
-use crate::un_derefer::UnDerefer;
 use rustc_index::IndexVec;
 use rustc_middle::mir::tcx::RvalueInitializationState;
 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;
@@ -21,7 +20,6 @@ struct MoveDataBuilder<'a, 'tcx> {
     param_env: ty::ParamEnv<'tcx>,
     data: MoveData<'tcx>,
     errors: Vec<(Place<'tcx>, MoveError<'tcx>)>,
-    un_derefer: UnDerefer<'tcx>,
 }
 
 impl<'a, 'tcx> MoveDataBuilder<'a, 'tcx> {
@@ -35,25 +33,29 @@ impl<'a, 'tcx> MoveDataBuilder<'a, 'tcx> {
             tcx,
             param_env,
             errors: Vec::new(),
-            un_derefer: UnDerefer { tcx: tcx, derefer_sidetable: Default::default() },
             data: MoveData {
                 moves: IndexVec::new(),
                 loc_map: LocationMap::new(body),
                 rev_lookup: MovePathLookup {
                     locals: body
                         .local_decls
-                        .indices()
-                        .map(|i| {
-                            Self::new_move_path(
-                                &mut move_paths,
-                                &mut path_map,
-                                &mut init_path_map,
-                                None,
-                                Place::from(i),
+                        .iter_enumerated()
+                        .filter(|(_, l)| !l.is_deref_temp())
+                        .map(|(i, _)| {
+                            (
+                                i,
+                                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(),
                 },
                 move_paths,
                 path_map,
@@ -98,13 +100,11 @@ 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>> {
-        if let Some(new_place) = self.builder.un_derefer.derefer(place.as_ref(), self.builder.body)
-        {
-            return self.move_path_for(new_place);
-        }
+        let deref_chain = self.builder.data.rev_lookup.deref_chain(place.as_ref());
 
         debug!("lookup({:?})", place);
-        let mut base = self.builder.data.rev_lookup.locals[place.local];
+        let mut base =
+            self.builder.data.rev_lookup.find_local(deref_chain.first().unwrap_or(&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,51 +113,55 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
         // from `*(u.f: &_)` isn't allowed.
         let mut union_path = None;
 
-        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(_) => {
-                    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 {
+        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,
-                            InteriorOfSliceOrArray { ty: place_ty, is_index: true },
+                            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: 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() {
+                    base = self
+                        .add_move_path(base, elem, |tcx| place_ref.project_deeper(&[elem], tcx));
+                }
             }
         }
 
@@ -198,10 +202,8 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
     }
 }
 
-pub type MoveDat<'tcx> = Result<
-    (FxHashMap<Local, Place<'tcx>>, MoveData<'tcx>),
-    (MoveData<'tcx>, Vec<(Place<'tcx>, MoveError<'tcx>)>),
->;
+pub type MoveDat<'tcx> =
+    Result<MoveData<'tcx>, (MoveData<'tcx>, Vec<(Place<'tcx>, MoveError<'tcx>)>)>;
 
 impl<'a, 'tcx> MoveDataBuilder<'a, 'tcx> {
     fn finalize(self) -> MoveDat<'tcx> {
@@ -217,11 +219,7 @@ impl<'a, 'tcx> MoveDataBuilder<'a, 'tcx> {
             "done dumping moves"
         });
 
-        if self.errors.is_empty() {
-            Ok((self.un_derefer.derefer_sidetable, self.data))
-        } else {
-            Err((self.data, self.errors))
-        }
+        if self.errors.is_empty() { Ok(self.data) } else { Err((self.data, self.errors)) }
     }
 }
 
@@ -250,7 +248,7 @@ pub(super) fn gather_moves<'tcx>(
 impl<'a, 'tcx> MoveDataBuilder<'a, 'tcx> {
     fn gather_args(&mut self) {
         for arg in self.body.args_iter() {
-            let path = self.data.rev_lookup.locals[arg];
+            let path = self.data.rev_lookup.find_local(arg);
 
             let init = self.data.inits.push(Init {
                 path,
@@ -286,7 +284,7 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
             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.un_derefer.derefer_sidetable.insert(place.local, *reffed);
+                    self.builder.data.rev_lookup.derefer_sidetable.insert(place.local, *reffed);
                 }
             }
             StatementKind::Assign(box (place, rval)) => {
@@ -308,7 +306,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.un_derefer.derefer_sidetable.contains_key(&local) {
+                if !self.builder.data.rev_lookup.derefer_sidetable.contains_key(&local) {
                     self.gather_move(Place::from(*local));
                 }
             }
@@ -450,12 +448,6 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
 
     fn gather_move(&mut self, place: Place<'tcx>) {
         debug!("gather_move({:?}, {:?})", self.loc, place);
-        if let Some(new_place) = self.builder.un_derefer.derefer(place.as_ref(), self.builder.body)
-        {
-            self.gather_move(new_place);
-            return;
-        }
-
         if let [ref base @ .., ProjectionElem::Subslice { from, to, from_end: false }] =
             **place.projection
         {
@@ -512,11 +504,6 @@ impl<'b, 'a, 'tcx> Gatherer<'b, 'a, 'tcx> {
     fn gather_init(&mut self, place: PlaceRef<'tcx>, kind: InitKind) {
         debug!("gather_init({:?}, {:?})", self.loc, place);
 
-        if let Some(new_place) = self.builder.un_derefer.derefer(place, self.builder.body) {
-            self.gather_init(new_place.as_ref(), kind);
-            return;
-        }
-
         let mut place = place;
 
         // Check if we are assigning into a field of a union, if so, lookup the place
diff --git a/compiler/rustc_mir_dataflow/src/move_paths/mod.rs b/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
index ab1a6715361..aa901f66d3f 100644
--- a/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
+++ b/compiler/rustc_mir_dataflow/src/move_paths/mod.rs
@@ -1,5 +1,5 @@
 use crate::move_paths::builder::MoveDat;
-use rustc_data_structures::fx::FxHashMap;
+use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
 use rustc_index::{IndexSlice, IndexVec};
 use rustc_middle::mir::*;
 use rustc_middle::ty::{ParamEnv, Ty, TyCtxt};
@@ -175,7 +175,7 @@ pub struct MoveData<'tcx> {
     /// particular path being moved.)
     pub loc_map: LocationMap<SmallVec<[MoveOutIndex; 4]>>,
     pub path_map: IndexVec<MovePathIndex, SmallVec<[MoveOutIndex; 4]>>,
-    pub rev_lookup: MovePathLookup,
+    pub rev_lookup: MovePathLookup<'tcx>,
     pub inits: IndexVec<InitIndex, Init>,
     /// Each Location `l` is mapped to the Inits that are effects
     /// of executing the code at `l`.
@@ -289,8 +289,8 @@ impl Init {
 
 /// Tables mapping from a place to its MovePathIndex.
 #[derive(Debug)]
-pub struct MovePathLookup {
-    locals: IndexVec<Local, MovePathIndex>,
+pub struct MovePathLookup<'tcx> {
+    locals: FxIndexMap<Local, MovePathIndex>,
 
     /// projections are made from a base-place and a projection
     /// elem. The base-place will have a unique MovePathIndex; we use
@@ -299,6 +299,9 @@ pub struct MovePathLookup {
     /// base-place). For the remaining lookup, we map the projection
     /// 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>>,
 }
 
 mod builder;
@@ -309,27 +312,59 @@ pub enum LookupResult {
     Parent(Option<MovePathIndex>),
 }
 
-impl MovePathLookup {
+impl<'tcx> MovePathLookup<'tcx> {
     // Unlike the builder `fn move_path_for` below, this lookup
     // 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 mut result = self.locals[place.local];
+        let deref_chain = self.deref_chain(place);
 
-        for elem in place.projection.iter() {
-            if let Some(&subpath) = self.projections.get(&(result, elem.lift())) {
-                result = subpath;
-            } else {
+        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()) {
                 return LookupResult::Parent(Some(result));
             }
         }
 
+        if let Some(result) = subpaths_for_place(place) {
+            return LookupResult::Parent(Some(result));
+        }
+
         LookupResult::Exact(result)
     }
 
     pub fn find_local(&self, local: Local) -> MovePathIndex {
-        self.locals[local]
+        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")
+        })
     }
 
     /// An enumerated iterator of `local`s and their associated
@@ -337,7 +372,22 @@ impl MovePathLookup {
     pub fn iter_locals_enumerated(
         &self,
     ) -> impl DoubleEndedIterator<Item = (Local, MovePathIndex)> + ExactSizeIterator + '_ {
-        self.locals.iter_enumerated().map(|(l, &idx)| (l, idx))
+        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
     }
 }
 
diff --git a/compiler/rustc_mir_dataflow/src/rustc_peek.rs b/compiler/rustc_mir_dataflow/src/rustc_peek.rs
index 0cbc7442cf1..156231c3ae1 100644
--- a/compiler/rustc_mir_dataflow/src/rustc_peek.rs
+++ b/compiler/rustc_mir_dataflow/src/rustc_peek.rs
@@ -34,7 +34,7 @@ impl<'tcx> MirPass<'tcx> for SanityCheck {
         }
 
         let param_env = tcx.param_env(def_id);
-        let (_, move_data) = MoveData::gather_moves(body, tcx, param_env).unwrap();
+        let move_data = MoveData::gather_moves(body, tcx, param_env).unwrap();
         let mdpe = MoveDataParamEnv { move_data, param_env };
 
         if has_rustc_mir_with(tcx, def_id, sym::rustc_peek_maybe_init).is_some() {
diff --git a/compiler/rustc_mir_dataflow/src/un_derefer.rs b/compiler/rustc_mir_dataflow/src/un_derefer.rs
deleted file mode 100644
index 7e6e25cc603..00000000000
--- a/compiler/rustc_mir_dataflow/src/un_derefer.rs
+++ /dev/null
@@ -1,22 +0,0 @@
-use rustc_data_structures::fx::FxHashMap;
-use rustc_middle::mir::*;
-use rustc_middle::ty::TyCtxt;
-
-/// Used for reverting changes made by `DerefSeparator`
-pub struct UnDerefer<'tcx> {
-    pub tcx: TyCtxt<'tcx>,
-    pub derefer_sidetable: FxHashMap<Local, Place<'tcx>>,
-}
-
-impl<'tcx> UnDerefer<'tcx> {
-    #[inline]
-    pub fn derefer(&self, place: PlaceRef<'tcx>, body: &Body<'tcx>) -> Option<Place<'tcx>> {
-        let reffed = self.derefer_sidetable.get(&place.local)?;
-
-        let new_place = reffed.project_deeper(place.projection, self.tcx);
-        if body.local_decls[new_place.local].is_deref_temp() {
-            return self.derefer(new_place.as_ref(), body);
-        }
-        Some(new_place)
-    }
-}