interpret: copy_provenance: avoid large intermediate buffer for large repeat counts

3 files changed, 25 insertions, 19 deletions

diff --git a/compiler/rustc_data_structures/src/lib.rs b/compiler/rustc_data_structures/src/lib.rs
index 17da3ea83c8..e4e86bcc41a 100644
--- a/compiler/rustc_data_structures/src/lib.rs
+++ b/compiler/rustc_data_structures/src/lib.rs
@@ -34,6 +34,7 @@
 #![feature(sized_hierarchy)]
 #![feature(test)]
 #![feature(thread_id_value)]
+#![feature(trusted_len)]
 #![feature(type_alias_impl_trait)]
 #![feature(unwrap_infallible)]
 // tidy-alphabetical-end
diff --git a/compiler/rustc_data_structures/src/sorted_map.rs b/compiler/rustc_data_structures/src/sorted_map.rs
index c002d47815b..c3c173b3fc4 100644
--- a/compiler/rustc_data_structures/src/sorted_map.rs
+++ b/compiler/rustc_data_structures/src/sorted_map.rs
@@ -1,6 +1,7 @@
 use std::borrow::Borrow;
 use std::cmp::Ordering;
 use std::fmt::Debug;
+use std::iter::TrustedLen;
 use std::mem;
 use std::ops::{Bound, Index, IndexMut, RangeBounds};
 
@@ -215,32 +216,36 @@ impl<K: Ord, V> SortedMap<K, V> {
     /// It is up to the caller to make sure that the elements are sorted by key
     /// and that there are no duplicates.
     #[inline]
-    pub fn insert_presorted(&mut self, elements: Vec<(K, V)>) {
-        if elements.is_empty() {
+    pub fn insert_presorted(
+        &mut self,
+        // We require `TrustedLen` to ensure that the `splice` below is actually efficient.
+        mut elements: impl Iterator<Item = (K, V)> + DoubleEndedIterator + TrustedLen,
+    ) {
+        let Some(first) = elements.next() else {
             return;
-        }
-
-        debug_assert!(elements.array_windows().all(|[fst, snd]| fst.0 < snd.0));
+        };
 
-        let start_index = self.lookup_index_for(&elements[0].0);
+        let start_index = self.lookup_index_for(&first.0);
 
         let elements = match start_index {
             Ok(index) => {
-                let mut elements = elements.into_iter();
-                self.data[index] = elements.next().unwrap();
-                elements
+                self.data[index] = first; // overwrite first element
+                elements.chain(None) // insert the rest below
             }
             Err(index) => {
-                if index == self.data.len() || elements.last().unwrap().0 < self.data[index].0 {
+                let last = elements.next_back();
+                if index == self.data.len()
+                    || last.as_ref().is_none_or(|l| l.0 < self.data[index].0)
+                {
                     // We can copy the whole range without having to mix with
                     // existing elements.
-                    self.data.splice(index..index, elements);
+                    self.data
+                        .splice(index..index, std::iter::once(first).chain(elements).chain(last));
                     return;
                 }
 
-                let mut elements = elements.into_iter();
-                self.data.insert(index, elements.next().unwrap());
-                elements
+                self.data.insert(index, first);
+                elements.chain(last) // insert the rest below
             }
         };
 
diff --git a/compiler/rustc_data_structures/src/sorted_map/tests.rs b/compiler/rustc_data_structures/src/sorted_map/tests.rs
index ea4d2f1feac..17d0d3cb170 100644
--- a/compiler/rustc_data_structures/src/sorted_map/tests.rs
+++ b/compiler/rustc_data_structures/src/sorted_map/tests.rs
@@ -171,7 +171,7 @@ fn test_insert_presorted_non_overlapping() {
     map.insert(2, 0);
     map.insert(8, 0);
 
-    map.insert_presorted(vec![(3, 0), (7, 0)]);
+    map.insert_presorted(vec![(3, 0), (7, 0)].into_iter());
 
     let expected = vec![2, 3, 7, 8];
     assert_eq!(keys(map), expected);
@@ -183,7 +183,7 @@ fn test_insert_presorted_first_elem_equal() {
     map.insert(2, 2);
     map.insert(8, 8);
 
-    map.insert_presorted(vec![(2, 0), (7, 7)]);
+    map.insert_presorted(vec![(2, 0), (7, 7)].into_iter());
 
     let expected = vec![(2, 0), (7, 7), (8, 8)];
     assert_eq!(elements(map), expected);
@@ -195,7 +195,7 @@ fn test_insert_presorted_last_elem_equal() {
     map.insert(2, 2);
     map.insert(8, 8);
 
-    map.insert_presorted(vec![(3, 3), (8, 0)]);
+    map.insert_presorted(vec![(3, 3), (8, 0)].into_iter());
 
     let expected = vec![(2, 2), (3, 3), (8, 0)];
     assert_eq!(elements(map), expected);
@@ -207,7 +207,7 @@ fn test_insert_presorted_shuffle() {
     map.insert(2, 2);
     map.insert(7, 7);
 
-    map.insert_presorted(vec![(1, 1), (3, 3), (8, 8)]);
+    map.insert_presorted(vec![(1, 1), (3, 3), (8, 8)].into_iter());
 
     let expected = vec![(1, 1), (2, 2), (3, 3), (7, 7), (8, 8)];
     assert_eq!(elements(map), expected);
@@ -219,7 +219,7 @@ fn test_insert_presorted_at_end() {
     map.insert(1, 1);
     map.insert(2, 2);
 
-    map.insert_presorted(vec![(3, 3), (8, 8)]);
+    map.insert_presorted(vec![(3, 3), (8, 8)].into_iter());
 
     let expected = vec![(1, 1), (2, 2), (3, 3), (8, 8)];
     assert_eq!(elements(map), expected);