Optimize empty provenance range checks.

compiler/rustc_data_structures/src/sorted_map.rs

@@ -1,4 +1,5 @@
 use std::borrow::Borrow;
+use std::cmp::Ordering;
 use std::fmt::Debug;
 use std::mem;
 use std::ops::{Bound, Index, IndexMut, RangeBounds};
@@ -156,6 +157,38 @@ impl<K: Ord, V> SortedMap<K, V> {
         &self.data[start..end]
     }
 
+    /// `sm.range_is_empty(r)` == `sm.range(r).is_empty()`, but is faster.
+    #[inline]
+    pub fn range_is_empty<R>(&self, range: R) -> bool
+    where
+        R: RangeBounds<K>,
+    {
+        // `range` must (via `range_slice_indices`) search for the start and
+        // end separately. But here we can do a single binary search for the
+        // entire range. If a single `x` matching `range` is found then the
+        // range is *not* empty.
+        self.data
+            .binary_search_by(|(x, _)| {
+                // Is `x` below `range`?
+                match range.start_bound() {
+                    Bound::Included(start) if x < start => return Ordering::Less,
+                    Bound::Excluded(start) if x <= start => return Ordering::Less,
+                    _ => {}
+                };
+
+                // Is `x` above `range`?
+                match range.end_bound() {
+                    Bound::Included(end) if x > end => return Ordering::Greater,
+                    Bound::Excluded(end) if x >= end => return Ordering::Greater,
+                    _ => {}
+                };
+
+                // `x` must be within `range`.
+                Ordering::Equal
+            })
+            .is_err()
+    }
+
     #[inline]
     pub fn remove_range<R>(&mut self, range: R)
     where

compiler/rustc_middle/src/mir/interpret/allocation/provenance_map.rs

@@ -2,6 +2,7 @@
 //! representation for the common case where PTR_SIZE consecutive bytes have the same provenance.
 
 use std::cmp;
+use std::ops::Range;
 
 use rustc_abi::{HasDataLayout, Size};
 use rustc_data_structures::sorted_map::SortedMap;
@@ -66,6 +67,15 @@ impl ProvenanceMap {
 }
 
 impl<Prov: Provenance> ProvenanceMap<Prov> {
+    fn adjusted_range(range: AllocRange, cx: &impl HasDataLayout) -> Range<Size> {
+        // We have to go back `pointer_size - 1` bytes, as that one would still overlap with
+        // the beginning of this range.
+        let adjusted_start = Size::from_bytes(
+            range.start.bytes().saturating_sub(cx.data_layout().pointer_size.bytes() - 1),
+        );
+        adjusted_start..range.end()
+    }
+
     /// Returns all ptr-sized provenance in the given range.
     /// If the range has length 0, returns provenance that crosses the edge between `start-1` and
     /// `start`.
@@ -74,12 +84,17 @@ impl<Prov: Provenance> ProvenanceMap<Prov> {
         range: AllocRange,
         cx: &impl HasDataLayout,
     ) -> &[(Size, Prov)] {
-        // We have to go back `pointer_size - 1` bytes, as that one would still overlap with
-        // the beginning of this range.
-        let adjusted_start = Size::from_bytes(
-            range.start.bytes().saturating_sub(cx.data_layout().pointer_size.bytes() - 1),
-        );
-        self.ptrs.range(adjusted_start..range.end())
+        self.ptrs.range(Self::adjusted_range(range, cx))
+    }
+
+    /// `pm.range_get_ptrs_is_empty(r, cx)` == `pm.range_get_ptrs(r, cx).is_empty()`, but is
+    /// faster.
+    pub(super) fn range_get_ptrs_is_empty(
+        &self,
+        range: AllocRange,
+        cx: &impl HasDataLayout,
+    ) -> bool {
+        self.ptrs.range_is_empty(Self::adjusted_range(range, cx))
     }
 
     /// Returns all byte-wise provenance in the given range.
@@ -117,7 +132,7 @@ impl<Prov: Provenance> ProvenanceMap<Prov> {
     /// limit access to provenance outside of the `Allocation` abstraction.
     ///
     pub fn range_empty(&self, range: AllocRange, cx: &impl HasDataLayout) -> bool {
-        self.range_get_ptrs(range, cx).is_empty() && self.range_get_bytes(range).is_empty()
+        self.range_get_ptrs_is_empty(range, cx) && self.range_get_bytes(range).is_empty()
     }
 
     /// Yields all the provenances stored in this map.
@@ -149,12 +164,14 @@ impl<Prov: Provenance> ProvenanceMap<Prov> {
         // provenance that overlaps with the given range.
         let (first, last) = {
             // Find all provenance overlapping the given range.
-            let provenance = self.range_get_ptrs(range, cx);
-            if provenance.is_empty() {
-                // No provenance in this range, we are done.
+            if self.range_get_ptrs_is_empty(range, cx) {
+                // No provenance in this range, we are done. This is the common case.
                 return Ok(());
             }
 
+            // This redoes some of the work of `range_get_ptrs_is_empty`, but this path is much
+            // colder than the early return above, so it's worth it.
+            let provenance = self.range_get_ptrs(range, cx);
             (
                 provenance.first().unwrap().0,
                 provenance.last().unwrap().0 + cx.data_layout().pointer_size,