Skip to main content

luaur_common/records/
small_vector.rs

1//! Faithful hand-port of `Luau::SmallVector<T, N>`.
2//!
3//! Reference: `luau/Common/include/Luau/SmallVector.h` (whole file), Luau upstream.
4//! A small-buffer-optimized vector: the first `N` elements live inline; on growth
5//! it spills to a heap block of `1.5x`-or-`+4` capacity.
6//!
7//! **Sound-representation deviation (required, not cosmetic):** the C++ caches a
8//! `ptr` that points either into its own inline `storage` or at the heap block.
9//! A Rust value can be *moved* freely, which would dangle a pointer into inline
10//! storage. So we do NOT cache a self-referential pointer: `heap` is null while
11//! inline and the data pointer is recomputed from `storage`/`heap` on demand.
12//! Behavior (SBO, capacities, element lifetimes) is otherwise identical.
13
14extern crate alloc;
15
16use alloc::alloc::{alloc, dealloc, handle_alloc_error, Layout};
17use core::fmt;
18use core::hash::{Hash, Hasher};
19use core::mem::MaybeUninit;
20use core::ops::{Deref, DerefMut};
21use core::ptr;
22use core::slice;
23
24#[allow(non_camel_case_types)]
25pub struct SmallVector<T, const N: usize> {
26    /// Inline storage for the first `N` elements (small-buffer optimization).
27    storage: [MaybeUninit<T>; N],
28    /// Heap block once we outgrow `N`; null while inline. Never a pointer into
29    /// `storage` — see the module note on move-safety.
30    heap: *mut T,
31    /// Number of live elements.
32    count: u32,
33    /// Capacity: `N` while inline, the heap block's capacity otherwise.
34    max: u32,
35}
36
37impl<T, const N: usize> SmallVector<T, N> {
38    pub fn new() -> Self {
39        SmallVector {
40            storage: [const { MaybeUninit::uninit() }; N],
41            heap: ptr::null_mut(),
42            count: 0,
43            max: N as u32,
44        }
45    }
46
47    fn is_heap(&self) -> bool {
48        !self.heap.is_null()
49    }
50
51    /// Current data pointer, recomputed (never cached) so moves stay sound.
52    fn data(&self) -> *const T {
53        if self.heap.is_null() {
54            self.storage.as_ptr() as *const T
55        } else {
56            self.heap
57        }
58    }
59
60    fn data_mut(&mut self) -> *mut T {
61        if self.heap.is_null() {
62            self.storage.as_mut_ptr() as *mut T
63        } else {
64            self.heap
65        }
66    }
67
68    pub fn as_slice(&self) -> &[T] {
69        // SAFETY: `data()` is valid for `count` initialized, contiguous elements.
70        unsafe { slice::from_raw_parts(self.data(), self.count as usize) }
71    }
72
73    pub fn as_mut_slice(&mut self) -> &mut [T] {
74        let len = self.count as usize;
75        let ptr = self.data_mut();
76        // SAFETY: as `as_slice`, with unique access from `&mut self`.
77        unsafe { slice::from_raw_parts_mut(ptr, len) }
78    }
79
80    pub fn size(&self) -> u32 {
81        self.count
82    }
83
84    pub fn capacity(&self) -> u32 {
85        self.max
86    }
87
88    pub fn empty(&self) -> bool {
89        self.count == 0
90    }
91
92    pub fn front(&self) -> &T {
93        assert!(self.count > 0);
94        &self.as_slice()[0]
95    }
96
97    pub fn back(&self) -> &T {
98        assert!(self.count > 0);
99        &self.as_slice()[self.count as usize - 1]
100    }
101
102    /// std-style alias for `push_back` — translations use Vec idioms.
103    pub fn push(&mut self, value: T) {
104        self.push_back(value);
105    }
106
107    pub fn push_back(&mut self, value: T) {
108        if self.count == self.max {
109            self.grow(self.count + 1);
110        }
111        let index = self.count as usize;
112        // SAFETY: we just ensured capacity for one more element at `index`.
113        unsafe { self.data_mut().add(index).write(value) };
114        self.count += 1;
115    }
116
117    /// `emplace_back` collapses to `push_back` of the constructed value; Rust has
118    /// no in-place variadic construction, and the move is free.
119    pub fn emplace_back(&mut self, value: T) -> &mut T {
120        self.push_back(value);
121        let index = self.count as usize - 1;
122        &mut self.as_mut_slice()[index]
123    }
124
125    pub fn pop_back(&mut self) {
126        assert!(self.count > 0);
127        self.count -= 1;
128        let index = self.count as usize;
129        // SAFETY: element at `index` was live and is now logically removed.
130        unsafe { ptr::drop_in_place(self.data_mut().add(index)) };
131    }
132
133    pub fn clear(&mut self) {
134        while self.count > 0 {
135            self.pop_back();
136        }
137    }
138
139    pub fn reserve(&mut self, reserve_size: u32) {
140        if reserve_size > self.max {
141            self.grow(reserve_size);
142        }
143    }
144
145    /// Faithful to the C++ growth policy: `1.5x`, or `+4` when that is too small.
146    fn grow(&mut self, new_size: u32) {
147        let new_size = if self.max + (self.max >> 1) > new_size {
148            self.max + (self.max >> 1)
149        } else {
150            new_size + 4
151        };
152        assert!(new_size < 0x4000_0000);
153
154        let layout = Layout::array::<T>(new_size as usize).expect("SmallVector capacity overflow");
155        // SAFETY: `Layout::array` rejects zero-size layouts for non-ZST `T`; null
156        // is handled immediately below.
157        let new_data = unsafe { alloc(layout) as *mut T };
158        if new_data.is_null() {
159            handle_alloc_error(layout);
160        }
161
162        // SAFETY: move (bitwise) the `count` live elements into the new block. The
163        // sources are then logically uninitialized and must NOT be dropped, which
164        // is exactly what switching `heap`/`is_heap` to the new block achieves.
165        unsafe {
166            ptr::copy_nonoverlapping(self.data(), new_data, self.count as usize);
167            if self.is_heap() {
168                let old_layout =
169                    Layout::array::<T>(self.max as usize).expect("SmallVector capacity overflow");
170                dealloc(self.heap as *mut u8, old_layout);
171            }
172        }
173
174        self.heap = new_data;
175        self.max = new_size;
176    }
177}
178
179impl<T: Default, const N: usize> SmallVector<T, N> {
180    pub fn resize(&mut self, new_size: u32) {
181        if new_size > self.count {
182            if new_size > self.max {
183                self.grow(new_size);
184            }
185            for index in self.count as usize..new_size as usize {
186                // SAFETY: capacity ensured above; writing a fresh element.
187                unsafe { self.data_mut().add(index).write(T::default()) };
188            }
189        } else {
190            while self.count > new_size {
191                self.pop_back();
192            }
193        }
194        self.count = new_size;
195    }
196}
197
198impl<T, const N: usize> Default for SmallVector<T, N> {
199    fn default() -> Self {
200        Self::new()
201    }
202}
203
204impl<T, const N: usize> Deref for SmallVector<T, N> {
205    type Target = [T];
206    fn deref(&self) -> &[T] {
207        self.as_slice()
208    }
209}
210
211impl<T, const N: usize> DerefMut for SmallVector<T, N> {
212    fn deref_mut(&mut self) -> &mut [T] {
213        self.as_mut_slice()
214    }
215}
216
217impl<T: Clone, const N: usize> Clone for SmallVector<T, N> {
218    fn clone(&self) -> Self {
219        let mut copy = Self::new();
220        copy.reserve(self.count);
221        for value in self.as_slice() {
222            copy.push_back(value.clone());
223        }
224        copy
225    }
226}
227
228impl<T, const N: usize> Drop for SmallVector<T, N> {
229    fn drop(&mut self) {
230        self.clear();
231        if self.is_heap() {
232            let layout =
233                Layout::array::<T>(self.max as usize).expect("SmallVector capacity overflow");
234            // SAFETY: `heap` was allocated with this layout in `grow`.
235            unsafe { dealloc(self.heap as *mut u8, layout) };
236        }
237    }
238}
239
240impl<T: PartialEq, const N: usize> PartialEq for SmallVector<T, N> {
241    fn eq(&self, other: &Self) -> bool {
242        self.as_slice() == other.as_slice()
243    }
244}
245
246impl<T: Eq, const N: usize> Eq for SmallVector<T, N> {}
247
248impl<T: Hash, const N: usize> Hash for SmallVector<T, N> {
249    fn hash<H: Hasher>(&self, state: &mut H) {
250        self.as_slice().hash(state);
251    }
252}
253
254impl<T: fmt::Debug, const N: usize> fmt::Debug for SmallVector<T, N> {
255    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
256        f.debug_list().entries(self.as_slice()).finish()
257    }
258}
259
260impl<'a, T, const N: usize> IntoIterator for &'a SmallVector<T, N> {
261    type Item = &'a T;
262    type IntoIter = slice::Iter<'a, T>;
263    fn into_iter(self) -> Self::IntoIter {
264        self.as_slice().iter()
265    }
266}
267
268impl<T, const N: usize> FromIterator<T> for SmallVector<T, N> {
269    fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
270        let mut out = Self::new();
271        for value in iter {
272            out.push_back(value);
273        }
274        out
275    }
276}
277
278// SAFETY: a `SmallVector<T>` owns its `T`s (inline or on the heap); the raw
279// `heap` pointer carries no shared state, so it is `Send`/`Sync` exactly when `T`
280// is, matching `alloc::vec::Vec`.
281unsafe impl<T: Send, const N: usize> Send for SmallVector<T, N> {}
282unsafe impl<T: Sync, const N: usize> Sync for SmallVector<T, N> {}
283
284impl<T, const N: usize> crate::records::dense_hash_table::DenseDefault for SmallVector<T, N> {
285    fn dense_default() -> Self {
286        Self::new()
287    }
288}