Implement a `find_descendant` method for `MovePath`

2 files changed, 47 insertions, 41 deletions

diff --git a/src/librustc_mir/dataflow/at_location.rs b/src/librustc_mir/dataflow/at_location.rs
index c6d29211d45..e4eb8506846 100644
--- a/src/librustc_mir/dataflow/at_location.rs
+++ b/src/librustc_mir/dataflow/at_location.rs
@@ -4,7 +4,6 @@
 use rustc::mir::{BasicBlock, Location};
 use rustc_index::bit_set::{BitIter, BitSet, HybridBitSet};
 
-use crate::dataflow::move_paths::{HasMoveData, MovePathIndex};
 use crate::dataflow::{BitDenotation, DataflowResults, GenKillSet};
 
 use std::borrow::Borrow;
@@ -168,43 +167,3 @@ where
         self.stmt_trans.apply(&mut self.curr_state)
     }
 }
-
-impl<'tcx, T, DR> FlowAtLocation<'tcx, T, DR>
-where
-    T: HasMoveData<'tcx> + BitDenotation<'tcx, Idx = MovePathIndex>,
-    DR: Borrow<DataflowResults<'tcx, T>>,
-{
-    pub fn has_any_child_of(&self, mpi: T::Idx) -> Option<T::Idx> {
-        // We process `mpi` before the loop below, for two reasons:
-        // - it's a little different from the loop case (we don't traverse its
-        //   siblings);
-        // - ~99% of the time the loop isn't reached, and this code is hot, so
-        //   we don't want to allocate `todo` unnecessarily.
-        if self.contains(mpi) {
-            return Some(mpi);
-        }
-        let move_data = self.operator().move_data();
-        let move_path = &move_data.move_paths[mpi];
-        let mut todo = if let Some(child) = move_path.first_child {
-            vec![child]
-        } else {
-            return None;
-        };
-
-        while let Some(mpi) = todo.pop() {
-            if self.contains(mpi) {
-                return Some(mpi);
-            }
-            let move_path = &move_data.move_paths[mpi];
-            if let Some(child) = move_path.first_child {
-                todo.push(child);
-            }
-            // After we've processed the original `mpi`, we should always
-            // traverse the siblings of any of its children.
-            if let Some(sibling) = move_path.next_sibling {
-                todo.push(sibling);
-            }
-        }
-        return None;
-    }
-}
diff --git a/src/librustc_mir/dataflow/move_paths/mod.rs b/src/librustc_mir/dataflow/move_paths/mod.rs
index 8d62b84bda8..614a2761643 100644
--- a/src/librustc_mir/dataflow/move_paths/mod.rs
+++ b/src/librustc_mir/dataflow/move_paths/mod.rs
@@ -72,6 +72,41 @@ impl<'tcx> MovePath<'tcx> {
 
         parents
     }
+
+    /// Finds the closest descendant of `self` for which `f` returns `true` using a breadth-first
+    /// search.
+    ///
+    /// `f` will **not** be called on `self`.
+    pub fn find_descendant(
+        &self,
+        move_paths: &IndexVec<MovePathIndex, MovePath<'_>>,
+        f: impl Fn(MovePathIndex) -> bool,
+    ) -> Option<MovePathIndex> {
+        let mut todo = if let Some(child) = self.first_child {
+            vec![child]
+        } else {
+            return None;
+        };
+
+        while let Some(mpi) = todo.pop() {
+            if f(mpi) {
+                return Some(mpi);
+            }
+
+            let move_path = &move_paths[mpi];
+            if let Some(child) = move_path.first_child {
+                todo.push(child);
+            }
+
+            // After we've processed the original `mpi`, we should always
+            // traverse the siblings of any of its children.
+            if let Some(sibling) = move_path.next_sibling {
+                todo.push(sibling);
+            }
+        }
+
+        None
+    }
 }
 
 impl<'tcx> fmt::Debug for MovePath<'tcx> {
@@ -333,4 +368,16 @@ impl<'tcx> MoveData<'tcx> {
             }
         }
     }
+
+    pub fn find_in_move_path_or_its_descendants(
+        &self,
+        root: MovePathIndex,
+        pred: impl Fn(MovePathIndex) -> bool,
+    ) -> Option<MovePathIndex> {
+        if pred(root) {
+            return Some(root);
+        }
+
+        self.move_paths[root].find_descendant(&self.move_paths, pred)
+    }
 }