From a36d8acd835c96de1b1184d97648a16648591607 Mon Sep 17 00:00:00 2001
From: Nicholas Nethercote <n.nethercote@gmail.com>
Date: Thu, 27 Feb 2025 18:10:17 +1100
Subject: [PATCH] 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.
---
 .../rustc_data_structures/src/sorted_map.rs   | 33 +++++++++++++++++
 .../interpret/allocation/provenance_map.rs    | 37 ++++++++++++++-----
 2 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/compiler/rustc_data_structures/src/sorted_map.rs b/compiler/rustc_data_structures/src/sorted_map.rs
index 066ea03b4ace4..a01a420dfbd66 100644
--- a/compiler/rustc_data_structures/src/sorted_map.rs
+++ b/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
diff --git a/compiler/rustc_middle/src/mir/interpret/allocation/provenance_map.rs b/compiler/rustc_middle/src/mir/interpret/allocation/provenance_map.rs
index 2f13f9635785a..82fb5f33b4c81 100644
--- a/compiler/rustc_middle/src/mir/interpret/allocation/provenance_map.rs
+++ b/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,