Use a safe BucketIndex abstraction in VecCache

The current code for indexing into bucket arrays is quite tricky and unsafe,
partly because it has to keep manually assuring the compiler that a bucket
index is always less than 21.

By encapsulating that knowledge in a 21-value enum, we can make the code
clearer and safer, without giving up performance.

Having a dedicated `BucketIndex` type could also help with further cleanups of
`VecCache` indexing.

Author: Zalathar

compiler/rustc_data_structures/src/vec_cache.rs

@@ -6,12 +6,16 @@
 //!
 //! This is currently used for query caching.
 
-use std::fmt::Debug;
+use std::fmt::{self, Debug};
 use std::marker::PhantomData;
+use std::ops::{Index, IndexMut};
 use std::sync::atomic::{AtomicPtr, AtomicU32, AtomicUsize, Ordering};
 
 use rustc_index::Idx;
 
+#[cfg(test)]
+mod tests;
+
 struct Slot<V> {
     // We never construct &Slot<V> so it's fine for this to not be in an UnsafeCell.
     value: V,
@@ -28,7 +32,7 @@ struct Slot<V> {
 #[derive(Copy, Clone, Debug)]
 struct SlotIndex {
     // the index of the bucket in VecCache (0 to 20)
-    bucket_idx: usize,
+    bucket_idx: BucketIndex,
     // the index of the slot within the bucket
     index_in_bucket: usize,
 }
@@ -42,7 +46,7 @@ const ENTRIES_BY_BUCKET: [usize; BUCKETS] = {
     let mut key = 0;
     loop {
         let si = SlotIndex::from_index(key);
-        entries[si.bucket_idx] = si.entries();
+        entries[si.bucket_idx.to_usize()] = si.bucket_idx.capacity();
         if key == 0 {
             key = 1;
         } else if key == (1 << 31) {
@@ -57,11 +61,6 @@ const ENTRIES_BY_BUCKET: [usize; BUCKETS] = {
 const BUCKETS: usize = 21;
 
 impl SlotIndex {
-    /// The total possible number of entries in the bucket
-    const fn entries(&self) -> usize {
-        if self.bucket_idx == 0 { 1 << 12 } else { 1 << (self.bucket_idx + 11) }
-    }
-
     // This unpacks a flat u32 index into identifying which bucket it belongs to and the offset
     // within that bucket. As noted in the VecCache docs, buckets double in size with each index.
     // Typically that would mean 31 buckets (2^0 + 2^1 ... + 2^31 = u32::MAX - 1), but to reduce
@@ -73,32 +72,21 @@ impl SlotIndex {
     // slots (see `slot_index_exhaustive` in tests).
     #[inline]
     const fn from_index(idx: u32) -> Self {
-        const FIRST_BUCKET_SHIFT: usize = 12;
-        if idx < (1 << FIRST_BUCKET_SHIFT) {
-            return SlotIndex { bucket_idx: 0, index_in_bucket: idx as usize };
-        }
-        // We already ruled out idx 0, so this `ilog2` never panics (and the check optimizes away)
-        let bucket = idx.ilog2() as usize;
-        let entries = 1 << bucket;
-        SlotIndex {
-            bucket_idx: bucket - FIRST_BUCKET_SHIFT + 1,
-            index_in_bucket: idx as usize - entries,
-        }
+        let (bucket_idx, index_in_bucket) = BucketIndex::from_flat_index(idx as usize);
+        SlotIndex { bucket_idx, index_in_bucket }
     }
 
     // SAFETY: Buckets must be managed solely by functions here (i.e., get/put on SlotIndex) and
     // `self` comes from SlotIndex::from_index
     #[inline]
     unsafe fn get<V: Copy>(&self, buckets: &[AtomicPtr<Slot<V>>; 21]) -> Option<(V, u32)> {
-        // SAFETY: `bucket_idx` is ilog2(u32).saturating_sub(11), which is at most 21, i.e.,
-        // in-bounds of buckets. See `from_index` for computation.
-        let bucket = unsafe { buckets.get_unchecked(self.bucket_idx) };
+        let bucket = &buckets[self.bucket_idx];
         let ptr = bucket.load(Ordering::Acquire);
         // Bucket is not yet initialized: then we obviously won't find this entry in that bucket.
         if ptr.is_null() {
             return None;
         }
-        debug_assert!(self.index_in_bucket < self.entries());
+        debug_assert!(self.index_in_bucket < self.bucket_idx.capacity());
         // SAFETY: `bucket` was allocated (so <= isize in total bytes) to hold `entries`, so this
         // must be inbounds.
         let slot = unsafe { ptr.add(self.index_in_bucket) };
@@ -131,7 +119,7 @@ impl SlotIndex {
 
     #[cold]
     #[inline(never)]
-    fn initialize_bucket<V>(bucket: &AtomicPtr<Slot<V>>, bucket_idx: usize) -> *mut Slot<V> {
+    fn initialize_bucket<V>(bucket: &AtomicPtr<Slot<V>>, bucket_idx: BucketIndex) -> *mut Slot<V> {
         static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
 
         // If we are initializing the bucket, then acquire a global lock.
@@ -145,8 +133,8 @@ impl SlotIndex {
         // OK, now under the allocator lock, if we're still null then it's definitely us that will
         // initialize this bucket.
         if ptr.is_null() {
-            let bucket_len = SlotIndex { bucket_idx, index_in_bucket: 0 }.entries();
-            let bucket_layout = std::alloc::Layout::array::<Slot<V>>(bucket_len).unwrap();
+            let bucket_layout =
+                std::alloc::Layout::array::<Slot<V>>(bucket_idx.capacity()).unwrap();
             // This is more of a sanity check -- this code is very cold, so it's safe to pay a
             // little extra cost here.
             assert!(bucket_layout.size() > 0);
@@ -167,12 +155,10 @@ impl SlotIndex {
     /// Returns true if this successfully put into the map.
     #[inline]
     fn put<V>(&self, buckets: &[AtomicPtr<Slot<V>>; 21], value: V, extra: u32) -> bool {
-        // SAFETY: `bucket_idx` is ilog2(u32).saturating_sub(11), which is at most 21, i.e.,
-        // in-bounds of buckets.
-        let bucket = unsafe { buckets.get_unchecked(self.bucket_idx) };
+        let bucket = &buckets[self.bucket_idx];
         let ptr = self.bucket_ptr(bucket);
 
-        debug_assert!(self.index_in_bucket < self.entries());
+        debug_assert!(self.index_in_bucket < self.bucket_idx.capacity());
         // SAFETY: `bucket` was allocated (so <= isize in total bytes) to hold `entries`, so this
         // must be inbounds.
         let slot = unsafe { ptr.add(self.index_in_bucket) };
@@ -209,12 +195,10 @@ impl SlotIndex {
     /// Inserts into the map, given that the slot is unique, so it won't race with other threads.
     #[inline]
     unsafe fn put_unique<V>(&self, buckets: &[AtomicPtr<Slot<V>>; 21], value: V, extra: u32) {
-        // SAFETY: `bucket_idx` is ilog2(u32).saturating_sub(11), which is at most 21, i.e.,
-        // in-bounds of buckets.
-        let bucket = unsafe { buckets.get_unchecked(self.bucket_idx) };
+        let bucket = &buckets[self.bucket_idx];
         let ptr = self.bucket_ptr(bucket);
 
-        debug_assert!(self.index_in_bucket < self.entries());
+        debug_assert!(self.index_in_bucket < self.bucket_idx.capacity());
         // SAFETY: `bucket` was allocated (so <= isize in total bytes) to hold `entries`, so this
         // must be inbounds.
         let slot = unsafe { ptr.add(self.index_in_bucket) };
@@ -289,7 +273,7 @@ unsafe impl<K: Idx, #[may_dangle] V, I> Drop for VecCache<K, V, I> {
         assert!(!std::mem::needs_drop::<K>());
         assert!(!std::mem::needs_drop::<V>());
 
-        for (idx, bucket) in self.buckets.iter().enumerate() {
+        for (idx, bucket) in BucketIndex::enumerate_buckets(&self.buckets) {
             let bucket = bucket.load(Ordering::Acquire);
             if !bucket.is_null() {
                 let layout = std::alloc::Layout::array::<Slot<V>>(ENTRIES_BY_BUCKET[idx]).unwrap();
@@ -299,7 +283,7 @@ unsafe impl<K: Idx, #[may_dangle] V, I> Drop for VecCache<K, V, I> {
             }
         }
 
-        for (idx, bucket) in self.present.iter().enumerate() {
+        for (idx, bucket) in BucketIndex::enumerate_buckets(&self.present) {
             let bucket = bucket.load(Ordering::Acquire);
             if !bucket.is_null() {
                 let layout = std::alloc::Layout::array::<Slot<()>>(ENTRIES_BY_BUCKET[idx]).unwrap();
@@ -365,5 +349,163 @@ where
     }
 }
 
-#[cfg(test)]
-mod tests;
+/// Index into an array of buckets.
+///
+/// Using an enum lets us tell the compiler that values range from 0 to 20,
+/// without having to resort to pattern types or other unstable features.
+#[derive(Clone, Copy, PartialEq, Eq)]
+#[repr(usize)]
+enum BucketIndex {
+    // tidy-alphabetical-start
+    Bucket00,
+    Bucket01,
+    Bucket02,
+    Bucket03,
+    Bucket04,
+    Bucket05,
+    Bucket06,
+    Bucket07,
+    Bucket08,
+    Bucket09,
+    Bucket10,
+    Bucket11,
+    Bucket12,
+    Bucket13,
+    Bucket14,
+    Bucket15,
+    Bucket16,
+    Bucket17,
+    Bucket18,
+    Bucket19,
+    Bucket20,
+    // tidy-alphabetical-end
+}
+
+impl Debug for BucketIndex {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        Debug::fmt(&self.to_usize(), f)
+    }
+}
+
+impl BucketIndex {
+    /// Capacity of bucket 0 (and also of bucket 1).
+    const BUCKET_0_CAPACITY: usize = 1 << (Self::NONZERO_BUCKET_SHIFT_ADJUST + 1);
+    /// Adjustment factor from the highest-set-bit-position of a flat index,
+    /// to its corresponding bucket number.
+    ///
+    /// For example, the first flat-index in bucket 2 is 8192.
+    /// Its highest-set-bit-position is `(8192).ilog2() == 13`, and subtracting
+    /// the adjustment factor of 11 gives the bucket number of 2.
+    const NONZERO_BUCKET_SHIFT_ADJUST: usize = 11;
+
+    #[inline(always)]
+    const fn to_usize(self) -> usize {
+        self as usize
+    }
+
+    #[inline(always)]
+    const fn from_raw(raw: usize) -> Self {
+        match raw {
+            // tidy-alphabetical-start
+            00 => Self::Bucket00,
+            01 => Self::Bucket01,
+            02 => Self::Bucket02,
+            03 => Self::Bucket03,
+            04 => Self::Bucket04,
+            05 => Self::Bucket05,
+            06 => Self::Bucket06,
+            07 => Self::Bucket07,
+            08 => Self::Bucket08,
+            09 => Self::Bucket09,
+            10 => Self::Bucket10,
+            11 => Self::Bucket11,
+            12 => Self::Bucket12,
+            13 => Self::Bucket13,
+            14 => Self::Bucket14,
+            15 => Self::Bucket15,
+            16 => Self::Bucket16,
+            17 => Self::Bucket17,
+            18 => Self::Bucket18,
+            19 => Self::Bucket19,
+            20 => Self::Bucket20,
+            // tidy-alphabetical-end
+            _ => panic!("bucket index out of range"),
+        }
+    }
+
+    /// Total number of slots in this bucket.
+    #[inline(always)]
+    const fn capacity(self) -> usize {
+        match self {
+            Self::Bucket00 => Self::BUCKET_0_CAPACITY,
+            // Bucket 1 has a capacity of `1 << (1 + 11) == pow(2, 12) == 4096`.
+            // Bucket 2 has a capacity of `1 << (2 + 11) == pow(2, 13) == 8192`.
+            _ => 1 << (self.to_usize() + Self::NONZERO_BUCKET_SHIFT_ADJUST),
+        }
+    }
+
+    /// Converts a flat index in the range `0..=u32::MAX` into a bucket index,
+    /// and a slot offset within that bucket.
+    ///
+    /// Panics if `flat > u32::MAX`.
+    #[inline(always)]
+    const fn from_flat_index(flat: usize) -> (Self, usize) {
+        if flat > u32::MAX as usize {
+            panic!();
+        }
+
+        // If the index is in bucket 0, the conversion is trivial.
+        // This also avoids calling `ilog2` when `flat == 0`.
+        if flat < Self::BUCKET_0_CAPACITY {
+            return (Self::Bucket00, flat);
+        }
+
+        // General-case conversion for a non-zero bucket index.
+        //
+        //              | bucket |   slot
+        // flat | ilog2 |  index | offset
+        // ------------------------------
+        // 4096 |    12 |      1 |      0
+        // 4097 |    12 |      1 |      1
+        // ...
+        // 8191 |    12 |      1 |   4095
+        // 8192 |    13 |      2 |      0
+        let highest_bit_pos = flat.ilog2() as usize;
+        let bucket_index =
+            BucketIndex::from_raw(highest_bit_pos - Self::NONZERO_BUCKET_SHIFT_ADJUST);
+
+        // Clear the highest-set bit (which selects the bucket) to get the
+        // slot offset within this bucket.
+        let slot_offset = flat & !(1 << highest_bit_pos);
+
+        (bucket_index, slot_offset)
+    }
+
+    #[inline(always)]
+    fn iter_all() -> impl ExactSizeIterator<Item = Self> {
+        (0usize..BUCKETS).map(BucketIndex::from_raw)
+    }
+
+    #[inline(always)]
+    fn enumerate_buckets<T>(buckets: &[T; BUCKETS]) -> impl ExactSizeIterator<Item = (Self, &T)> {
+        BucketIndex::iter_all().zip(buckets)
+    }
+}
+
+impl<T> Index<BucketIndex> for [T; BUCKETS] {
+    type Output = T;
+
+    #[inline(always)]
+    fn index(&self, index: BucketIndex) -> &Self::Output {
+        // The optimizer should be able to see that see that a bucket index is
+        // always in-bounds, and omit the runtime bounds check.
+        &self[index.to_usize()]
+    }
+}
+
+impl<T> IndexMut<BucketIndex> for [T; BUCKETS] {
+    #[inline(always)]
+    fn index_mut(&mut self, index: BucketIndex) -> &mut Self::Output {
+        &mut self[index.to_usize()]
+    }
+}

compiler/rustc_data_structures/src/vec_cache/tests.rs

@@ -1,5 +1,39 @@
 use super::*;
 
+#[test]
+#[should_panic(expected = "bucket index out of range")]
+fn bucket_index_n_buckets() {
+    BucketIndex::from_raw(BUCKETS);
+}
+
+#[test]
+fn bucket_index_round_trip() {
+    for i in 0..BUCKETS {
+        assert_eq!(BucketIndex::from_raw(i).to_usize(), i);
+    }
+}
+
+#[test]
+fn bucket_index_iter_all_len() {
+    let len = BucketIndex::iter_all().len();
+    assert_eq!(len, BUCKETS);
+
+    let len = BucketIndex::iter_all().collect::<Vec<_>>().len();
+    assert_eq!(len, BUCKETS);
+
+    let len = BucketIndex::enumerate_buckets(&[(); BUCKETS]).len();
+    assert_eq!(len, BUCKETS);
+}
+
+#[test]
+fn bucket_index_capacity() {
+    let mut total = 0u128;
+    for i in BucketIndex::iter_all() {
+        total += u128::try_from(i.capacity()).unwrap();
+    }
+    assert_eq!(total, (u32::MAX as u128) + 1);
+}
+
 #[test]
 #[cfg(not(miri))]
 fn vec_cache_empty() {
@@ -70,8 +104,8 @@ fn slot_entries_table() {
 
 #[test]
 fn bucket_entries_matches() {
-    for i in 0..BUCKETS {
-        assert_eq!(SlotIndex { bucket_idx: i, index_in_bucket: 0 }.entries(), ENTRIES_BY_BUCKET[i]);
+    for i in BucketIndex::iter_all() {
+        assert_eq!(i.capacity(), ENTRIES_BY_BUCKET[i]);
     }
 }
 
@@ -84,22 +118,22 @@ fn slot_index_exhaustive() {
     }
     let slot_idx = SlotIndex::from_index(0);
     assert_eq!(slot_idx.index_in_bucket, 0);
-    assert_eq!(slot_idx.bucket_idx, 0);
+    assert_eq!(slot_idx.bucket_idx, BucketIndex::Bucket00);
     let mut prev = slot_idx;
     for idx in 1..=u32::MAX {
         let slot_idx = SlotIndex::from_index(idx);
 
         // SAFETY: Ensure indices don't go out of bounds of buckets.
-        assert!(slot_idx.index_in_bucket < slot_idx.entries());
+        assert!(slot_idx.index_in_bucket < slot_idx.bucket_idx.capacity());
 
         if prev.bucket_idx == slot_idx.bucket_idx {
             assert_eq!(prev.index_in_bucket + 1, slot_idx.index_in_bucket);
         } else {
             assert_eq!(slot_idx.index_in_bucket, 0);
         }
 
-        assert_eq!(buckets[slot_idx.bucket_idx], slot_idx.entries() as u32);
-        assert_eq!(ENTRIES_BY_BUCKET[slot_idx.bucket_idx], slot_idx.entries(), "{}", idx);
+        assert_eq!(buckets[slot_idx.bucket_idx], slot_idx.bucket_idx.capacity() as u32);
+        assert_eq!(ENTRIES_BY_BUCKET[slot_idx.bucket_idx], slot_idx.bucket_idx.capacity(), "{idx}",);
 
         prev = slot_idx;
     }