1
//! Authors: Maurice Laveaux, Flip van Spaendonck and Jan Friso Groote
2

            
3
use std::alloc;
4
use std::alloc::Layout;
5
use std::cmp::max;
6
use std::hint::spin_loop;
7
use std::ptr;
8
use std::ptr::NonNull;
9
use std::sync::atomic::AtomicUsize;
10
use std::sync::atomic::Ordering;
11

            
12
use crate::BfSharedMutex;
13

            
14
/// An implementation of [`Vec<T>`] based on the [BfSharedMutex] implementation
15
/// that can be safely sent between threads. Elements can be appended
16
/// concurrently from multiple threads holding a [`BfVec::share`] of the vector.
17
///
18
/// Zero-sized element types are not supported.
19
pub struct BfVec<T> {
20
    shared: BfSharedMutex<BfVecShared<T>>,
21
}
22

            
23
/// The internal shared data of the [BfVec].
24
pub struct BfVecShared<T> {
25
    buffer: Option<NonNull<T>>,
26
    capacity: usize,
27

            
28
    /// The number of slots reserved by writers; can exceed `len` while pushes
29
    /// are in flight.
30
    reserved: AtomicUsize,
31

            
32
    /// The number of initialised elements; reads may only access slots below
33
    /// this bound.
34
    len: AtomicUsize,
35
}
36

            
37
impl<T> BfVec<T> {
38
    /// Create a new vector with zero capacity.
39
2
    pub fn new() -> BfVec<T> {
40
        const {
41
            assert!(std::mem::size_of::<T>() > 0, "Zero-sized types are not supported");
42
        }
43

            
44
2
        BfVec {
45
2
            shared: BfSharedMutex::new(BfVecShared::<T> {
46
2
                buffer: None,
47
2
                capacity: 0,
48
2
                reserved: AtomicUsize::new(0),
49
2
                len: AtomicUsize::new(0),
50
2
            }),
51
2
        }
52
2
    }
53

            
54
    /// Append a new element to the vector.
55
100003
    pub fn push(&self, value: T) {
56
100003
        let mut read = self.shared.read().unwrap();
57

            
58
        // Reserve an index for the new element with a compare-and-swap so that a reservation
59
        // never has to be rolled back; a rollback could hand the same slot to two threads.
60
100003
        let last_index = loop {
61
116477
            let current = read.reserved.load(Ordering::Relaxed);
62

            
63
116477
            if current >= read.capacity {
64
                // Vector needs to be resized.
65
23
                let new_capacity = max(read.capacity * 2, 8);
66

            
67
23
                drop(read);
68
23
                self.reserve(new_capacity);
69

            
70
                // Acquire read access and try a new position.
71
23
                read = self.shared.read().unwrap();
72
23
                continue;
73
116454
            }
74

            
75
116454
            if read
76
116454
                .reserved
77
116454
                .compare_exchange_weak(current, current + 1, Ordering::Relaxed, Ordering::Relaxed)
78
116454
                .is_ok()
79
            {
80
100003
                break current;
81
16451
            }
82
        };
83

            
84
        // Write the element on the specified index.
85

            
86
        // SAFETY: The CAS above reserved `last_index < capacity` while holding
87
        // the read lock, so the buffer is allocated and the slot is exclusively
88
        // ours and uninitialised.
89
100003
        unsafe {
90
100003
            let end = read.buffer.unwrap().as_ptr().add(last_index);
91
100003
            ptr::write(end, value);
92
100003
        }
93

            
94
        // Publish the element. `len` may only be raised past our slot once all
95
        // earlier slots have been initialised, so wait for the slots below us.
96
783449830
        while read
97
783449830
            .len
98
783449830
            .compare_exchange_weak(last_index, last_index + 1, Ordering::Release, Ordering::Relaxed)
99
783449830
            .is_err()
100
783349827
        {
101
783349827
            spin_loop();
102
783349827
        }
103
100003
    }
104

            
105
    /// Obtain another view on the vector to share among threads.
106
10
    pub fn share(&self) -> BfVec<T> {
107
10
        BfVec {
108
10
            shared: self.shared.clone(),
109
10
        }
110
10
    }
111

            
112
    /// Obtain the number of elements in the vector.
113
4
    pub fn len(&self) -> usize {
114
4
        self.shared.read().unwrap().len.load(Ordering::Relaxed)
115
4
    }
116

            
117
    /// Returns true iff the vector is empty.
118
1
    pub fn is_empty(&self) -> bool {
119
1
        self.len() == 0
120
1
    }
121

            
122
    /// Returns a clone of the element at the given index.
123
    ///
124
    /// The value is cloned while holding the read lock; handing out a
125
    /// reference instead would let it dangle when another thread's `push`
126
    /// reallocates the buffer.
127
100001
    pub fn at(&self, index: usize) -> T
128
100001
    where
129
100001
        T: Clone,
130
    {
131
100001
        let read = self.shared.read().unwrap();
132

            
133
        // Acquire pairs with the publishing store in `push`, so the slots below `len` are initialised.
134
100001
        assert!(
135
100001
            index < read.len.load(Ordering::Acquire),
136
            "Index {index} is out of bounds"
137
        );
138

            
139
        // SAFETY: `index < len` (Acquire) guarantees the slot is initialised
140
        // and published, and the read lock keeps the buffer alive for the
141
        // clone.
142
100001
        unsafe { (*read.buffer.unwrap().as_ptr().add(index)).clone() }
143
100001
    }
144

            
145
    /// Drops the elements in the Vec, but keeps the capacity.
146
1
    pub fn clear(&self) {
147
1
        let mut write = self.shared.write().unwrap();
148
1
        write.clear();
149
1
    }
150

            
151
    /// Reserve the given capacity.
152
23
    fn reserve(&self, capacity: usize) {
153
23
        let mut write = self.shared.write().unwrap();
154

            
155
        // A reserve could have happened in the meantime which makes this call obsolete
156
23
        if capacity <= write.capacity {
157
7
            return;
158
16
        }
159

            
160
        // Reservation, write and publish all happen under the read lock, so once the write lock
161
        // is held no push is in flight and every reserved slot has been initialised.
162
16
        debug_assert_eq!(
163
16
            write.reserved.load(Ordering::Relaxed),
164
16
            write.len.load(Ordering::Relaxed),
165
            "All reserved slots must be published before a resize"
166
        );
167

            
168
16
        let old_layout = Layout::array::<T>(write.capacity).unwrap();
169
16
        let layout = Layout::array::<T>(capacity).unwrap();
170

            
171
        // SAFETY: Held under the write lock, so no push is in flight and every reserved slot is
172
        // initialised. We allocate a larger buffer, move the `len` live elements into it, and
173
        // free the old allocation with the layout it was created with.
174
        unsafe {
175
16
            let new_buffer = alloc::alloc(layout) as *mut T;
176
16
            if new_buffer.is_null() {
177
                alloc::handle_alloc_error(layout);
178
16
            }
179

            
180
16
            if let Some(old_buffer) = write.buffer {
181
14
                debug_assert!(
182
14
                    write.len.load(Ordering::Relaxed) <= write.capacity,
183
                    "Length {} should be less than capacity {}",
184
                    write.len.load(Ordering::Relaxed),
185
                    write.capacity
186
                );
187

            
188
14
                ptr::copy_nonoverlapping(old_buffer.as_ptr(), new_buffer, write.len.load(Ordering::Relaxed));
189

            
190
                // Clean up the old buffer.
191
14
                alloc::dealloc(old_buffer.as_ptr() as *mut u8, old_layout);
192
2
            }
193

            
194
16
            write.capacity = capacity;
195
16
            write.buffer = NonNull::new(new_buffer);
196
        }
197
23
    }
198
}
199

            
200
impl<T> Default for BfVec<T> {
201
    fn default() -> Self {
202
        Self::new()
203
    }
204
}
205

            
206
impl<T> BfVecShared<T> {
207
    /// Clears the vector by dropping all elements.
208
3
    pub fn clear(&mut self) {
209
        // Only drop items within the 0..len range since the other values are not initialised.
210
100003
        for i in 0..self.len.load(Ordering::Relaxed) {
211
            // SAFETY: `&mut self` gives exclusive access, and slots `0..len` are initialised, so
212
            // each element can be dropped exactly once in place.
213
100003
            unsafe {
214
100003
                let ptr = self.buffer.unwrap().as_ptr().add(i);
215
100003

            
216
100003
                ptr::drop_in_place(ptr);
217
100003
            }
218
        }
219

            
220
        // Reset the length so the dropped elements cannot be observed or dropped a second time.
221
3
        self.len.store(0, Ordering::Relaxed);
222
3
        self.reserved.store(0, Ordering::Relaxed);
223
3
    }
224
}
225

            
226
impl<T> Drop for BfVecShared<T> {
227
2
    fn drop(&mut self) {
228
2
        self.clear();
229

            
230
        // SAFETY: `clear` dropped every live element, so the buffer holds no initialised values;
231
        // we free it with the same layout it was allocated with.
232
        unsafe {
233
2
            let layout = Layout::array::<T>(self.capacity).unwrap();
234
2
            if let Some(buffer) = self.buffer {
235
2
                alloc::dealloc(buffer.as_ptr() as *mut u8, layout);
236
2
            }
237
        }
238
2
    }
239
}
240
// SAFETY: The shared state is accessed concurrently by every thread holding a
241
// read guard, so the element type must be both `Send` and `Sync` for these
242
// accesses to be free of data races.
243
unsafe impl<T: Send + Sync> Send for BfVecShared<T> {}
244

            
245
// SAFETY: A `BfVec` is shared across threads by handing each thread a `share()`
246
// clone, which requires `T: Send + Sync` (see `BfVecShared` above).
247
unsafe impl<T: Send + Sync> Send for BfVec<T> {}
248

            
249
#[cfg(test)]
250
mod tests {
251
    use std::thread;
252

            
253
    use crate::BfVec;
254

            
255
    // These are just simple tests.
256
    #[test]
257
    #[cfg_attr(miri, ignore)]
258
1
    fn test_push() {
259
1
        let mut threads = vec![];
260

            
261
1
        let shared_vector = BfVec::<u32>::new();
262
1
        let num_threads = 10;
263
1
        let num_iterations = 10000;
264

            
265
10
        for t in 0..num_threads {
266
10
            let shared_vector = shared_vector.share();
267
10
            threads.push(thread::spawn(move || {
268
100000
                for _ in 0..num_iterations {
269
100000
                    shared_vector.push(t);
270
100000
                }
271
10
            }));
272
        }
273

            
274
        // Check whether threads have completed successfully.
275
10
        for thread in threads {
276
10
            thread.join().unwrap();
277
10
        }
278

            
279
        // Check the vector for some kind of consistency, correct total
280
1
        let mut total = 0;
281
100000
        for i in 0..shared_vector.len() {
282
100000
            total += shared_vector.at(i);
283
100000
        }
284

            
285
1
        assert_eq!(total, num_threads * (num_threads - 1) * num_iterations / 2);
286
1
        assert_eq!(shared_vector.len(), (num_threads * num_iterations) as usize);
287
1
    }
288

            
289
    #[test]
290
1
    fn test_clear_resets_length() {
291
1
        let vector = BfVec::<String>::new();
292
1
        vector.push("a".to_string());
293
1
        vector.push("b".to_string());
294
1
        assert_eq!(vector.len(), 2);
295

            
296
        // Dropping the vector afterwards must not drop the elements again.
297
1
        vector.clear();
298
1
        assert!(vector.is_empty());
299

            
300
1
        vector.push("c".to_string());
301
1
        assert_eq!(vector.at(0), "c");
302
1
    }
303
}