[experiment] Try to make index vec smaller when possible

compiler/rustc_abi/src/layout/coroutine.rs

@@ -186,7 +186,8 @@ pub(super) fn layout<
             // "a" (`0..b_start`) and "b" (`b_start..`) correspond to
             // "outer" and "promoted" fields respectively.
             let b_start = tag_index.plus(1);
-            let offsets_b = IndexVec::from_raw(offsets.raw.split_off(b_start.index()));
+            let offsets_b =
+                IndexVec::from_raw(offsets.mutate(|raw| raw.split_off(b_start.index())));
             let offsets_a = offsets;
 
             // Disentangle the "a" and "b" components of `in_memory_order`
@@ -271,7 +272,7 @@ pub(super) fn layout<
 
             // Remove the unused slots to obtain the combined `in_memory_order`
             // (also see previous comment).
-            combined_in_memory_order.raw.retain(|&i| i.index() != invalid_field_idx);
+            combined_in_memory_order.mutate(|raw| raw.retain(|&i| i.index() != invalid_field_idx));
 
             variant.fields = FieldsShape::Arbitrary {
                 offsets: combined_offsets,

compiler/rustc_ast_lowering/src/lib.rs

@@ -494,8 +494,8 @@ fn index_ast<'tcx>(
     impl Indexer<'_, '_> {
         fn insert(&mut self, id: NodeId, node: AstOwner) {
             let def_id = self.owners[&id].def_id;
-            self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner);
-            self.index[def_id] = node;
+            let elem = self.index.ensure_contains_elem(def_id, || AstOwner::NonOwner);
+            *elem = node;
         }
 
         fn make_dummy<K>(

compiler/rustc_codegen_cranelift/src/base.rs

@@ -862,7 +862,7 @@ fn codegen_stmt<'tcx>(fx: &mut FunctionCx<'_, '_, 'tcx>, cur_block: Block, stmt:
                 {
                     let ty = to_place_and_rval.1.ty(&fx.mir.local_decls, fx.tcx);
                     let layout = fx.layout_of(fx.monomorphize(ty));
-                    let [data, meta] = &*operands.raw else {
+                    let [data, meta] = &operands.as_slice().raw else {
                         bug!("RawPtr fields: {operands:?}");
                     };
                     let data = codegen_operand(fx, data);

compiler/rustc_data_structures/src/graph/linked_graph/mod.rs

@@ -91,10 +91,12 @@ impl NodeIndex {
 }
 
 impl Idx for NodeIndex {
+    #[inline]
     fn new(idx: usize) -> NodeIndex {
         NodeIndex(idx)
     }
 
+    #[inline]
     fn index(self) -> usize {
         self.0
     }

compiler/rustc_data_structures/src/marker.rs

@@ -95,6 +95,7 @@ impl_dyn_send!(
     [indexmap::IndexSet<V, S> where V: DynSend, S: DynSend]
     [indexmap::IndexMap<K, V, S> where K: DynSend, V: DynSend, S: DynSend]
     [thin_vec::ThinVec<T> where T: DynSend]
+    [rustc_index::IndexVec<I, T> where I: rustc_index::Idx, T: DynSend]
     [smallvec::SmallVec<A> where A: smallvec::Array + DynSend]
 );
 
@@ -182,6 +183,7 @@ impl_dyn_sync!(
     [indexmap::IndexMap<K, V, S> where K: DynSync, V: DynSync, S: DynSync]
     [smallvec::SmallVec<A> where A: smallvec::Array + DynSync]
     [thin_vec::ThinVec<T> where T: DynSync]
+    [rustc_index::IndexVec<I, T> where I: rustc_index::Idx, T: DynSync]
 );
 
 pub fn assert_dyn_sync<T: ?Sized + PointeeSized + DynSync>() {}

compiler/rustc_index/src/lib.rs

@@ -1,6 +1,7 @@
 // tidy-alphabetical-start
 #![cfg_attr(all(feature = "nightly", test), feature(stmt_expr_attributes))]
 #![cfg_attr(all(feature = "nightly", test), feature(test))]
+#![cfg_attr(feature = "nightly", feature(dropck_eyepatch))]
 #![cfg_attr(feature = "nightly", feature(extend_one, step_trait))]
 // tidy-alphabetical-end
 

compiler/rustc_index/src/slice.rs

@@ -254,7 +254,10 @@ impl<I: Idx, T: Clone> ToOwned for IndexSlice<I, T> {
     }
 
     fn clone_into(&self, target: &mut IndexVec<I, T>) {
-        self.raw.clone_into(&mut target.raw)
+        target.mutate(|v| {
+            // todo this coerces to a slice now, is that correct?
+            self.raw.clone_into(v)
+        });
     }
 }
 

compiler/rustc_index/src/vec.rs

@@ -1,8 +1,10 @@
 use std::borrow::{Borrow, BorrowMut};
-use std::hash::Hash;
+use std::hash::{Hash, Hasher};
 use std::marker::PhantomData;
+use std::mem::ManuallyDrop;
 use std::ops::{Deref, DerefMut, RangeBounds};
-use std::{fmt, slice, vec};
+use std::ptr::NonNull;
+use std::{fmt, ptr, slice, vec};
 
 #[cfg(feature = "nightly")]
 use rustc_serialize::{Decodable, Decoder, Encodable, Encoder};
@@ -35,24 +37,85 @@ use crate::{Idx, IndexSlice};
 /// This allows to index the IndexVec with the new index type.
 ///
 /// [`newtype_index!`]: ../macro.newtype_index.html
-#[derive(Clone, PartialEq, Eq, Hash)]
-#[repr(transparent)]
 pub struct IndexVec<I: Idx, T> {
-    pub raw: Vec<T>,
+    data: NonNull<T>,
+    len: I,
+    capacity: I,
+
     _marker: PhantomData<fn(&I)>,
+    _marker2: PhantomData<T>,
+}
+
+impl<I: Idx, T: Clone> Clone for IndexVec<I, T> {
+    fn clone(&self) -> Self {
+        IndexVec::from_raw(self.as_slice().raw.to_vec())
+    }
+}
+
+impl<I: Idx, T: PartialEq> PartialEq for IndexVec<I, T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.as_slice().eq(other.as_slice())
+    }
+}
+
+impl<I: Idx, T: PartialEq> Eq for IndexVec<I, T> {}
+
+impl<I: Idx, T: Hash> Hash for IndexVec<I, T> {
+    fn hash<H: Hasher>(&self, state: &mut H) {
+        self.as_slice().hash(state);
+    }
+}
+
+#[cfg(feature = "nightly")]
+unsafe impl<I: Idx, #[may_dangle] T> Drop for IndexVec<I, T> {
+    #[inline]
+    fn drop(&mut self) {
+        if self.capacity.index() == 0 {
+            return;
+        }
+        std::mem::take(self).into_vec();
+    }
+}
+
+#[cfg(not(feature = "nightly"))]
+impl<I: Idx, T> Drop for IndexVec<I, T> {
+    fn drop(&mut self) {
+        std::mem::take(self).into_vec();
+    }
 }
 
 impl<I: Idx, T> IndexVec<I, T> {
     /// Constructs a new, empty `IndexVec<I, T>`.
     #[inline]
-    pub const fn new() -> Self {
+    pub fn new() -> Self {
         IndexVec::from_raw(Vec::new())
     }
 
     /// Constructs a new `IndexVec<I, T>` from a `Vec<T>`.
     #[inline]
-    pub const fn from_raw(raw: Vec<T>) -> Self {
-        IndexVec { raw, _marker: PhantomData }
+    pub fn from_raw(raw: Vec<T>) -> Self {
+        let mut me = ManuallyDrop::new(raw);
+
+        IndexVec {
+            data: unsafe { NonNull::new_unchecked(me.as_mut_ptr()) },
+            len: I::new(me.len()),
+            capacity: I::new(me.capacity()),
+
+            _marker: PhantomData,
+            _marker2: PhantomData,
+        }
+    }
+
+    #[inline]
+    pub fn into_vec(self) -> Vec<T> {
+        let me = ManuallyDrop::new(self);
+        // fixme this is unsound because we rely on correct Idx trait impls
+        let len = me.len.index();
+        let cap = me.capacity.index();
+        unsafe {
+            core::hint::assert_unchecked(len <= cap);
+        }
+        unsafe { Vec::from_raw_parts(me.data.as_ptr(), len, cap) }
     }
 
     #[inline]
@@ -98,32 +161,52 @@ impl<I: Idx, T> IndexVec<I, T> {
         IndexVec::from_raw((0..n).map(I::new).map(func).collect())
     }
 
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.len.index()
+    }
+
     #[inline]
     pub fn as_slice(&self) -> &IndexSlice<I, T> {
-        IndexSlice::from_raw(&self.raw)
+        IndexSlice::from_raw(unsafe {
+            &*std::ptr::slice_from_raw_parts(self.data.as_ptr(), self.len.index())
+        })
     }
 
     #[inline]
     pub fn as_mut_slice(&mut self) -> &mut IndexSlice<I, T> {
-        IndexSlice::from_raw_mut(&mut self.raw)
+        IndexSlice::from_raw_mut(unsafe {
+            &mut *ptr::slice_from_raw_parts_mut(self.data.as_ptr(), self.len.index())
+        })
     }
 
     /// Pushes an element to the array returning the index where it was pushed to.
     #[inline]
     pub fn push(&mut self, d: T) -> I {
         let idx = self.next_index();
-        self.raw.push(d);
+        let len = self.len.index();
+
+        if len < self.capacity.index() {
+            unsafe {
+                // todo copy pasta from Vec::push_mut
+                let end = self.data.add(len);
+                ptr::write(end.as_ptr(), d);
+                self.len = I::new(len + 1);
+            }
+        } else {
+            self.mutate(|vec| vec.push(d));
+        }
         idx
     }
 
     #[inline]
     pub fn pop(&mut self) -> Option<T> {
-        self.raw.pop()
+        self.mutate(|raw| raw.pop())
     }
 
     #[inline]
     pub fn into_iter(self) -> vec::IntoIter<T> {
-        self.raw.into_iter()
+        self.into_vec().into_iter()
     }
 
     #[inline]
@@ -132,35 +215,39 @@ impl<I: Idx, T> IndexVec<I, T> {
     ) -> impl DoubleEndedIterator<Item = (I, T)> + ExactSizeIterator {
         // Allow the optimizer to elide the bounds checking when creating each index.
         let _ = I::new(self.len());
-        self.raw.into_iter().enumerate().map(|(n, t)| (I::new(n), t))
+        self.into_iter().enumerate().map(|(n, t)| (I::new(n), t))
     }
 
     #[inline]
-    pub fn drain<R: RangeBounds<usize>>(&mut self, range: R) -> impl Iterator<Item = T> {
-        self.raw.drain(range)
+    pub fn drain_into<R: RangeBounds<usize>>(&mut self, range: R, target: &mut IndexVec<I, T>) {
+        self.mutate(|raw| target.extend(raw.drain(range)))
     }
 
-    #[inline]
-    pub fn drain_enumerated<R: RangeBounds<usize>>(
-        &mut self,
-        range: R,
-    ) -> impl Iterator<Item = (I, T)> {
-        let begin = match range.start_bound() {
-            std::ops::Bound::Included(i) => *i,
-            std::ops::Bound::Excluded(i) => i.checked_add(1).unwrap(),
-            std::ops::Bound::Unbounded => 0,
+    #[inline(always)]
+    pub fn mutate<U, F: FnOnce(&mut Vec<T>) -> U>(&mut self, f: F) -> U {
+        let sentinel = IndexVec {
+            data: NonNull::dangling(),
+            len: I::new(0),
+            capacity: I::new(0),
+            _marker: PhantomData,
+            _marker2: PhantomData,
         };
-        self.raw.drain(range).enumerate().map(move |(n, t)| (I::new(begin + n), t))
+        let mut vec = std::mem::replace(self, sentinel).into_vec();
+
+        // let mut vec = std::mem::take(self).into_vec();
+        let v = f(&mut vec);
+        let _ = std::mem::replace(self, IndexVec::from_raw(vec));
+        v
     }
 
     #[inline]
     pub fn shrink_to_fit(&mut self) {
-        self.raw.shrink_to_fit()
+        self.mutate(|vec| vec.shrink_to_fit());
     }
 
     #[inline]
     pub fn truncate(&mut self, a: usize) {
-        self.raw.truncate(a)
+        self.mutate(|vec| vec.truncate(a))
     }
 
     /// Grows the index vector so that it contains an entry for
@@ -173,7 +260,7 @@ impl<I: Idx, T> IndexVec<I, T> {
     pub fn ensure_contains_elem(&mut self, elem: I, fill_value: impl FnMut() -> T) -> &mut T {
         let min_new_len = elem.index() + 1;
         if self.len() < min_new_len {
-            self.raw.resize_with(min_new_len, fill_value);
+            self.mutate(|vec| vec.resize_with(min_new_len, fill_value));
         }
 
         &mut self[elem]
@@ -184,18 +271,18 @@ impl<I: Idx, T> IndexVec<I, T> {
     where
         T: Clone,
     {
-        self.raw.resize(new_len, value)
+        self.mutate(|vec| vec.resize(new_len, value))
     }
 
     #[inline]
     pub fn resize_to_elem(&mut self, elem: I, fill_value: impl FnMut() -> T) {
         let min_new_len = elem.index() + 1;
-        self.raw.resize_with(min_new_len, fill_value);
+        self.mutate(|vec| vec.resize_with(min_new_len, fill_value));
     }
 
     #[inline]
     pub fn append(&mut self, other: &mut Self) {
-        self.raw.append(&mut other.raw);
+        self.mutate(|vec| other.mutate(|other| vec.append(other)));
     }
 }
 
@@ -259,19 +346,19 @@ impl<I: Idx, T> BorrowMut<IndexSlice<I, T>> for IndexVec<I, T> {
 impl<I: Idx, T> Extend<T> for IndexVec<I, T> {
     #[inline]
     fn extend<J: IntoIterator<Item = T>>(&mut self, iter: J) {
-        self.raw.extend(iter);
+        self.mutate(|vec| vec.extend(iter));
     }
 
     #[inline]
     #[cfg(feature = "nightly")]
     fn extend_one(&mut self, item: T) {
-        self.raw.push(item);
+        self.mutate(|vec| vec.push(item));
     }
 
     #[inline]
     #[cfg(feature = "nightly")]
     fn extend_reserve(&mut self, additional: usize) {
-        self.raw.reserve(additional);
+        self.mutate(|vec| vec.reserve(additional));
     }
 }
 
@@ -291,7 +378,7 @@ impl<I: Idx, T> IntoIterator for IndexVec<I, T> {
 
     #[inline]
     fn into_iter(self) -> vec::IntoIter<T> {
-        self.raw.into_iter()
+        self.into_vec().into_iter()
     }
 }
 
@@ -347,5 +434,15 @@ impl<D: Decoder, I: Idx, T: Decodable<D>> Decodable<D> for IndexVec<I, T> {
 // not the phantom data.
 unsafe impl<I: Idx, T> Send for IndexVec<I, T> where T: Send {}
 
+#[cfg(target_pointer_width = "64")]
+mod size_asserts {
+    use super::IndexVec;
+    use crate::static_assert_size;
+    static_assert_size!(IndexVec<u32, u8>, 16);
+    static_assert_size!(IndexVec<usize, u8>, 24);
+    static_assert_size!(Option<IndexVec<u32, u8>>, 16);
+    static_assert_size!(Option<IndexVec<usize, u8>>, 24);
+}
+
 #[cfg(test)]
 mod tests;