tiny_vec/
lib.rs

1/*  Copyright (C) 2025 Saúl Valdelvira
2 *
3 *  This program is free software: you can redistribute it and/or modify
4 *  it under the terms of the GNU General Public License as published by
5 *  the Free Software Foundation, version 3.
6 *
7 *  This program is distributed in the hope that it will be useful,
8 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
9 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
10 *  GNU General Public License for more details.
11 *
12 *  You should have received a copy of the GNU General Public License
13 *  along with this program.  If not, see <https://www.gnu.org/licenses/>. */
14
15//! Tiny Vec
16//!
17//! A dynamic array that can store a small amount of elements on the stack.
18//!
19//! This struct provides a vec-like API, but performs small-vector optimization.
20//! This means that a `TinyVec<T, N>` stores up to N elements on the stack.
21//! If the vector grows bigger than that, it moves the contents to the heap.
22//!
23//! # Example
24//! ```
25//! use tiny_vec::TinyVec;
26//!
27//! let mut tv = TinyVec::<u8, 16>::new();
28//!
29//! for n in 0..16 {
30//!     tv.push(n);
31//! }
32//!
33//! // Up to this point, no heap allocations are needed.
34//! // All the elements are stored on the stack.
35//!
36//! tv.push(123); // This moves the vector to the heap
37//!
38//! assert_eq!(&tv[..], &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9,
39//!                       10, 11, 12, 13, 14, 15, 123])
40//! ```
41//!
42//! # Memory layout
43//! For a TinyVec<T, N>
44//!
45//! On the stack (length <= N)
46//! - [T; N] : Data
47//! - usize  : Length
48//!
49//! On the heap (length > N)
50//! - T*    : Data
51//! - usize : Capacity
52//! - usize : Length
53//!
54//! If N is equal to `sizeof (T*, usize) / sizeof T`, the
55//! TinyVec is the same size as a regular vector \
56//! NOTE: The [n_elements_for_stack] function returns the maximun
57//! number of elements for a type, such that it doesn't waste extra
58//! space when moved to the heap
59//!
60
61#![allow(incomplete_features)]
62#![cfg_attr(feature = "nightly-const-generics", feature(generic_const_exprs))]
63
64use core::mem::{self, ManuallyDrop, MaybeUninit};
65use core::ops::{Deref, DerefMut};
66use core::ptr::NonNull;
67use core::{fmt, ptr};
68use core::slice;
69
70mod raw;
71use raw::{next_cap, RawVec};
72
73union TinyVecInner<T, const N: usize> {
74    stack: ManuallyDrop<[MaybeUninit<T>; N]>,
75    raw: RawVec<T>,
76}
77
78impl<T, const N: usize> TinyVecInner<T, N> {
79    unsafe fn as_ptr_stack(&self) -> *const T {
80        unsafe { &self.stack as *const _ as *const T }
81    }
82
83    unsafe fn as_ptr_stack_mut(&mut self) -> *mut T {
84        unsafe { &mut self.stack as *mut _ as *mut T }
85    }
86
87    unsafe fn as_ptr_heap(&self) -> *const T {
88        unsafe { self.raw.ptr.as_ptr() }
89    }
90
91    unsafe fn as_ptr_heap_mut(&mut self) -> *mut T {
92        unsafe { self.raw.ptr.as_ptr() }
93    }
94}
95
96#[repr(transparent)]
97struct Length(usize);
98
99impl Length {
100    const fn new_heap(len: usize) -> Self {
101        Self(len << 1 | 0b1)
102    }
103
104    #[inline]
105    const fn is_stack(&self) -> bool {
106        (self.0 & 0b1) == 0
107    }
108
109    #[inline]
110    const fn set_heap(&mut self) {
111        self.0 |= 0b1;
112    }
113
114    #[inline]
115    const fn get(&self) -> usize {
116        self.0 >> 1
117    }
118
119    #[inline]
120    const fn add(&mut self, n: usize) {
121        self.0 += n << 1;
122    }
123
124    #[inline]
125    const fn sub(&mut self, n: usize) {
126        self.0 -= n << 1;
127    }
128}
129
130/// The maximun number of elements that can be stored in the stack
131/// for the vector, without incrementing it's size
132///
133/// This means, that [`n_elements_for_stack`] for T returns the max
134/// number of elements, so that when switching to a heap allocated
135/// buffer, no stack size is wasted
136///
137/// # Examples
138/// ```
139/// use tiny_vec::n_elements_for_stack;
140///
141/// assert_eq!(n_elements_for_stack::<u8>(), 16);
142/// assert_eq!(n_elements_for_stack::<u16>(), 8);
143/// assert_eq!(n_elements_for_stack::<i32>(), 4);
144/// ```
145pub const fn n_elements_for_stack<T>() -> usize {
146    mem::size_of::<RawVec<T>>() / mem::size_of::<T>()
147}
148
149/// A dynamic array that can store a small amount of elements on the stack.
150pub struct TinyVec<T,
151    #[cfg(not(feature = "nightly-const-generics"))]
152    const N: usize,
153
154    #[cfg(feature = "nightly-const-generics")]
155    const N: usize = { n_elements_for_stack::<T>() },
156> {
157    inner: TinyVecInner<T, N>,
158    len: Length,
159}
160
161impl<T, const N: usize> TinyVec<T, N> {
162
163    fn ptr(&self) -> *const T {
164        unsafe {
165            if self.len.is_stack() {
166                self.inner.as_ptr_stack()
167            } else {
168                self.inner.as_ptr_heap()
169            }
170        }
171    }
172
173    fn ptr_mut(&mut self) -> *mut T {
174        unsafe {
175            if self.len.is_stack() {
176                self.inner.as_ptr_stack_mut()
177            } else {
178                self.inner.as_ptr_heap_mut()
179            }
180        }
181    }
182
183    unsafe fn switch(&mut self, n: usize) {
184        let cap = next_cap(self.len.get() + n);
185        let vec = RawVec::with_capacity(cap);
186        unsafe {
187            let src = self.inner.as_ptr_stack();
188            let dst = vec.ptr.as_ptr();
189            ptr::copy(
190                src,
191                dst,
192                self.len.get()
193            );
194            self.inner.raw = vec;
195        }
196        self.len.set_heap();
197    }
198}
199
200impl<T, const N: usize> TinyVec<T, N> {
201
202    pub const fn new() -> Self {
203        let stack = [ const { MaybeUninit::uninit() }; N ];
204        Self {
205            inner: TinyVecInner { stack: ManuallyDrop::new(stack) },
206            len: Length(0),
207        }
208    }
209
210    pub fn with_capacity(cap: usize) -> Self {
211        let mut len = Length(0);
212        let inner = if cap <= N {
213            let s = [ const { MaybeUninit::uninit() }; N ];
214            TinyVecInner { stack: ManuallyDrop::new(s) }
215        } else {
216            len.set_heap();
217            TinyVecInner { raw: RawVec::with_capacity(cap) }
218        };
219
220        Self { inner, len }
221    }
222
223    /// Returns the number of elements inside this vec
224    #[inline]
225    pub const fn len(&self) -> usize { self.len.get() }
226
227    /// Returns true if the vector is empty
228    #[inline]
229    pub const fn is_empty(&self) -> bool { self.len.get() == 0 }
230
231    /// Returns the allocated capacity for this vector
232    pub const fn capacity(&self) -> usize {
233        if self.len.is_stack() {
234            N
235        } else {
236            unsafe { self.inner.raw.cap }
237        }
238    }
239
240    /// Reserves space for, at least, n elements
241    pub fn reserve(&mut self, n: usize) {
242        if self.len.is_stack() {
243            if self.len.get() + n > N {
244                unsafe {
245                    self.switch(n);
246                }
247            }
248        } else {
249            unsafe {
250                self.inner.raw.expand_if_needed(self.len.get(), n);
251            }
252        }
253    }
254
255    /// Appends an element to the back of the vector
256    pub fn push(&mut self, elem: T) {
257        self.reserve(1);
258        unsafe {
259            let dst = self.ptr_mut().add(self.len.get());
260            dst.write(elem);
261        }
262        self.len.add(1);
263    }
264
265    /// Try to push an element inside the vector, only if
266    /// there's room for it. If the push would've caused a
267    /// reallocation of the buffer, returns the value.
268    pub fn push_within_capacity(&mut self, val: T) -> Result<(),T> {
269        if self.len.get() < self.capacity() {
270            unsafe {
271                // TODO CHECK
272                self.ptr_mut().add(self.len.get()).write(val)
273            }
274            self.len.add(1);
275            Ok(())
276        } else {
277            Err(val)
278        }
279    }
280    /// Removes the last element of this vector (if present)
281    pub fn pop(&mut self) -> Option<T> {
282        if self.len.get() == 0 {
283            None
284        } else {
285            self.len.sub(1);
286            let val =
287            unsafe {
288                self.ptr().add(self.len.get()).read()
289            };
290            Some(val)
291        }
292    }
293
294    /// Inserts an element in the given index position
295    #[inline]
296    pub fn insert(&mut self, index: usize, elem: T) -> Result<(),T> {
297        if index > self.len.get() { return Err(elem) }
298        unsafe { self.insert_unckecked(index, elem); }
299        Ok(())
300    }
301
302    /// # Safety
303    ///
304    /// The index should be valid.
305    pub unsafe fn insert_unckecked(&mut self, index: usize, elem: T) {
306        self.reserve(1);
307
308        unsafe {
309            ptr::copy(
310                self.ptr().add(index),
311                self.ptr_mut().add(index + 1),
312                self.len.get() - index,
313            );
314            self.ptr_mut().add(index).write(elem);
315        }
316        self.len.add(1);
317    }
318
319    /// Reserves space for exactly n elements more
320    pub fn reserve_exact(&mut self, n: usize) {
321        if self.len.is_stack() {
322            if self.len.get() + n > N {
323                unsafe { self.switch(n); }
324            }
325        } else {
326            let vec = unsafe { &mut self.inner.raw };
327            let new_cap = vec.cap.max(self.len.get() + n);
328            vec.expand(new_cap);
329        }
330    }
331
332    /// # Safety
333    /// Index should be < Vec::len()
334    pub unsafe fn remove_unchecked(&mut self, index: usize) -> T {
335        debug_assert!(index < self.len.get(), "Index is >= than {}, this will trigger UB", self.len.get());
336
337        unsafe {
338            self.len.sub(1);
339            let result = self.ptr_mut().add(index).read();
340            ptr::copy(
341                self.ptr().add(index + 1),
342                self.ptr_mut().add(index),
343                self.len.get() - index,
344            );
345            result
346        }
347    }
348
349    #[inline]
350    pub fn remove(&mut self, index: usize) -> Option<T> {
351        if index >= self.len.get() { return None }
352        /* Safety: We've just checked the invariant. Index will always
353         * be < self.len, so it's 100% safe to call this function */
354        Some( unsafe { self.remove_unchecked(index) } )
355    }
356
357    #[inline]
358    pub fn shrink_to_fit(&mut self) {
359        if !self.len.is_stack() {
360            unsafe { self.inner.raw.shrink_to_fit(self.len.get()); }
361        }
362    }
363
364    pub fn push_slice(&mut self, s: &[T]) {
365        let len = s.len();
366        let src = s.as_ptr();
367
368        self.reserve(len);
369
370        unsafe {
371            ptr::copy(
372                src,
373                self.ptr_mut().add(self.len.get()),
374                len
375            )
376        }
377
378        self.len.add(len);
379    }
380
381    pub fn extend_from<I>(&mut self, it: I)
382    where
383        I: IntoIterator<Item = T>,
384    {
385        let it = it.into_iter();
386
387        let (min, max) = it.size_hint();
388        let reserve = match max {
389            Some(max) => max,
390            None => min,
391        };
392
393        if reserve > 0 {
394            self.reserve(reserve);
395        }
396
397        for elem in it {
398            self.push(elem);
399        }
400    }
401}
402
403impl<T, const N: usize> Default for TinyVec<T, N> {
404    fn default() -> Self {
405        Self::new()
406    }
407}
408
409impl<T: fmt::Debug, const N: usize> fmt::Debug for TinyVec<T, N> {
410    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
411        fmt::Debug::fmt(&self[..], f)
412    }
413}
414
415impl<T: PartialEq, const N: usize> PartialEq for TinyVec<T, N> {
416    fn eq(&self, other: &Self) -> bool {
417        self.deref() == other.deref()
418    }
419}
420
421impl<T, const N: usize> Deref for TinyVec<T, N> {
422    type Target = [T];
423
424    fn deref(&self) -> &Self::Target {
425        unsafe { slice::from_raw_parts(self.ptr(), self.len.get()) }
426    }
427}
428
429
430impl<T, const N: usize> DerefMut for TinyVec<T, N> {
431    fn deref_mut(&mut self) -> &mut Self::Target {
432        unsafe { slice::from_raw_parts_mut(self.ptr_mut(), self.len.get()) }
433    }
434}
435
436impl<T, const N: usize> Drop for TinyVec<T, N> {
437    fn drop(&mut self) {
438        if mem::needs_drop::<T>() {
439            for e in self.deref_mut() {
440                unsafe { ptr::drop_in_place(e) };
441            }
442        }
443        if !self.len.is_stack() {
444            unsafe { self.inner.raw.destroy(); }
445        }
446    }
447}
448
449impl<T, const N: usize> From<Vec<T>> for TinyVec<T, N> {
450    fn from(value: Vec<T>) -> Self {
451        let mut md = ManuallyDrop::new(value);
452        let ptr = unsafe { NonNull::new_unchecked( md.as_mut_ptr() ) };
453        let inner = TinyVecInner {
454            raw: RawVec {
455                cap: md.capacity(),
456                ptr,
457            }
458        };
459
460        Self {
461            inner,
462            len: Length::new_heap(md.len()),
463        }
464    }
465}
466
467#[cfg(test)]
468mod test;
tiny_vec/lib.rs

tiny_vec/
lib.rs
