66//!
77//! This is currently used for query caching.
88
9- use std:: fmt:: Debug ;
9+ use std:: fmt:: { self , Debug } ;
1010use std:: marker:: PhantomData ;
11+ use std:: ops:: { Index , IndexMut } ;
1112use std:: sync:: atomic:: { AtomicPtr , AtomicU32 , AtomicUsize , Ordering } ;
1213
1314use rustc_index:: Idx ;
1415
16+ #[ cfg( test) ]
17+ mod tests;
18+
1519struct Slot < V > {
1620 // We never construct &Slot<V> so it's fine for this to not be in an UnsafeCell.
1721 value : V ,
@@ -28,7 +32,7 @@ struct Slot<V> {
2832#[ derive( Copy , Clone , Debug ) ]
2933struct SlotIndex {
3034 // the index of the bucket in VecCache (0 to 20)
31- bucket_idx : usize ,
35+ bucket_idx : BucketIndex ,
3236 // the index of the slot within the bucket
3337 index_in_bucket : usize ,
3438}
@@ -42,7 +46,7 @@ const ENTRIES_BY_BUCKET: [usize; BUCKETS] = {
4246 let mut key = 0 ;
4347 loop {
4448 let si = SlotIndex :: from_index ( key) ;
45- entries[ si. bucket_idx ] = si. entries ( ) ;
49+ entries[ si. bucket_idx . to_usize ( ) ] = si. bucket_idx . capacity ( ) ;
4650 if key == 0 {
4751 key = 1 ;
4852 } else if key == ( 1 << 31 ) {
@@ -57,48 +61,24 @@ const ENTRIES_BY_BUCKET: [usize; BUCKETS] = {
5761const BUCKETS : usize = 21 ;
5862
5963impl SlotIndex {
60- /// The total possible number of entries in the bucket
61- const fn entries ( & self ) -> usize {
62- if self . bucket_idx == 0 { 1 << 12 } else { 1 << ( self . bucket_idx + 11 ) }
63- }
64-
65- // This unpacks a flat u32 index into identifying which bucket it belongs to and the offset
66- // within that bucket. As noted in the VecCache docs, buckets double in size with each index.
67- // Typically that would mean 31 buckets (2^0 + 2^1 ... + 2^31 = u32::MAX - 1), but to reduce
68- // the size of the VecCache struct and avoid uselessly small allocations, we instead have the
69- // first bucket have 2**12 entries. To simplify the math, the second bucket also 2**12 entries,
70- // and buckets double from there.
71- //
72- // We assert that [0, 2**32 - 1] uniquely map through this function to individual, consecutive
73- // slots (see `slot_index_exhaustive` in tests).
64+ /// Unpacks a flat 32-bit index into a [`BucketIndex`] and a slot offset within that bucket.
7465 #[ inline]
7566 const fn from_index ( idx : u32 ) -> Self {
76- const FIRST_BUCKET_SHIFT : usize = 12 ;
77- if idx < ( 1 << FIRST_BUCKET_SHIFT ) {
78- return SlotIndex { bucket_idx : 0 , index_in_bucket : idx as usize } ;
79- }
80- // We already ruled out idx 0, so this `ilog2` never panics (and the check optimizes away)
81- let bucket = idx. ilog2 ( ) as usize ;
82- let entries = 1 << bucket;
83- SlotIndex {
84- bucket_idx : bucket - FIRST_BUCKET_SHIFT + 1 ,
85- index_in_bucket : idx as usize - entries,
86- }
67+ let ( bucket_idx, index_in_bucket) = BucketIndex :: from_flat_index ( idx as usize ) ;
68+ SlotIndex { bucket_idx, index_in_bucket }
8769 }
8870
8971 // SAFETY: Buckets must be managed solely by functions here (i.e., get/put on SlotIndex) and
9072 // `self` comes from SlotIndex::from_index
9173 #[ inline]
9274 unsafe fn get < V : Copy > ( & self , buckets : & [ AtomicPtr < Slot < V > > ; 21 ] ) -> Option < ( V , u32 ) > {
93- // SAFETY: `bucket_idx` is ilog2(u32).saturating_sub(11), which is at most 21, i.e.,
94- // in-bounds of buckets. See `from_index` for computation.
95- let bucket = unsafe { buckets. get_unchecked ( self . bucket_idx ) } ;
75+ let bucket = & buckets[ self . bucket_idx ] ;
9676 let ptr = bucket. load ( Ordering :: Acquire ) ;
9777 // Bucket is not yet initialized: then we obviously won't find this entry in that bucket.
9878 if ptr. is_null ( ) {
9979 return None ;
10080 }
101- debug_assert ! ( self . index_in_bucket < self . entries ( ) ) ;
81+ debug_assert ! ( self . index_in_bucket < self . bucket_idx . capacity ( ) ) ;
10282 // SAFETY: `bucket` was allocated (so <= isize in total bytes) to hold `entries`, so this
10383 // must be inbounds.
10484 let slot = unsafe { ptr. add ( self . index_in_bucket ) } ;
@@ -131,7 +111,7 @@ impl SlotIndex {
131111
132112 #[ cold]
133113 #[ inline( never) ]
134- fn initialize_bucket < V > ( bucket : & AtomicPtr < Slot < V > > , bucket_idx : usize ) -> * mut Slot < V > {
114+ fn initialize_bucket < V > ( bucket : & AtomicPtr < Slot < V > > , bucket_idx : BucketIndex ) -> * mut Slot < V > {
135115 static LOCK : std:: sync:: Mutex < ( ) > = std:: sync:: Mutex :: new ( ( ) ) ;
136116
137117 // If we are initializing the bucket, then acquire a global lock.
@@ -145,8 +125,8 @@ impl SlotIndex {
145125 // OK, now under the allocator lock, if we're still null then it's definitely us that will
146126 // initialize this bucket.
147127 if ptr. is_null ( ) {
148- let bucket_len = SlotIndex { bucket_idx , index_in_bucket : 0 } . entries ( ) ;
149- let bucket_layout = std:: alloc:: Layout :: array :: < Slot < V > > ( bucket_len ) . unwrap ( ) ;
128+ let bucket_layout =
129+ std:: alloc:: Layout :: array :: < Slot < V > > ( bucket_idx . capacity ( ) ) . unwrap ( ) ;
150130 // This is more of a sanity check -- this code is very cold, so it's safe to pay a
151131 // little extra cost here.
152132 assert ! ( bucket_layout. size( ) > 0 ) ;
@@ -167,12 +147,10 @@ impl SlotIndex {
167147 /// Returns true if this successfully put into the map.
168148 #[ inline]
169149 fn put < V > ( & self , buckets : & [ AtomicPtr < Slot < V > > ; 21 ] , value : V , extra : u32 ) -> bool {
170- // SAFETY: `bucket_idx` is ilog2(u32).saturating_sub(11), which is at most 21, i.e.,
171- // in-bounds of buckets.
172- let bucket = unsafe { buckets. get_unchecked ( self . bucket_idx ) } ;
150+ let bucket = & buckets[ self . bucket_idx ] ;
173151 let ptr = self . bucket_ptr ( bucket) ;
174152
175- debug_assert ! ( self . index_in_bucket < self . entries ( ) ) ;
153+ debug_assert ! ( self . index_in_bucket < self . bucket_idx . capacity ( ) ) ;
176154 // SAFETY: `bucket` was allocated (so <= isize in total bytes) to hold `entries`, so this
177155 // must be inbounds.
178156 let slot = unsafe { ptr. add ( self . index_in_bucket ) } ;
@@ -209,12 +187,10 @@ impl SlotIndex {
209187 /// Inserts into the map, given that the slot is unique, so it won't race with other threads.
210188 #[ inline]
211189 unsafe fn put_unique < V > ( & self , buckets : & [ AtomicPtr < Slot < V > > ; 21 ] , value : V , extra : u32 ) {
212- // SAFETY: `bucket_idx` is ilog2(u32).saturating_sub(11), which is at most 21, i.e.,
213- // in-bounds of buckets.
214- let bucket = unsafe { buckets. get_unchecked ( self . bucket_idx ) } ;
190+ let bucket = & buckets[ self . bucket_idx ] ;
215191 let ptr = self . bucket_ptr ( bucket) ;
216192
217- debug_assert ! ( self . index_in_bucket < self . entries ( ) ) ;
193+ debug_assert ! ( self . index_in_bucket < self . bucket_idx . capacity ( ) ) ;
218194 // SAFETY: `bucket` was allocated (so <= isize in total bytes) to hold `entries`, so this
219195 // must be inbounds.
220196 let slot = unsafe { ptr. add ( self . index_in_bucket ) } ;
@@ -254,7 +230,7 @@ pub struct VecCache<K: Idx, V, I> {
254230 // ...
255231 // Bucket 19: 1073741824
256232 // Bucket 20: 2147483648
257- // The total number of entries if all buckets are initialized is u32::MAX-1 .
233+ // The total number of entries if all buckets are initialized is 2^32 .
258234 buckets : [ AtomicPtr < Slot < V > > ; BUCKETS ] ,
259235
260236 // In the compiler's current usage these are only *read* during incremental and self-profiling.
@@ -289,7 +265,7 @@ unsafe impl<K: Idx, #[may_dangle] V, I> Drop for VecCache<K, V, I> {
289265 assert ! ( !std:: mem:: needs_drop:: <K >( ) ) ;
290266 assert ! ( !std:: mem:: needs_drop:: <V >( ) ) ;
291267
292- for ( idx, bucket) in self . buckets . iter ( ) . enumerate ( ) {
268+ for ( idx, bucket) in BucketIndex :: enumerate_buckets ( & self . buckets ) {
293269 let bucket = bucket. load ( Ordering :: Acquire ) ;
294270 if !bucket. is_null ( ) {
295271 let layout = std:: alloc:: Layout :: array :: < Slot < V > > ( ENTRIES_BY_BUCKET [ idx] ) . unwrap ( ) ;
@@ -299,7 +275,7 @@ unsafe impl<K: Idx, #[may_dangle] V, I> Drop for VecCache<K, V, I> {
299275 }
300276 }
301277
302- for ( idx, bucket) in self . present . iter ( ) . enumerate ( ) {
278+ for ( idx, bucket) in BucketIndex :: enumerate_buckets ( & self . present ) {
303279 let bucket = bucket. load ( Ordering :: Acquire ) ;
304280 if !bucket. is_null ( ) {
305281 let layout = std:: alloc:: Layout :: array :: < Slot < ( ) > > ( ENTRIES_BY_BUCKET [ idx] ) . unwrap ( ) ;
@@ -365,5 +341,164 @@ where
365341 }
366342}
367343
368- #[ cfg( test) ]
369- mod tests;
344+ /// Index into an array of buckets.
345+ ///
346+ /// Using an enum lets us tell the compiler that values range from 0 to 20,
347+ /// allowing array bounds checks to be optimized away,
348+ /// without having to resort to pattern types or other unstable features.
349+ #[ derive( Clone , Copy , PartialEq , Eq ) ]
350+ #[ repr( usize ) ]
351+ enum BucketIndex {
352+ // tidy-alphabetical-start
353+ Bucket00 ,
354+ Bucket01 ,
355+ Bucket02 ,
356+ Bucket03 ,
357+ Bucket04 ,
358+ Bucket05 ,
359+ Bucket06 ,
360+ Bucket07 ,
361+ Bucket08 ,
362+ Bucket09 ,
363+ Bucket10 ,
364+ Bucket11 ,
365+ Bucket12 ,
366+ Bucket13 ,
367+ Bucket14 ,
368+ Bucket15 ,
369+ Bucket16 ,
370+ Bucket17 ,
371+ Bucket18 ,
372+ Bucket19 ,
373+ Bucket20 ,
374+ // tidy-alphabetical-end
375+ }
376+
377+ impl Debug for BucketIndex {
378+ fn fmt ( & self , f : & mut fmt:: Formatter < ' _ > ) -> fmt:: Result {
379+ Debug :: fmt ( & self . to_usize ( ) , f)
380+ }
381+ }
382+
383+ impl BucketIndex {
384+ /// Capacity of bucket 0 (and also of bucket 1).
385+ const BUCKET_0_CAPACITY : usize = 1 << ( Self :: NONZERO_BUCKET_SHIFT_ADJUST + 1 ) ;
386+ /// Adjustment factor from the highest-set-bit-position of a flat index,
387+ /// to its corresponding bucket number.
388+ ///
389+ /// For example, the first flat-index in bucket 2 is 8192.
390+ /// Its highest-set-bit-position is `(8192).ilog2() == 13`, and subtracting
391+ /// the adjustment factor of 11 gives the bucket number of 2.
392+ const NONZERO_BUCKET_SHIFT_ADJUST : usize = 11 ;
393+
394+ #[ inline( always) ]
395+ const fn to_usize ( self ) -> usize {
396+ self as usize
397+ }
398+
399+ #[ inline( always) ]
400+ const fn from_raw ( raw : usize ) -> Self {
401+ match raw {
402+ // tidy-alphabetical-start
403+ 00 => Self :: Bucket00 ,
404+ 01 => Self :: Bucket01 ,
405+ 02 => Self :: Bucket02 ,
406+ 03 => Self :: Bucket03 ,
407+ 04 => Self :: Bucket04 ,
408+ 05 => Self :: Bucket05 ,
409+ 06 => Self :: Bucket06 ,
410+ 07 => Self :: Bucket07 ,
411+ 08 => Self :: Bucket08 ,
412+ 09 => Self :: Bucket09 ,
413+ 10 => Self :: Bucket10 ,
414+ 11 => Self :: Bucket11 ,
415+ 12 => Self :: Bucket12 ,
416+ 13 => Self :: Bucket13 ,
417+ 14 => Self :: Bucket14 ,
418+ 15 => Self :: Bucket15 ,
419+ 16 => Self :: Bucket16 ,
420+ 17 => Self :: Bucket17 ,
421+ 18 => Self :: Bucket18 ,
422+ 19 => Self :: Bucket19 ,
423+ 20 => Self :: Bucket20 ,
424+ // tidy-alphabetical-end
425+ _ => panic ! ( "bucket index out of range" ) ,
426+ }
427+ }
428+
429+ /// Total number of slots in this bucket.
430+ #[ inline( always) ]
431+ const fn capacity ( self ) -> usize {
432+ match self {
433+ Self :: Bucket00 => Self :: BUCKET_0_CAPACITY ,
434+ // Bucket 1 has a capacity of `1 << (1 + 11) == pow(2, 12) == 4096`.
435+ // Bucket 2 has a capacity of `1 << (2 + 11) == pow(2, 13) == 8192`.
436+ _ => 1 << ( self . to_usize ( ) + Self :: NONZERO_BUCKET_SHIFT_ADJUST ) ,
437+ }
438+ }
439+
440+ /// Converts a flat index in the range `0..=u32::MAX` into a bucket index,
441+ /// and a slot offset within that bucket.
442+ ///
443+ /// Panics if `flat > u32::MAX`.
444+ #[ inline( always) ]
445+ const fn from_flat_index ( flat : usize ) -> ( Self , usize ) {
446+ if flat > u32:: MAX as usize {
447+ panic ! ( ) ;
448+ }
449+
450+ // If the index is in bucket 0, the conversion is trivial.
451+ // This also avoids calling `ilog2` when `flat == 0`.
452+ if flat < Self :: BUCKET_0_CAPACITY {
453+ return ( Self :: Bucket00 , flat) ;
454+ }
455+
456+ // General-case conversion for a non-zero bucket index.
457+ //
458+ // | bucket | slot
459+ // flat | ilog2 | index | offset
460+ // ------------------------------
461+ // 4096 | 12 | 1 | 0
462+ // 4097 | 12 | 1 | 1
463+ // ...
464+ // 8191 | 12 | 1 | 4095
465+ // 8192 | 13 | 2 | 0
466+ let highest_bit_pos = flat. ilog2 ( ) as usize ;
467+ let bucket_index =
468+ BucketIndex :: from_raw ( highest_bit_pos - Self :: NONZERO_BUCKET_SHIFT_ADJUST ) ;
469+
470+ // Clear the highest-set bit (which selects the bucket) to get the
471+ // slot offset within this bucket.
472+ let slot_offset = flat - ( 1 << highest_bit_pos) ;
473+
474+ ( bucket_index, slot_offset)
475+ }
476+
477+ #[ inline( always) ]
478+ fn iter_all ( ) -> impl ExactSizeIterator < Item = Self > {
479+ ( 0usize ..BUCKETS ) . map ( BucketIndex :: from_raw)
480+ }
481+
482+ #[ inline( always) ]
483+ fn enumerate_buckets < T > ( buckets : & [ T ; BUCKETS ] ) -> impl ExactSizeIterator < Item = ( Self , & T ) > {
484+ BucketIndex :: iter_all ( ) . zip ( buckets)
485+ }
486+ }
487+
488+ impl < T > Index < BucketIndex > for [ T ; BUCKETS ] {
489+ type Output = T ;
490+
491+ #[ inline( always) ]
492+ fn index ( & self , index : BucketIndex ) -> & Self :: Output {
493+ // The optimizer should be able to see that see that a bucket index is
494+ // always in-bounds, and omit the runtime bounds check.
495+ & self [ index. to_usize ( ) ]
496+ }
497+ }
498+
499+ impl < T > IndexMut < BucketIndex > for [ T ; BUCKETS ] {
500+ #[ inline( always) ]
501+ fn index_mut ( & mut self , index : BucketIndex ) -> & mut Self :: Output {
502+ & mut self [ index. to_usize ( ) ]
503+ }
504+ }
0 commit comments