Add supporting APIs for MIR storage reconstruction

Author: Amanieu

compiler/rustc_index/src/interval.rs

@@ -1,6 +1,7 @@
 use std::iter::Step;
 use std::marker::PhantomData;
-use std::ops::{Bound, Range, RangeBounds};
+use std::ops::{Bound, RangeBounds};
+use std::range::RangeInclusive;
 
 use smallvec::SmallVec;
 
@@ -59,11 +60,14 @@ impl<I: Idx> IntervalSet<I> {
     }
 
     /// Iterates through intervals stored in the set, in order.
-    pub fn iter_intervals(&self) -> impl Iterator<Item = std::ops::Range<I>>
+    pub fn iter_intervals(&self) -> impl Iterator<Item = RangeInclusive<I>>
     where
         I: Step,
     {
-        self.map.iter().map(|&(start, end)| I::new(start as usize)..I::new(end as usize + 1))
+        self.map.iter().map(|&(start, end)| RangeInclusive {
+            start: I::new(start as usize),
+            last: I::new(end as usize),
+        })
     }
 
     /// Returns true if we increased the number of elements present.
@@ -204,17 +208,38 @@ impl<I: Idx> IntervalSet<I> {
         needle <= *prev_end
     }
 
+    /// Returns whether any point in `range` is contained in the set.
+    pub fn intersects_range(&self, range: impl RangeBounds<I> + Clone) -> bool {
+        let start = inclusive_start(range.clone());
+        let Some(end) = inclusive_end(self.domain, range) else {
+            // empty range
+            return false;
+        };
+        if start > end {
+            return false;
+        }
+
+        // Find the last interval whose start is <= end.
+        let Some(last) = self.map.partition_point(|r| r.0 <= end).checked_sub(1) else {
+            // All ranges in the map start after the new range's end
+            return false;
+        };
+        let (_, prev_end) = &self.map[last];
+        start <= *prev_end
+    }
+
     pub fn superset(&self, other: &IntervalSet<I>) -> bool
     where
         I: Step,
     {
         let mut sup_iter = self.iter_intervals();
         let mut current = None;
-        let contains = |sup: Range<I>, sub: Range<I>, current: &mut Option<Range<I>>| {
-            if sup.end < sub.start {
-                // if `sup.end == sub.start`, the next sup doesn't contain `sub.start`
+        let contains = |sup: RangeInclusive<I>,
+                        sub: RangeInclusive<I>,
+                        current: &mut Option<RangeInclusive<I>>| {
+            if sup.last < sub.start {
                 None // continue to the next sup
-            } else if sup.end >= sub.end && sup.start <= sub.start {
+            } else if sup.last >= sub.last && sup.start <= sub.start {
                 *current = Some(sup); // save the current sup
                 Some(true)
             } else {
@@ -224,8 +249,8 @@ impl<I: Idx> IntervalSet<I> {
         other.iter_intervals().all(|sub| {
             current
                 .take()
-                .and_then(|sup| contains(sup, sub.clone(), &mut current))
-                .or_else(|| sup_iter.find_map(|sup| contains(sup, sub.clone(), &mut current)))
+                .and_then(|sup| contains(sup, sub, &mut current))
+                .or_else(|| sup_iter.find_map(|sup| contains(sup, sub, &mut current)))
                 .unwrap_or(false)
         })
     }
@@ -242,11 +267,11 @@ impl<I: Idx> IntervalSet<I> {
             let mut other_current = other_iter.next()?;
 
             loop {
-                if self_current.end <= other_current.start {
+                if self_current.last < other_current.start {
                     self_current = self_iter.next()?;
                     continue;
                 }
-                if other_current.end <= self_current.start {
+                if other_current.last < self_current.start {
                     other_current = other_iter.next()?;
                     continue;
                 }

compiler/rustc_index/src/interval/tests.rs

@@ -5,7 +5,7 @@ fn insert_collapses() {
     let mut set = IntervalSet::<u32>::new(10000);
     set.insert_range(9831..=9837);
     set.insert_range(43..=9830);
-    assert_eq!(set.iter_intervals().collect::<Vec<_>>(), [43..9838]);
+    assert_eq!(set.iter_intervals().collect::<Vec<_>>(), [(43..=9837).into()]);
 }
 
 #[test]

compiler/rustc_mir_dataflow/src/points.rs

@@ -61,6 +61,15 @@ impl DenseLocationMap {
         PointIndex::new(start_index)
     }
 
+    /// Returns the `PointIndex` for the terminator in the given `BasicBlock`. O(1).
+    #[inline]
+    pub fn terminator(&self, block: BasicBlock) -> PointIndex {
+        let next_block = BasicBlock::new(block.index() + 1);
+        let next_start_index =
+            *self.statements_before_block.get(next_block).unwrap_or(&self.num_points);
+        PointIndex::new(next_start_index - 1)
+    }
+
     /// Return the PointIndex for the block start of this index.
     #[inline]
     pub fn to_block_start(&self, index: PointIndex) -> PointIndex {

compiler/rustc_mir_transform/src/patch.rs

@@ -215,6 +215,21 @@ impl<'tcx> MirPatch<'tcx> {
         self.term_patch_map.insert(block, new);
     }
 
+    /// Modifies the terminator of a block, reading the existing patch if one exists or
+    /// cloning from the body otherwise.
+    pub(crate) fn mutate_terminator(
+        &mut self,
+        body: &Body<'tcx>,
+        bb: BasicBlock,
+        f: impl FnOnce(&mut TerminatorKind<'tcx>),
+    ) {
+        let kind = self
+            .term_patch_map
+            .entry(bb)
+            .or_insert_with(|| body.basic_blocks[bb].terminator().kind.clone());
+        f(kind);
+    }
+
     /// Mark given statement to be replaced by a `Nop`.
     ///
     /// This method only works on statements from the initial body, and cannot be used to remove