commonlibsse_ng\re\b\BSTArray/
small_heap.rs

1// Copyright (c) 2018 The Servo Project Developers
2// SPDX-License-Identifier: Apache-2.0 OR MIT
3// https://github.com/servo/rust-smallvec/blob/v2/src/lib.rs#L3
4
5use core::{
6    alloc::Layout,
7    marker::PhantomData,
8    mem::{ManuallyDrop, MaybeUninit},
9    ops::{Index, IndexMut, Range, RangeBounds},
10    ptr::{self, NonNull, drop_in_place},
11};
12use std::alloc::handle_alloc_error;
13
14use std_fork::alloc::SelflessAllocator;
15use stdx::ptr::{ConstNonNull, Unique};
16
17use crate::re::MemoryManager::TESGlobalAlloc;
18
19/// Use stack while within the specified size, and use heap if it is larger.
20///
21/// This is the same purpose as `smallvec` crate and other optimizations, except that the memory layout is for TES.
22///
23/// It is effective when most of the memory can be fit in the stack, except for some exceptions, but it slows down the process if there are frequent fallbacks to the heap.
24///
25/// - [`smallvec` crate](https://crates.io/crates/smallvec)
26///
27/// - `N`: element length (not bytes size)
28///
29/// # Examples
30/// ```
31/// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
32/// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
33///
34/// let mut array = BSTSmallArray::<i32, 1>::new();
35/// array.push(1); // push on stack.
36/// array.push(2); // change to heap.
37/// assert_eq!(array[0], 1);
38/// assert_eq!(array[1], 2);
39/// assert_eq!(array.len(), 2);
40/// array.clear(); // reuse heap. grow to capacity 4.
41/// assert_eq!(array.len(), 0);
42/// assert_eq!(array.capacity(), 4); // Capacity is preserved
43/// ```
44#[repr(C)]
45pub struct BSTSmallArray<T, const N: usize = 1, A = TESGlobalAlloc>
46where
47    A: SelflessAllocator,
48{
49    // BSTSmallArrayHeapAllocator
50    /// Capacity is at least 4 and increases in multiples of 2.(In `self.grow`)
51    capacity: u32, // 0x00,
52
53    /// Indicates whether the data is stored locally (on the stack).
54    storage_type: StorageType_CEnum, // 0x04,
55
56    // The union of local stack data and heap pointer.
57    data: RawBSTSmallArray<T, N>, // 0x08
58
59    // BSTArrayBase
60    // length of elements
61    size: u32, // 0x10
62
63    // Assumed Zero size type.
64    alloc: PhantomData<A>,
65}
66const _: () = assert!(core::mem::size_of::<BSTSmallArray<(), 10>>() == 0x18);
67
68/// Indicates how the data in `BSTSmallArray` is stored.
69///
70/// - `Heap`: Data is allocated on the heap.
71/// - `Inline`: Data is stored inline (e.g., on the stack or within the object).
72#[commonlibsse_ng_derive_internal::ffi_enum]
73#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74#[repr(u32)]
75pub enum StorageType {
76    /// Heap-allocated storage.
77    Heap = 0,
78    /// Inline storage (e.g., stack or inline buffer).
79    Inline = 1,
80}
81
82/// A union that stores either a heap pointer or a fixed-size array for local storage.
83#[repr(C)]
84union RawBSTSmallArray<T, const N: usize> {
85    /// Pointer to heap memory. Same as `*mut T` memory layout
86    heap: Option<Unique<T>>,
87    /// Fixed-size array for local (stack) storage.
88    inline: ManuallyDrop<MaybeUninit<[T; N]>>,
89}
90
91impl<T, const N: usize> RawBSTSmallArray<T, N> {
92    #[inline]
93    const fn new() -> Self {
94        Self::new_inline(MaybeUninit::uninit())
95    }
96    #[inline]
97    const fn new_inline(inline: MaybeUninit<[T; N]>) -> Self {
98        Self { inline: ManuallyDrop::new(inline) }
99    }
100    #[inline]
101    const fn new_heap(ptr: NonNull<T>) -> Self {
102        Self { heap: Some(unsafe { Unique::new_unchecked(ptr.as_ptr()) }) }
103    }
104}
105
106const _: () = {
107    const SIZE: usize = core::mem::size_of::<RawBSTSmallArray<u32, 4>>();
108    assert!(SIZE == 0x10);
109};
110
111impl<T, const N: usize, A> BSTSmallArray<T, N, A>
112where
113    A: SelflessAllocator,
114{
115    /// Creates a new, empty `BSTSmallArray<T, N, A>` with the specified allocator.
116    ///
117    /// The array will not allocate until elements are pushed.
118    ///
119    /// # Example
120    /// ```
121    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
122    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
123    ///
124    /// let array = BSTSmallArray::<i32, 10>::new();
125    /// assert!(array.is_empty());
126    /// ```
127    #[inline]
128    pub const fn new() -> Self {
129        Self {
130            data: RawBSTSmallArray::new(),
131            capacity: N as u32, // ?INFO: The NG branch of the VR version is at 0.
132            storage_type: StorageType_CEnum::from_enum(StorageType::Inline),
133            size: 0,
134            alloc: PhantomData,
135        }
136    }
137
138    /// Returns the number of elements in the array.
139    ///
140    /// This is also referred to as the array’s "length".
141    ///
142    /// # Example
143    /// ```
144    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
145    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
146    ///
147    /// let array = BSTSmallArray::<i32, 10>::new();
148    /// assert_eq!(array.len(), 0);
149    /// ```
150    #[inline]
151    pub const fn len(&self) -> usize {
152        self.size as usize
153    }
154
155    /// Returns `true` if the array contains no elements.
156    ///
157    /// # Example
158    /// ```
159    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
160    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
161    ///
162    /// let array = BSTSmallArray::<i32, 10>::new();
163    /// assert!(array.is_empty());
164    /// ```
165    #[inline]
166    pub const fn is_empty(&self) -> bool {
167        self.len() == 0
168    }
169
170    /// Returns the total number of elements the array can hold without reallocating.
171    ///
172    /// This is the allocated capacity, which may be larger than the current length.
173    ///
174    /// # Example
175    /// ```
176    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
177    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
178    ///
179    /// let mut array = BSTSmallArray::<i32, 10>::new();
180    /// assert!(array.capacity() == 10); // stack size
181    /// ```
182    #[inline]
183    pub const fn capacity(&self) -> usize {
184        self.capacity as usize
185    }
186
187    /// Shrinks the capacity of the array as much as possible.
188    ///
189    /// It will drop any excess capacity not used by the current elements.
190    ///
191    /// # Examples
192    ///
193    /// ```
194    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
195    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
196    ///
197    /// let mut array = BSTSmallArray::<_, 4>::new(); // auto i32
198    /// array.push(1);
199    /// array.push(2);
200    /// array.push(3);
201    /// array.push(4);
202    /// array.push(5); // change to heap
203    /// assert_eq!(array.len(), 5);
204    /// array.shrink_to_fit();
205    /// assert!(array.capacity() == array.len());
206    /// ```
207    #[inline]
208    pub fn shrink_to_fit(&mut self) {
209        self.change_capacity(self.size);
210    }
211
212    /// Appends an element to the back of the array.
213    ///
214    /// # Panics
215    /// Panics if the array is at fixed capacity and cannot grow.
216    ///
217    /// # Example
218    /// ```
219    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
220    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
221    ///
222    /// let mut array = BSTSmallArray::<i32, 10>::new();
223    /// array.push(5);
224    /// assert_eq!(array[0], 5);
225    /// ```
226    #[inline]
227    pub fn push(&mut self, value: T) {
228        let size = self.size;
229        dbg!(size, self.capacity);
230        if size == self.capacity {
231            dbg!("Glow");
232            self.grow();
233        }
234        dbg!(size, self.capacity);
235        unsafe {
236            if let Some(ptr) = self.as_non_null_ptr() {
237                ptr.add(size as usize).write(value);
238            }
239        }
240
241        self.size += 1;
242    }
243
244    /// Removes the last element from the array and returns it, or `None` if it's empty.
245    ///
246    /// # Example
247    /// ```
248    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
249    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
250    ///
251    /// let mut array = BSTSmallArray::<i32, 10>::new();
252    /// array.push(1);
253    /// assert_eq!(array.pop(), Some(1));
254    /// assert_eq!(array.pop(), None);
255    /// ```
256    #[inline]
257    pub fn pop(&mut self) -> Option<T> {
258        let len = self.len();
259        if len == 0 {
260            None
261        } else {
262            self.size -= 1;
263            unsafe { Some(ptr::read(self.as_non_null_ptr()?.add(len - 1).as_ptr())) }
264        }
265    }
266
267    /// Returns a reference to the element at the given index, if it exists.
268    ///
269    /// # Example
270    /// ```
271    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
272    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
273    ///
274    /// let mut array = BSTSmallArray::<i32, 10>::new();
275    /// array.push(42);
276    /// assert_eq!(array.get(0), Some(&42));
277    /// assert_eq!(array.get(1), None);
278    /// ```
279    #[inline]
280    pub fn get(&self, index: usize) -> Option<&T> {
281        if index < self.len() {
282            return unsafe { Some(self.as_const_non_null_ptr()?.add(index).as_ref()) };
283        }
284        None
285    }
286
287    /// Returns a mutable reference to the element at the given index, if it exists.
288    ///
289    /// # Example
290    /// ```
291    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
292    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
293    ///
294    /// let mut array = BSTSmallArray::<i32, 1>::new();
295    /// array.push(10);
296    /// if let Some(x) = array.get_mut(0) {
297    ///     *x += 1;
298    /// }
299    /// assert_eq!(array[0], 11);
300    /// ```
301    #[inline]
302    pub fn get_mut(&mut self, index: usize) -> Option<&mut T> {
303        if index < self.len() {
304            return unsafe { Some(self.as_non_null_ptr()?.add(index).as_mut()) };
305        }
306        None
307    }
308
309    /// Clears the array, removing all elements but preserving the capacity.
310    ///
311    /// # Examples
312    /// ```
313    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
314    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
315    ///
316    /// let mut array = BSTSmallArray::<i32, 1>::new();
317    /// array.push(1); // push on stack.
318    /// array.push(2); // change to heap.
319    /// assert_eq!(array.len(), 2);
320    /// array.clear(); // reuse heap.
321    /// assert_eq!(array.len(), 0);
322    /// assert_eq!(array.capacity(), 4); // Capacity is preserved
323    /// ```
324    #[inline]
325    pub fn clear(&mut self) {
326        // Drop all elements in the array without changing capacity
327        if let Some(non_null) = self.as_non_null_ptr() {
328            for i in 0..self.len() {
329                unsafe {
330                    // SAFETY: we're dropping each element in place
331                    ptr::drop_in_place(non_null.add(i).as_ptr());
332                }
333            }
334        }
335
336        self.size = 0; // Reset the length, but keep the allocated capacity
337    }
338
339    /// Returns a non null pointer of the array’s buffer.
340    #[inline]
341    #[allow(clippy::missing_const_for_fn)] // Wrong lint. cannot use const when use deref
342    pub fn as_non_null_ptr(&mut self) -> Option<NonNull<T>> {
343        match self.storage_type() {
344            StorageType::Heap => unsafe {
345                self.data.heap.map(|ptr| NonNull::new_unchecked(ptr.as_ptr()))
346            },
347            StorageType::Inline => {
348                let ptr = unsafe { (*self.data.inline).as_mut_ptr() }.cast::<T>();
349                Some(unsafe { NonNull::new_unchecked(ptr) })
350            }
351        }
352    }
353
354    /// Returns a non null pointer of the array’s buffer.
355    #[inline]
356    pub fn as_const_non_null_ptr(&self) -> Option<ConstNonNull<T>> {
357        match self.storage_type() {
358            StorageType::Heap => unsafe { self.data.heap.map(ConstNonNull::from_unique) },
359            StorageType::Inline => {
360                let ptr = (unsafe { self.data.inline.as_ptr() }).cast::<T>();
361                Some(unsafe { ConstNonNull::new_unchecked(ptr) })
362            }
363        }
364    }
365
366    /// Checks if the array contains the given element.
367    ///
368    /// Returns `true` if the element is present in the array, and `false` otherwise.
369    ///
370    /// # Examples
371    ///
372    /// ```
373    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
374    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
375    ///
376    /// let mut array = BSTSmallArray::<i32, 10>::new();
377    /// array.push(1);
378    /// array.push(2);
379    /// assert!(array.contains(&1));
380    /// assert!(!array.contains(&3));
381    /// ```
382    #[inline]
383    pub fn contains(&self, value: &T) -> bool
384    where
385        T: PartialEq,
386    {
387        for i in 0..self.len() {
388            if let Some(item) = self.get(i) {
389                if item == value {
390                    return true;
391                }
392            }
393        }
394        false
395    }
396
397    /// Retains only the elements that satisfy the predicate.
398    ///
399    /// This method takes a closure that accepts an element of the array and returns a boolean.
400    /// Elements for which the closure returns `true` will be kept, while elements for which
401    /// it returns `false` will be removed.
402    ///
403    /// # Examples
404    ///
405    /// ```
406    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
407    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
408    ///
409    /// let mut array = BSTSmallArray::<i32, 10>::new();
410    /// array.push(1);
411    /// array.push(2);
412    /// array.push(3);
413    /// array.retain(|&x| x > 1);
414    /// assert_eq!(array.len(), 2);
415    /// assert!(array.contains(&2));
416    /// assert!(array.contains(&3));
417    /// ```
418    ///
419    /// # Panics
420    /// array ptr is null
421    #[inline]
422    pub fn retain<F>(&mut self, mut f: F)
423    where
424        F: FnMut(&T) -> bool,
425    {
426        let mut retained = 0;
427
428        for i in 0..self.len() {
429            let src = match self.as_non_null_ptr() {
430                Some(elem) => unsafe { elem.add(i) },
431                None => continue,
432            };
433
434            if f(unsafe { src.as_ref() }) {
435                if retained != i {
436                    unsafe {
437                        let dst = self.as_non_null_ptr().unwrap().add(retained).as_ptr();
438                        ptr::copy_nonoverlapping(src.as_ptr(), dst, 1);
439                    }
440                }
441                retained += 1;
442            } else {
443                unsafe { ptr::drop_in_place(src.as_ptr()) }; // Drop elements that do not match the predicate
444            }
445        }
446
447        self.size = retained as u32;
448    }
449
450    /// Resizes the array to the specified length.
451    ///
452    /// If the array is resized to a larger length, the new elements will be initialized
453    /// using the default constructor for `T`. If the array is resized to a smaller length,
454    /// elements at the end will be dropped.
455    ///
456    /// # Examples
457    ///
458    /// ```
459    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
460    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
461    ///
462    /// let mut array = BSTSmallArray::<i32, 10>::new();
463    /// array.push(1);
464    /// array.push(2);
465    /// array.resize(5, 0);
466    /// assert_eq!(array.len(), 5);
467    /// assert_eq!(array[3], 0);
468    /// ```
469    #[inline]
470    pub fn resize(&mut self, new_size: usize, value: T)
471    where
472        T: Clone,
473    {
474        let prev_size = self.len();
475        if new_size > prev_size {
476            for _ in prev_size..new_size {
477                self.push(value.clone());
478            }
479        } else {
480            for i in new_size..prev_size {
481                if let Some(src) = self.as_non_null_ptr() {
482                    unsafe { ptr::drop_in_place(src.add(i).as_ptr()) };
483                }
484            }
485        }
486        self.size = new_size as u32;
487    }
488
489    /// Removes a range of elements from the array, returning them as a vector.
490    ///
491    /// This method removes the elements within the specified range and returns them as
492    /// a `Vec<T>`. The range must be within the bounds of the array.
493    ///
494    /// # Examples
495    ///
496    /// ```
497    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
498    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
499    ///
500    /// let mut array = BSTSmallArray::<i32, 10>::new();
501    /// array.push(1);
502    /// array.push(2);
503    /// array.push(3);
504    /// array.push(4);
505    /// array.push(5);
506    /// let drained = array.drain(0..2);
507    /// assert_eq!(drained.collect::<Vec<_>>(), vec![1, 2]);
508    /// assert_eq!(array.len(), 3);
509    /// assert_eq!(array[0], 3);
510    /// assert_eq!(array[1], 4);
511    /// assert_eq!(array[2], 5);
512    /// ```
513    ///
514    /// # Panics
515    /// Panics if the range is out of bounds.
516    #[inline]
517    pub fn drain<R>(&mut self, range: R) -> BSTDrain<'_, T, N, A>
518    where
519        R: RangeBounds<usize>,
520    {
521        let len = self.len();
522        let Range { start, end } = stdx::slice::range(range, ..len);
523        debug_assert!(start <= end);
524        debug_assert!(end <= len);
525
526        // Need this.
527        // If the size is not changed before creating iter for Drain, inconsistencies will occur.
528        self.size = start as u32;
529
530        let iter = match self.as_non_null_ptr() {
531            Some(src) => unsafe {
532                core::slice::from_raw_parts(src.add(start).as_ptr(), end - start)
533            },
534            None => &[],
535        }
536        .iter();
537
538        BSTDrain {
539            iter,
540            tail_start: end,
541            tail_len: len - end,
542            array: unsafe { NonNull::new_unchecked(self as *mut Self) },
543        }
544    }
545
546    /// Returns a slice of all elements in the array.
547    #[inline]
548    pub fn as_slice(&self) -> &[T] {
549        let ptr = self.as_const_non_null_ptr();
550        let len = self.len();
551
552        if ptr.is_none() || (len == 0) {
553            return &[];
554        }
555
556        if let Some(src) = self.as_const_non_null_ptr() {
557            unsafe { core::slice::from_raw_parts(src.as_ptr(), len) }
558        } else {
559            &[]
560        }
561    }
562
563    /// Returns a mutable slice of all elements in the array.
564    #[inline]
565    pub fn as_mut_slice(&mut self) -> &mut [T] {
566        let ptr = self.as_non_null_ptr();
567        let len = self.len();
568
569        if ptr.is_none() || (len == 0) {
570            return &mut [];
571        }
572
573        if let Some(src) = self.as_non_null_ptr() {
574            unsafe { core::slice::from_raw_parts_mut(src.as_ptr(), len) }
575        } else {
576            &mut []
577        }
578    }
579
580    /// Returns an iterator over the elements of the array.
581    ///
582    /// This iterator yields references to the elements in the array.
583    ///
584    /// # Examples
585    ///
586    /// ```
587    /// # use commonlibsse_ng::re::BSTArray::BSTSmallArray as BSTSmallArray_;
588    /// # type BSTSmallArray<T, const N: usize> = BSTSmallArray_<T, N, stdx::alloc::Global>;
589    ///
590    /// let mut array = BSTSmallArray::<i32, 10>::new();
591    /// array.push(1);
592    /// array.push(2);
593    /// let sum: i32 = array.iter().sum();
594    /// assert_eq!(sum, 3);
595    /// ```
596    #[inline]
597    pub const fn iter(&self) -> BSTSmallArrayIterator<'_, T, N, A> {
598        BSTSmallArrayIterator { array: self, index: 0 }
599    }
600
601    /// Grows the array's capacity when the current size exceeds the limit.
602    ///
603    /// Doubles the current capacity, or sets it to a minimum threshold if it's too small.
604    ///
605    /// # Note
606    /// - If `new_capacity` <= `N` and in stack mode, then do nothing
607    fn grow(&mut self) {
608        const MIN_CAPACITY: u32 = 4;
609        const GROWTH_FACTOR: u32 = 2;
610
611        let old_capacity = self.capacity;
612        let new_capacity =
613            if old_capacity < MIN_CAPACITY { MIN_CAPACITY } else { old_capacity * GROWTH_FACTOR };
614        self.change_capacity(new_capacity);
615    }
616
617    /// Changes the capacity of the array to the specified value.
618    ///
619    /// # Note
620    /// - If `new_capacity` is 0, the function does nothing.
621    /// - If `new_capacity` <= `N` and in stack mode, then do nothing
622    fn change_capacity(&mut self, new_capacity: u32) {
623        if new_capacity == 0 {
624            return;
625        }
626
627        let storage_type = self.storage_type();
628
629        if new_capacity <= (N as u32) && storage_type == StorageType::Inline {
630            return;
631        }
632
633        let copy_count = core::cmp::min(self.size, new_capacity) as usize;
634
635        let new_layout = Self::new_layout(new_capacity);
636        let new_data = A::allocate(new_layout)
637            .map_or_else(|_| handle_alloc_error(new_layout), |data| data.cast::<T>());
638
639        unsafe {
640            let dst = new_data.as_ptr();
641
642            match storage_type {
643                StorageType::Inline => {
644                    // Inline -> heap
645                    let src = self.data.inline.as_ptr().cast::<T>();
646                    ptr::copy_nonoverlapping(src, dst, copy_count);
647                }
648                StorageType::Heap => {
649                    if let Some(old_ptr) = self.data.heap {
650                        ptr::copy_nonoverlapping(old_ptr.as_ptr(), dst, copy_count);
651                        A::deallocate(old_ptr.as_non_null_ptr().cast::<u8>(), self.layout());
652                    }
653                }
654            }
655        }
656
657        self.data = RawBSTSmallArray::new_heap(new_data);
658        self.storage_type = StorageType::Heap.into();
659        self.capacity = new_capacity;
660    }
661
662    /// Creates a layout describing the record for a `[T; n]`.
663    ///
664    /// # Panics
665    /// On arithmetic overflow or when the total size would exceed
666    /// `isize::MAX`, panic.
667    fn new_layout(n: u32) -> Layout {
668        Layout::array::<T>(n as usize).expect("BSTSmallArray need: alloc size < isize::MAX")
669    }
670
671    /// Gets a current layout self.
672    ///
673    /// # Panics
674    /// On arithmetic overflow or when the total size would exceed
675    /// `isize::MAX`, panic.
676    fn layout(&self) -> Layout {
677        Self::new_layout(self.capacity)
678    }
679
680    const fn storage_type(&self) -> StorageType {
681        match self.storage_type.to_enum() {
682            Some(value) => value,
683            None => StorageType::Inline,
684        }
685    }
686}
687
688impl<T, const N: usize, A> Drop for BSTSmallArray<T, N, A>
689where
690    A: SelflessAllocator,
691{
692    fn drop(&mut self) {
693        unsafe {
694            match self.storage_type() {
695                StorageType::Heap => {
696                    if let Some(ptr) = self.data.heap {
697                        // Drop each element
698                        let ptr = ptr.as_non_null_ptr();
699                        for i in 0..self.size as usize {
700                            drop_in_place(ptr.add(i).as_ptr());
701                        }
702                        A::deallocate(ptr.cast(), self.layout());
703                    }
704                }
705                StorageType::Inline => {
706                    debug_assert!(
707                        self.size <= N as u32,
708                        "[BSTSmallArray] size is larger than stack capacity. Wrong Implementation."
709                    );
710                    for i in self.as_mut_slice() {
711                        drop_in_place(i);
712                    }
713                    // No need to deallocate inline storage
714                }
715            }
716        }
717    }
718}
719
720impl<T, const N: usize, A> Index<usize> for BSTSmallArray<T, N, A>
721where
722    A: SelflessAllocator,
723{
724    type Output = T;
725
726    #[inline]
727    fn index(&self, index: usize) -> &Self::Output {
728        assert!(index < self.len(), "Index out of bounds");
729        unsafe { self.as_const_non_null_ptr().unwrap().add(index).as_ref() }
730    }
731}
732
733impl<T, const N: usize, A> IndexMut<usize> for BSTSmallArray<T, N, A>
734where
735    A: SelflessAllocator,
736{
737    #[inline]
738    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
739        assert!(index < self.len(), "Index out of bounds");
740        unsafe { self.as_non_null_ptr().unwrap().add(index).as_mut() }
741    }
742}
743
744////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
745// Iterator
746
747/// Iterator returned by `BSTSmallArray::drain()`
748pub struct BSTDrain<'a, T, const N: usize, A>
749where
750    A: SelflessAllocator,
751{
752    tail_start: usize, // = range.end
753    tail_len: usize,   // = original_len - range.end
754    iter: core::slice::Iter<'a, T>,
755    array: NonNull<BSTSmallArray<T, N, A>>,
756}
757
758impl<T, const N: usize, A> Iterator for BSTDrain<'_, T, N, A>
759where
760    A: SelflessAllocator,
761{
762    type Item = T;
763
764    #[inline]
765    fn next(&mut self) -> Option<Self::Item> {
766        self.iter.next().map(|item| unsafe { ptr::read(item) })
767    }
768
769    #[inline]
770    fn size_hint(&self) -> (usize, Option<usize>) {
771        self.iter.size_hint()
772    }
773}
774
775impl<T, const N: usize, A> DoubleEndedIterator for BSTDrain<'_, T, N, A>
776where
777    A: SelflessAllocator,
778{
779    #[inline]
780    fn next_back(&mut self) -> Option<Self::Item> {
781        self.iter.next_back().map(|item| unsafe { ptr::read(item) })
782    }
783}
784
785impl<T, const N: usize, A> ExactSizeIterator for BSTDrain<'_, T, N, A>
786where
787    A: SelflessAllocator,
788{
789    #[inline]
790    fn len(&self) -> usize {
791        self.iter.len()
792    }
793}
794
795impl<T, const N: usize, A: SelflessAllocator> Drop for BSTDrain<'_, T, N, A> {
796    fn drop(&mut self) {
797        // Copyright (c) 2018 The Servo Project Developers
798        // SPDX-License-Identifier: Apache-2.0 OR MIT
799        // https://github.com/servo/rust-smallvec/blob/v2/src/lib.rs#L3
800
801        if core::mem::needs_drop::<T>() {
802            self.for_each(drop);
803        }
804
805        // Copy backward data not subject to drain to the drained start location
806        if self.tail_len > 0 {
807            unsafe {
808                let array = self.array.as_mut();
809
810                let start = array.len();
811                let tail = self.tail_start;
812                if let Some(ptr) = array.as_non_null_ptr() {
813                    if tail != start {
814                        let src = ptr.add(tail).as_ptr();
815                        let dst = ptr.add(start).as_ptr();
816                        ptr::copy(src, dst, self.tail_len);
817                    }
818                }
819                array.size = (start + self.tail_len) as u32;
820            }
821        }
822    }
823}
824
825////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
826
827pub struct BSTSmallArrayIterator<'a, T, const N: usize, A>
828where
829    A: SelflessAllocator,
830{
831    array: &'a BSTSmallArray<T, N, A>,
832    index: usize,
833}
834
835impl<'a, T, const N: usize, A> Iterator for BSTSmallArrayIterator<'a, T, N, A>
836where
837    A: SelflessAllocator,
838{
839    type Item = &'a T;
840
841    #[inline]
842    fn next(&mut self) -> Option<Self::Item> {
843        if self.index < self.array.len() {
844            let item = unsafe { self.array.as_const_non_null_ptr()?.add(self.index).as_ref() };
845            self.index += 1;
846            Some(item)
847        } else {
848            None
849        }
850    }
851}
852
853pub struct BSTSmallArrayIntoIterator<T, const N: usize, A>
854where
855    A: SelflessAllocator,
856{
857    array: BSTSmallArray<T, N, A>,
858    index: usize,
859}
860
861impl<T, const N: usize, A> Iterator for BSTSmallArrayIntoIterator<T, N, A>
862where
863    A: SelflessAllocator,
864{
865    type Item = T;
866
867    #[inline]
868    fn next(&mut self) -> Option<Self::Item> {
869        if self.index < self.array.len() {
870            let item = unsafe { ptr::read(self.array.as_non_null_ptr()?.add(self.index).as_ptr()) };
871            self.index += 1;
872            Some(item)
873        } else {
874            None
875        }
876    }
877}
878
879impl<T, const N: usize, A> IntoIterator for BSTSmallArray<T, N, A>
880where
881    A: SelflessAllocator,
882{
883    type Item = T;
884    type IntoIter = BSTSmallArrayIntoIterator<T, N, A>;
885
886    #[inline]
887    fn into_iter(self) -> Self::IntoIter {
888        BSTSmallArrayIntoIterator { array: self, index: 0 }
889    }
890}
891
892impl<'a, T, const N: usize, A> IntoIterator for &'a BSTSmallArray<T, N, A>
893where
894    A: SelflessAllocator,
895{
896    type Item = &'a T;
897    type IntoIter = BSTSmallArrayIterator<'a, T, N, A>;
898
899    #[inline]
900    fn into_iter(self) -> Self::IntoIter {
901        BSTSmallArrayIterator { array: self, index: 0 }
902    }
903}
904
905pub struct BSTSmallArrayIterMut<'a, T, const N: usize, A>
906where
907    A: SelflessAllocator,
908{
909    array: &'a mut BSTSmallArray<T, N, A>,
910    index: usize,
911}
912
913impl<'a, T, const N: usize, A> Iterator for BSTSmallArrayIterMut<'a, T, N, A>
914where
915    A: SelflessAllocator,
916{
917    type Item = &'a mut T;
918
919    #[inline]
920    fn next(&mut self) -> Option<Self::Item> {
921        if self.index < self.array.len() {
922            unsafe {
923                let mut ptr = self.array.as_non_null_ptr()?.add(self.index);
924                self.index += 1;
925                Some(ptr.as_mut())
926            }
927        } else {
928            None
929        }
930    }
931
932    #[inline]
933    fn size_hint(&self) -> (usize, Option<usize>) {
934        let len = self.array.len();
935        (len, Some(len))
936    }
937}
938
939impl<'a, T, const N: usize, A> IntoIterator for &'a mut BSTSmallArray<T, N, A>
940where
941    A: SelflessAllocator,
942{
943    type Item = &'a mut T;
944    type IntoIter = BSTSmallArrayIterMut<'a, T, N, A>;
945
946    #[inline]
947    fn into_iter(self) -> Self::IntoIter {
948        BSTSmallArrayIterMut { array: self, index: 0 }
949    }
950}
951
952impl<T, const N: usize, A> Extend<T> for BSTSmallArray<T, N, A>
953where
954    A: SelflessAllocator,
955{
956    #[inline]
957    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {
958        for elem in iter {
959            self.push(elem);
960        }
961    }
962}
963
964////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
965// Standard derive
966
967impl<T, const N: usize, A> core::fmt::Debug for BSTSmallArray<T, N, A>
968where
969    T: core::fmt::Debug,
970    A: SelflessAllocator,
971{
972    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
973        f.debug_list().entries(self.as_slice()).finish()
974    }
975}
976
977impl<T, const N: usize> Default for BSTSmallArray<T, N> {
978    #[inline]
979    fn default() -> Self {
980        Self::new()
981    }
982}
983
984impl<T: Clone, const N: usize, A> Clone for BSTSmallArray<T, N, A>
985where
986    A: SelflessAllocator + Clone,
987{
988    fn clone(&self) -> Self {
989        let mut new = if self.capacity as usize > N {
990            // allocate heap
991            let layout = Self::new_layout(self.capacity);
992            let new_ptr = A::allocate(layout).map_or_else(
993                |_| handle_alloc_error(layout),
994                |p| unsafe { NonNull::new_unchecked(p.cast::<T>().as_ptr()) },
995            );
996
997            Self {
998                data: RawBSTSmallArray::new_heap(new_ptr),
999                capacity: self.capacity,
1000                storage_type: StorageType::Heap.into(),
1001                size: self.size,
1002                alloc: PhantomData,
1003            }
1004        } else {
1005            Self {
1006                data: RawBSTSmallArray::new(),
1007                capacity: self.capacity,
1008                storage_type: StorageType::Inline.into(),
1009                size: self.size,
1010                alloc: PhantomData,
1011            }
1012        };
1013
1014        // Clone elements
1015        let count = self.size as usize;
1016        for i in 0..count {
1017            let src = &self[i];
1018            new.push(src.clone());
1019        }
1020
1021        new
1022    }
1023}
1024
1025impl<T, const N: usize, A> PartialEq for BSTSmallArray<T, N, A>
1026where
1027    T: PartialEq,
1028    A: SelflessAllocator,
1029{
1030    #[inline]
1031    fn eq(&self, other: &Self) -> bool {
1032        self.as_slice() == other.as_slice()
1033    }
1034}
1035impl<T, const N: usize, A> PartialEq<Vec<T>> for BSTSmallArray<T, N, A>
1036where
1037    T: PartialEq,
1038    A: SelflessAllocator,
1039{
1040    #[inline]
1041    fn eq(&self, other: &Vec<T>) -> bool {
1042        self.as_slice() == *other
1043    }
1044}
1045
1046impl<T, const N: usize, A> PartialEq<&[T]> for BSTSmallArray<T, N, A>
1047where
1048    T: PartialEq,
1049    A: SelflessAllocator,
1050{
1051    #[inline]
1052    fn eq(&self, other: &&[T]) -> bool {
1053        self.as_slice() == *other
1054    }
1055}
1056
1057impl<T, const N: usize, A> Eq for BSTSmallArray<T, N, A>
1058where
1059    T: Eq,
1060    A: SelflessAllocator,
1061{
1062}
1063
1064impl<T, const N: usize, A> PartialOrd for BSTSmallArray<T, N, A>
1065where
1066    T: PartialOrd,
1067    A: SelflessAllocator,
1068{
1069    #[inline]
1070    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
1071        self.as_slice().partial_cmp(other.as_slice())
1072    }
1073}
1074
1075impl<T, const N: usize, A> Ord for BSTSmallArray<T, N, A>
1076where
1077    T: Ord,
1078    A: SelflessAllocator,
1079{
1080    #[inline]
1081    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
1082        self.as_slice().cmp(other.as_slice())
1083    }
1084}
1085
1086impl<T, const N: usize, A> core::hash::Hash for BSTSmallArray<T, N, A>
1087where
1088    T: core::hash::Hash,
1089    A: SelflessAllocator,
1090{
1091    #[inline]
1092    fn hash<H: core::hash::Hasher>(&self, state: &mut H) {
1093        self.as_slice().hash(state);
1094    }
1095}
1096
1097////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
1098
1099#[cfg(test)]
1100mod tests {
1101    use stdx::alloc::Global;
1102
1103    use super::BSTSmallArray as BSTSmallArray_;
1104
1105    type BSTSmallArray<T, const N: usize = 10, A = Global> = BSTSmallArray_<T, N, A>;
1106
1107    #[test]
1108    fn test_drain() {
1109        let mut array = BSTSmallArray::<_, 10>::new();
1110        array.push(1);
1111        array.push(2);
1112        array.push(3);
1113        array.push(4);
1114        array.push(5);
1115
1116        let drained = array.drain(1..3);
1117        // for i in drained {
1118        //     println!("{:?}", i);
1119        // }
1120        assert_eq!(drained.collect::<Vec<_>>(), vec![2, 3]);
1121        assert_eq!(array.len(), 3);
1122        assert_eq!(array[0], 1);
1123        assert_eq!(array[1], 4);
1124        assert_eq!(array[2], 5);
1125    }
1126
1127    #[test]
1128    fn test_shrink_to_fit_on_stack() {
1129        let mut array = BSTSmallArray::<_, 10>::new();
1130        array.push(1);
1131        array.push(2);
1132        array.push(3);
1133        array.push(4);
1134        array.push(5);
1135
1136        assert_eq!(array.len(), 5);
1137        assert_eq!(array.capacity(), 10);
1138
1139        array.shrink_to_fit();
1140
1141        assert_eq!(array.len(), 5);
1142        // In the case of stack, it is determined at compile time, so nothing is done.
1143        assert_eq!(array.capacity(), 10);
1144    }
1145
1146    #[test]
1147    fn test_shrink_to_fit_on_heap() {
1148        let mut array = BSTSmallArray::<_, 4>::new();
1149        array.push(1);
1150        array.push(2);
1151        array.push(3);
1152        array.push(4);
1153        array.push(5); // heap
1154
1155        assert_eq!(array.len(), 5);
1156        assert_eq!(array.capacity(), 4 * 2); // grow to 8
1157
1158        array.shrink_to_fit();
1159
1160        assert_eq!(array.len(), 5);
1161        assert_eq!(array.capacity(), 5);
1162    }
1163}