turbo_tasks/
no_move_vec.rs

1use std::{
2    ptr::null_mut,
3    slice::from_raw_parts_mut,
4    sync::{
5        Mutex,
6        atomic::{AtomicPtr, Ordering},
7    },
8};
9
10const BUCKETS: usize = (usize::BITS + 1) as usize;
11
12/// An `Option`-like type that guarantees that a fully zeroed value is a valid
13/// `None` variant.
14#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
15#[repr(u8)]
16enum COption<T> {
17    // TODO(alexkirsz) We need a way to guarantee that a fully zeroed value is a
18    // valid `None` variant. This is theoretically possible when the wrapped
19    // type has no valid value that can be represented by all zeros, but there
20    // is no way to enforce this at the type level. For now, we just use a custom
21    // option type with explicit discriminant for the `None` variant.
22    // The issue with this implementation is that it disables niche optimization.
23    None = 0,
24    Some(T),
25}
26
27impl<T> Default for COption<T> {
28    fn default() -> Self {
29        Self::None
30    }
31}
32
33impl<T> COption<T> {
34    /// Returns a slice of the given size filled with the `None` variant.
35    fn new_none_slice(size: usize) -> Box<[Self]> {
36        let slice = Box::<[COption<T>]>::new_zeroed_slice(size);
37        // Safety:
38        // We know that a zeroed COption<T> is a valid COption::None value.
39        unsafe { slice.assume_init() }
40    }
41
42    /// Returns a reference to the contained value, or `None` if it is `None`.
43    fn as_option_ref(&self) -> Option<&T> {
44        match self {
45            COption::None => None,
46            COption::Some(t) => Some(t),
47        }
48    }
49}
50
51pub struct NoMoveVec<T, const INITIAL_CAPACITY_BITS: u32 = 6> {
52    buckets: [(AtomicPtr<COption<T>>, Mutex<()>); BUCKETS],
53}
54
55fn get_bucket_index<const INITIAL_CAPACITY_BITS: u32>(idx: usize) -> u32 {
56    (usize::BITS - idx.leading_zeros()).saturating_sub(INITIAL_CAPACITY_BITS)
57}
58
59fn get_bucket_size<const INITIAL_CAPACITY_BITS: u32>(bucket_index: u32) -> usize {
60    if bucket_index != 0 {
61        1 << (bucket_index + INITIAL_CAPACITY_BITS - 1)
62    } else {
63        1 << INITIAL_CAPACITY_BITS
64    }
65}
66
67fn get_index_in_bucket<const INITIAL_CAPACITY_BITS: u32>(idx: usize, bucket_index: u32) -> usize {
68    if bucket_index != 0 {
69        idx ^ (1 << (bucket_index + INITIAL_CAPACITY_BITS - 1))
70    } else {
71        idx
72    }
73}
74
75/// Allocates a new bucket of `COption<T>`s, all initialized to `None`.
76fn allocate_bucket<const INITIAL_CAPACITY_BITS: u32, T>(bucket_index: u32) -> *mut COption<T> {
77    let size = get_bucket_size::<INITIAL_CAPACITY_BITS>(bucket_index);
78    let slice = COption::<T>::new_none_slice(size);
79    Box::into_raw(slice) as *mut COption<T>
80}
81
82impl<T, const INITIAL_CAPACITY_BITS: u32> Default for NoMoveVec<T, INITIAL_CAPACITY_BITS> {
83    fn default() -> Self {
84        Self::new()
85    }
86}
87
88impl<T, const INITIAL_CAPACITY_BITS: u32> NoMoveVec<T, INITIAL_CAPACITY_BITS> {
89    pub fn new() -> Self {
90        let mut buckets = [null_mut(); BUCKETS];
91        buckets[0] = allocate_bucket::<INITIAL_CAPACITY_BITS, T>(0);
92        let buckets = buckets.map(|p| (AtomicPtr::new(p), Mutex::new(())));
93        NoMoveVec { buckets }
94    }
95
96    pub fn get(&self, idx: usize) -> Option<&T> {
97        let bucket_idx = get_bucket_index::<INITIAL_CAPACITY_BITS>(idx);
98        let bucket_ptr = unsafe { self.buckets.get_unchecked(bucket_idx as usize) }
99            .0
100            .load(Ordering::Acquire);
101        if bucket_ptr.is_null() {
102            return None;
103        }
104        let index = get_index_in_bucket::<INITIAL_CAPACITY_BITS>(idx, bucket_idx);
105        unsafe { &*bucket_ptr.add(index) }.as_option_ref()
106    }
107
108    /// # Safety
109    /// There must not be a concurrent operation to this idx
110    pub unsafe fn insert(&self, idx: usize, value: T) -> &T {
111        let bucket_idx = get_bucket_index::<INITIAL_CAPACITY_BITS>(idx);
112        let bucket = unsafe { self.buckets.get_unchecked(bucket_idx as usize) };
113        // SAFETY: This is safe to be relaxed as the bucket will never become null
114        // again. We perform a acquire load when it's null.
115        let mut bucket_ptr = bucket.0.load(Ordering::Relaxed);
116        if bucket_ptr.is_null() {
117            bucket_ptr = bucket.0.load(Ordering::Acquire);
118            if bucket_ptr.is_null() {
119                let lock = bucket.1.lock();
120                let guarded_bucket_ptr = bucket.0.load(Ordering::Acquire);
121                if guarded_bucket_ptr.is_null() {
122                    let new_bucket = allocate_bucket::<INITIAL_CAPACITY_BITS, T>(bucket_idx);
123                    bucket_ptr = match bucket.0.compare_exchange(
124                        null_mut(),
125                        new_bucket,
126                        Ordering::AcqRel,
127                        Ordering::Relaxed,
128                    ) {
129                        Ok(_) => new_bucket,
130                        Err(current_bucket) => {
131                            drop(unsafe { Box::from_raw(new_bucket) });
132                            current_bucket
133                        }
134                    };
135                    drop(lock);
136                } else {
137                    bucket_ptr = guarded_bucket_ptr;
138                }
139            }
140        }
141        let index = get_index_in_bucket::<INITIAL_CAPACITY_BITS>(idx, bucket_idx);
142        let item = unsafe { &mut *bucket_ptr.add(index) };
143        *item = COption::Some(value);
144        // To sync with any acquire load of the bucket ptr
145        bucket.0.store(bucket_ptr, Ordering::Release);
146        item.as_option_ref().unwrap()
147    }
148}
149
150impl<T, const INITIAL_CAPACITY_BITS: u32> Drop for NoMoveVec<T, INITIAL_CAPACITY_BITS> {
151    fn drop(&mut self) {
152        for (bucket_index, (bucket, _)) in self.buckets.iter_mut().enumerate() {
153            if bucket_index < (usize::BITS + 1 - INITIAL_CAPACITY_BITS) as usize {
154                let bucket_size = get_bucket_size::<INITIAL_CAPACITY_BITS>(bucket_index as u32);
155                let bucket_ptr = *bucket.get_mut();
156
157                if !bucket_ptr.is_null() {
158                    drop(unsafe { Box::from_raw(from_raw_parts_mut(bucket_ptr, bucket_size)) });
159                }
160            }
161        }
162    }
163}
164
165#[cfg(test)]
166mod tests {
167    use super::NoMoveVec;
168
169    #[test]
170    fn basic_operations() {
171        let v = NoMoveVec::<(usize, usize)>::new();
172        assert_eq!(v.get(0), None);
173        assert_eq!(v.get(1), None);
174        assert_eq!(v.get(8), None);
175        assert_eq!(v.get(9), None);
176        assert_eq!(v.get(15), None);
177        assert_eq!(v.get(16), None);
178        assert_eq!(v.get(100), None);
179        assert_eq!(v.get(1000), None);
180
181        for i in 0..1000 {
182            unsafe {
183                v.insert(i, (i, i));
184            }
185            assert_eq!(v.get(i), Some(&(i, i)));
186        }
187        for i in 0..1000 {
188            assert_eq!(v.get(i), Some(&(i, i)));
189        }
190        assert_eq!(v.get(1001), None);
191
192        unsafe {
193            v.insert(1000000, (0, 0));
194        }
195        assert_eq!(v.get(1000000), Some(&(0, 0)));
196        assert_eq!(v.get(10000), None);
197    }
198}
turbo_tasks/no_move_vec.rs

turbo_tasks/
no_move_vec.rs
