simplify MoveDataBuilder::move_path_for with new UnDerefer::co_iter_projections

compiler/rustc_mir_dataflow/src/move_paths/builder.rs

@@ -108,10 +108,8 @@ impl<'a, 'tcx, F: Fn(Ty<'tcx>) -> bool> MoveDataBuilder<'a, 'tcx, F> {
     where
         G: FnMut(&mut Self, MovePathIndex),
     {
-        let data = &mut self.data;
-
         debug!("lookup({:?})", place);
-        let Some(mut base) = data.rev_lookup.find_local(place.local) else {
+        let Some(mut base) = self.data.rev_lookup.find_local(place.local) else {
             return;
         };
 
@@ -122,8 +120,8 @@ impl<'a, 'tcx, F: Fn(Ty<'tcx>) -> bool> MoveDataBuilder<'a, 'tcx, F> {
         // from `*(u.f: &_)` isn't allowed.
         let mut union_path = None;
 
-        let mut iter = data.rev_lookup.un_derefer.iter_projections(place.as_ref());
-        while let Some((place_ref, elem)) = iter.next() {
+        let mut iter = self.data.rev_lookup.un_derefer.co_iter_projections(place.as_ref());
+        while let Some((place_ref, elem)) = iter.next(&self.data.rev_lookup.un_derefer) {
             let body = self.body;
             let tcx = self.tcx;
             let place_ty = place_ref.ty(body, tcx).ty;
@@ -232,12 +230,9 @@ impl<'a, 'tcx, F: Fn(Ty<'tcx>) -> bool> MoveDataBuilder<'a, 'tcx, F> {
                 // `ConstIndex` patterns. This is done to ensure that all move paths
                 // are disjoint, which is expected by drop elaboration.
                 MoveSubPathResult::Subslice { from, to } => {
-                    assert!(
-                        iter.all(
-                            |(_, elem)| MoveSubPath::of(elem.kind()) == MoveSubPathResult::Skip
-                        )
-                    );
-                    drop(iter); // drop for borrowck
+                    assert!(iter.into_iter(&self.data.rev_lookup.un_derefer).all(|(_, elem)| {
+                        MoveSubPath::of(elem.kind()) == MoveSubPathResult::Skip
+                    }));
 
                     let (&elem_ty, len) = match place_ty.kind() {
                         ty::Array(ty, size) => (
@@ -275,21 +270,11 @@ impl<'a, 'tcx, F: Fn(Ty<'tcx>) -> bool> MoveDataBuilder<'a, 'tcx, F> {
                 return;
             }
             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, move_elem)).or_insert_with(|| {
-                    new_move_path(
-                        &mut data.move_paths,
-                        &mut data.move_out_path_map,
-                        &mut data.init_path_map,
-                        Some(base),
-                        place_ref.project_deeper(&[elem], tcx),
-                    )
-                })
+                base = self
+                    .add_move_path(base, move_elem, |tcx| place_ref.project_deeper(&[elem], tcx))
             }
         }
 
-        drop(iter); // drop for borrowck
-
         if let Some(base) = union_path {
             // Move out of union - always move the entire union.
             on_move(self, base);

compiler/rustc_mir_dataflow/src/un_derefer.rs

@@ -29,72 +29,87 @@ impl<'tcx> UnDerefer<'tcx> {
         &self,
         place: PlaceRef<'tcx>,
     ) -> impl Iterator<Item = (PlaceRef<'tcx>, PlaceElem<'tcx>)> {
-        ProjectionIter::new(self.deref_chain(place.local), place)
+        self.co_iter_projections(place).into_iter(self)
+    }
+
+    /// Like [`UnDerefer::iter_projections`], but doesn't capture the self reference in the returned type.
+    /// Instead, getting the next element requires passing a reference to this `UnDerefer` for each iteration.
+    #[inline]
+    pub(crate) fn co_iter_projections(&self, place: PlaceRef<'tcx>) -> ProjectionCoroutine<'tcx> {
+        ProjectionCoroutine::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,
+pub(crate) enum ProjectionCoroutine<'tcx> {
+    InChain { current: PlaceRef<'tcx>, proj_idx: usize, last: PlaceRef<'tcx>, chain_idx: usize },
+    Last { last: PlaceRef<'tcx>, proj_idx: usize },
+    Finished,
 }
 
-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)
+impl<'tcx> ProjectionCoroutine<'tcx> {
+    fn new(deref_chain: &[PlaceRef<'tcx>], place: PlaceRef<'tcx>) -> Self {
+        if let &[first, ..] = deref_chain {
+            Self::InChain { current: first, proj_idx: 0, last: place, chain_idx: 0 }
         } else {
-            debug_assert!(deref_chain.is_empty());
-            None
-        };
-
-        ProjectionIter { places: SlicePlusOne { slice: deref_chain, last }, proj_idx: 0 }
+            if place.as_local().is_none() {
+                Self::Last { last: place, proj_idx: 0 }
+            } else {
+                Self::Finished
+            }
+        }
     }
-}
 
-impl<'tcx> Iterator for ProjectionIter<'_, 'tcx> {
-    type Item = (PlaceRef<'tcx>, PlaceElem<'tcx>);
+    fn advance_chain(&mut self, un_derefer: &UnDerefer<'tcx>) {
+        *self = match self {
+            &mut Self::InChain { last, chain_idx, .. } => {
+                let chain = un_derefer.deref_chain(last.local);
 
-    #[inline]
-    fn next(&mut self) -> Option<(PlaceRef<'tcx>, PlaceElem<'tcx>)> {
-        let place = self.places.read()?;
+                if let Some(&next) = chain.get(chain_idx + 1) {
+                    Self::InChain { current: next, proj_idx: 0, last, chain_idx: chain_idx + 1 }
+                } else {
+                    Self::Last { last, proj_idx: 0 }
+                }
+            }
+            &mut Self::Last { .. } => Self::Finished,
+            &mut Self::Finished => unreachable!(),
+        }
+    }
+
+    /// Returns the next `PlaceRef` and `PlaceElem` pair,
+    /// or `None` if the entire place has been iterated through.
+    ///
+    /// `un_derefer` must be the same instance that produced `self`.
+    pub(crate) fn next(
+        &mut self,
+        un_derefer: &UnDerefer<'tcx>,
+    ) -> Option<(PlaceRef<'tcx>, PlaceElem<'tcx>)> {
+        let (place, proj_idx) = match self {
+            &mut Self::InChain { current, ref mut proj_idx, .. } => (current, proj_idx),
+            &mut Self::Last { last, ref mut proj_idx } => (last, proj_idx),
+            &mut Self::Finished => return None,
+        };
 
         // 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];
+            PlaceRef { local: place.local, projection: &place.projection[..*proj_idx] };
+        let elem = place.projection[*proj_idx];
 
-        if self.proj_idx == place.projection.len() - 1 {
-            self.proj_idx = 0;
-            self.places.advance();
+        if *proj_idx == place.projection.len() - 1 {
+            self.advance_chain(un_derefer);
         } else {
-            self.proj_idx += 1;
+            *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 {
-            [_, remainder @ ..] => {
-                self.slice = remainder;
-            }
-            [] => self.last = None,
-        }
+    /// Returns a normal iterator over the remaining elements.
+    ///
+    /// `un_derefer` must be the same instance that produced `self`.
+    pub(crate) fn into_iter(
+        mut self,
+        un_derefer: &UnDerefer<'tcx>,
+    ) -> impl Iterator<Item = (PlaceRef<'tcx>, PlaceElem<'tcx>)> {
+        std::iter::from_fn(move || self.next(un_derefer))
     }
 }