Auto merge of rust-lang#137704 - nnethercote:opt-empty-prov-range-checks, r=oli-obk

Optimize empty provenance range checks.

Currently it gets the pointers in the range and checks if the result is empty, but it can be done faster if you combine those two steps.

r? `@oli-obk`

Author: bors

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,