From 7d2cfcab9d8291e412f41e43cdadf587d700d577 Mon Sep 17 00:00:00 2001
From: Michał Kostrubiec <fractalfirdev@gmail.com>
Date: Tue, 25 Feb 2025 00:43:25 +0100
Subject: Changed interners to start with preallocated capacity

---
 compiler/rustc_data_structures/src/sharded.rs | 3 +++
 1 file changed, 3 insertions(+)

(limited to 'compiler/rustc_data_structures/src')

diff --git a/compiler/rustc_data_structures/src/sharded.rs b/compiler/rustc_data_structures/src/sharded.rs
index 65488c73d3c..5a53f8af5f8 100644
--- a/compiler/rustc_data_structures/src/sharded.rs
+++ b/compiler/rustc_data_structures/src/sharded.rs
@@ -143,6 +143,9 @@ pub fn shards() -> usize {
 pub type ShardedHashMap<K, V> = Sharded<FxHashMap<K, V>>;
 
 impl<K: Eq, V> ShardedHashMap<K, V> {
+    pub fn with_capacity(cap: usize) -> Self {
+        Self::new(|| FxHashMap::with_capacity_and_hasher(cap, rustc_hash::FxBuildHasher::default()))
+    }
     pub fn len(&self) -> usize {
         self.lock_shards().map(|shard| shard.len()).sum()
     }
-- 
cgit 1.4.1-3-g733a5


From 3477297c374f2b12a5846b8b522b1e69ea08cc5e Mon Sep 17 00:00:00 2001
From: bjorn3 <17426603+bjorn3@users.noreply.github.com>
Date: Fri, 24 Jan 2025 15:40:41 +0000
Subject: Fix UB in ThinVec::flat_map_in_place

thin_vec.as_ptr() goes through the Deref impl of ThinVec, which will
not allow access to any memory as we did call set_len(0) first.
---
 .../rustc_data_structures/src/flat_map_in_place.rs | 42 +++++++++++++---------
 1 file changed, 25 insertions(+), 17 deletions(-)

(limited to 'compiler/rustc_data_structures/src')

diff --git a/compiler/rustc_data_structures/src/flat_map_in_place.rs b/compiler/rustc_data_structures/src/flat_map_in_place.rs
index e66b00b7557..6d718059f9c 100644
--- a/compiler/rustc_data_structures/src/flat_map_in_place.rs
+++ b/compiler/rustc_data_structures/src/flat_map_in_place.rs
@@ -1,4 +1,4 @@
-use std::ptr;
+use std::{mem, ptr};
 
 use smallvec::{Array, SmallVec};
 use thin_vec::ThinVec;
@@ -13,39 +13,44 @@ pub trait FlatMapInPlace<T>: Sized {
 // The implementation of this method is syntactically identical for all the
 // different vector types.
 macro_rules! flat_map_in_place {
-    () => {
+    ($vec:ident $( where T: $bound:path)?) => {
         fn flat_map_in_place<F, I>(&mut self, mut f: F)
         where
             F: FnMut(T) -> I,
             I: IntoIterator<Item = T>,
         {
+            struct LeakGuard<'a, T $(: $bound)?>(&'a mut $vec<T>);
+
+            impl<'a, T $(: $bound)?> Drop for LeakGuard<'a, T> {
+                fn drop(&mut self) {
+                    unsafe {
+                        self.0.set_len(0); // make sure we just leak elements in case of panic
+                    }
+                }
+            }
+
+            let this = LeakGuard(self);
+
             let mut read_i = 0;
             let mut write_i = 0;
             unsafe {
-                let mut old_len = self.len();
-                self.set_len(0); // make sure we just leak elements in case of panic
-
-                while read_i < old_len {
+                while read_i < this.0.len() {
                     // move the read_i'th item out of the vector and map it
                     // to an iterator
-                    let e = ptr::read(self.as_ptr().add(read_i));
+                    let e = ptr::read(this.0.as_ptr().add(read_i));
                     let iter = f(e).into_iter();
                     read_i += 1;
 
                     for e in iter {
                         if write_i < read_i {
-                            ptr::write(self.as_mut_ptr().add(write_i), e);
+                            ptr::write(this.0.as_mut_ptr().add(write_i), e);
                             write_i += 1;
                         } else {
                             // If this is reached we ran out of space
                             // in the middle of the vector.
                             // However, the vector is in a valid state here,
                             // so we just do a somewhat inefficient insert.
-                            self.set_len(old_len);
-                            self.insert(write_i, e);
-
-                            old_len = self.len();
-                            self.set_len(0);
+                            this.0.insert(write_i, e);
 
                             read_i += 1;
                             write_i += 1;
@@ -54,20 +59,23 @@ macro_rules! flat_map_in_place {
                 }
 
                 // write_i tracks the number of actually written new items.
-                self.set_len(write_i);
+                this.0.set_len(write_i);
+
+                // The ThinVec is in a sane state again. Prevent the LeakGuard from leaking the data.
+                mem::forget(this);
             }
         }
     };
 }
 
 impl<T> FlatMapInPlace<T> for Vec<T> {
-    flat_map_in_place!();
+    flat_map_in_place!(Vec);
 }
 
 impl<T, A: Array<Item = T>> FlatMapInPlace<T> for SmallVec<A> {
-    flat_map_in_place!();
+    flat_map_in_place!(SmallVec where T: Array);
 }
 
 impl<T> FlatMapInPlace<T> for ThinVec<T> {
-    flat_map_in_place!();
+    flat_map_in_place!(ThinVec);
 }
-- 
cgit 1.4.1-3-g733a5