Skip to main content

turbo_trace_server/
lazy_sorted_vec.rs

1use std::{cell::UnsafeCell, ops::Deref, sync::Once};
2
3use smallvec::SmallVec;
4
5/// Backing storage for `LazySortedVec`. Inlines a single element so that
6/// the common case (e.g. a leaf span with one self-time event) does not
7/// pay a heap allocation. With one inline slot plus the workspace's
8/// `union` smallvec feature, this is the same size as a `Vec`.
9type Inner<T> = SmallVec<[T; 1]>;
10
11pub struct LazySortedVec<T> {
12    vec: UnsafeCell<Inner<T>>,
13    once: Once,
14}
15
16unsafe impl<T> Send for LazySortedVec<T> where T: Send {}
17unsafe impl<T> Sync for LazySortedVec<T> where T: Sync {}
18
19impl<T> LazySortedVec<T> {
20    pub fn new() -> Self {
21        Self {
22            vec: UnsafeCell::new(SmallVec::new()),
23            once: Once::new(),
24        }
25    }
26
27    pub fn push(&mut self, value: T) {
28        self.once = Once::new();
29        self.vec.get_mut().push(value);
30    }
31
32    pub fn retain_unordered(&mut self, mut f: impl FnMut(&T) -> bool) {
33        self.vec.get_mut().retain(|t| f(t));
34    }
35
36    pub fn iter_mut_unordered(&mut self) -> std::slice::IterMut<'_, T> {
37        self.vec.get_mut().iter_mut()
38    }
39}
40
41impl<T: Ord> Deref for LazySortedVec<T> {
42    type Target = [T];
43
44    fn deref(&self) -> &Self::Target {
45        let ptr = self.vec.get();
46        self.once.call_once(|| {
47            // SAFETY: The only access to the `vec` is through this `Deref` implementation, or we
48            // have a `&mut self` which prevents a simultaneous `Deref`. So we can guarantee that
49            // there are no other accesses to the `vec` while we sort it.
50            let inner = unsafe { &mut *ptr };
51            inner.sort();
52            inner.shrink_to_fit();
53        });
54        // SAFETY: Returning this reference is safe because the lifetime guarantees that there is no
55        // `&mut self` that could cause a simultaneous access to the `vec`, and the `Once`
56        // guarantees that the sorting is complete before we return the reference.
57        unsafe { &*ptr }.as_slice()
58    }
59}
60
61impl<T> Default for LazySortedVec<T> {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl<T> From<Inner<T>> for LazySortedVec<T> {
68    fn from(vec: Inner<T>) -> Self {
69        Self {
70            vec: UnsafeCell::new(vec),
71            once: Once::new(),
72        }
73    }
74}
75
76impl<T> From<Vec<T>> for LazySortedVec<T> {
77    fn from(vec: Vec<T>) -> Self {
78        Self {
79            vec: UnsafeCell::new(SmallVec::from_vec(vec)),
80            once: Once::new(),
81        }
82    }
83}