iter_merge/storage/
array.rs

1use core::{
2    cell::Cell,
3    fmt::{Debug, Display},
4    marker::{PhantomData, PhantomPinned},
5    mem::MaybeUninit,
6    pin::Pin,
7};
8
9use crate::{
10    internal::{BaseStorage, PeekIter},
11    merge_iter::{DefaultBuilder, DefaultMergeIter},
12    storage::{Storage as _, debug_formatter},
13};
14
15/// Error signaling an overflow of the array's capacity
16#[derive(Debug, Clone, Copy)]
17pub struct ArrayCapacityOverflow;
18
19impl Display for ArrayCapacityOverflow {
20    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
21        f.write_str("Capacity overflow")
22    }
23}
24
25#[rustversion::since(1.81)]
26impl core::error::Error for ArrayCapacityOverflow {}
27
28/// Fixed-capacity array-based storage for [`MergeIter`](crate::MergeIter)
29pub struct ArrayStorage<const CAP: usize, IT: Iterator> {
30    storage: [MaybeUninit<PeekIter<IT>>; CAP],
31    heap: [MaybeUninit<*mut PeekIter<IT>>; CAP],
32    len: Cell<usize>,
33    _p: PhantomPinned,
34}
35
36impl<const CAP: usize, IT: Iterator> Debug for ArrayStorage<CAP, IT>
37where
38    PeekIter<IT>: Debug,
39{
40    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
41        f.debug_struct("ArrayStorage")
42            .field("CAP", &CAP)
43            .field("len", &self.len)
44            .field("storage", &
45                // SAFETY: array is initialized up to self.len()
46                unsafe {
47                    core::slice::from_raw_parts(
48                        self.storage.as_ptr().cast::<PeekIter<IT>>(),
49                        self.len(),
50                    )
51                })
52            .finish_non_exhaustive()
53    }
54}
55
56#[inline(always)]
57const fn uninit_array<const CAP: usize, T>() -> [MaybeUninit<T>; CAP] {
58    // SAFETY: array of MaybeUninit does not need initialization
59    unsafe { MaybeUninit::<[MaybeUninit<T>; CAP]>::uninit().assume_init() }
60}
61
62impl<IT: Iterator> ArrayStorage<0, IT> {
63    /// Create [`ArrayStorage`] with given capacity and inferred iterator type
64    #[must_use]
65    #[inline(always)]
66    pub const fn with_capacity<const CAP: usize>() -> ArrayStorage<CAP, IT> {
67        ArrayStorage::new()
68    }
69}
70
71impl<const CAP: usize, IT: Iterator> ArrayStorage<CAP, IT> {
72    /// Create a new [`ArrayStorage`]
73    ///
74    /// # Example
75    /// Building a merge iterator from an `ArrayStorage`.
76    ///
77    /// ```
78    /// use core::{iter, pin::pin};
79    ///
80    /// use iter_merge::ArrayStorage;
81    ///
82    /// let mut storage: ArrayStorage<5, _> = ArrayStorage::new();
83    /// storage.push(iter::once(2));
84    /// storage.push(iter::once(1));
85    /// let storage = pin!(storage);
86    /// let it = storage.build();
87    /// assert!(it.eq([1, 2]));
88    /// ```
89    #[must_use]
90    #[inline(always)]
91    pub const fn new() -> Self {
92        Self {
93            storage: uninit_array(),
94            heap: uninit_array(),
95            len: Cell::new(0),
96            _p: PhantomPinned,
97        }
98    }
99
100    /// Creates a new [`ArrayStorage`] with the same `CAP` as
101    /// the provided array.
102    ///
103    /// # Example
104    /// Building a merge iterator from an `ArrayStorage`.
105    ///
106    /// ```
107    /// use core::{iter, pin::pin};
108    ///
109    /// use iter_merge::ArrayStorage;
110    /// let storage = ArrayStorage::from_arr([[1, 3], [2, 4]]);
111    /// assert_eq!(storage.capacity(), 2)
112    /// ```
113    #[must_use]
114    #[inline]
115    pub fn from_arr<T: IntoIterator<IntoIter = IT>>(iters: [T; CAP]) -> Self {
116        let mut res = Self::new();
117        res.extend(iters);
118        res
119    }
120
121    /// Returns the number of non-empty iterators stored in [`ArrayStorage`]
122    #[inline]
123    pub fn len(&self) -> usize {
124        self.len.get()
125    }
126
127    /// Returns the (fixed) capacity of [`ArrayStorage`]
128    #[inline]
129    pub const fn capacity(&self) -> usize {
130        CAP
131    }
132
133    /// Returns `true` if this [`ArrayStorage`] is empty
134    #[inline]
135    pub fn is_empty(&self) -> bool {
136        self.len() == 0
137    }
138
139    /// Appends an element to the back of a collection.
140    ///
141    /// # Panics
142    ///
143    /// Panics if the collection is full.
144    pub fn push<Iter>(&mut self, iter: Iter)
145    where
146        Iter: IntoIterator<IntoIter = IT>,
147    {
148        self.try_push(iter).unwrap();
149    }
150
151    /// Tries to append an element to the back of a collection.
152    /// # Errors
153    /// Returns error if the [`ArrayStorage`] is full
154    pub fn try_push<Iter>(&mut self, iter: Iter) -> Result<(), ArrayCapacityOverflow>
155    where
156        Iter: IntoIterator<IntoIter = IT>,
157    {
158        if let Some(peek_iter) = PeekIter::new_from_iter(iter) {
159            let len = self.len.get();
160            if len >= CAP {
161                return Err(ArrayCapacityOverflow);
162            }
163            self.storage[len].write(peek_iter);
164            self.len.set(len.checked_add(1).expect("unreachable"));
165        }
166        Ok(())
167    }
168
169    /// Constructs a [`Builder`] from this storage.
170    ///
171    /// Note: the storage cannot move for [`MergeIter`](crate::MergeIter) to work, thus
172    /// you need to call this method on a pinned mutable reference.
173    ///
174    /// Example:
175    /// ```
176    /// use core::{iter, pin::pin};
177    ///
178    /// use iter_merge::ArrayStorage;
179    /// let mut storage = ArrayStorage::<5, _>::new();
180    /// storage.push(iter::once(1));
181    /// let storage = pin!(storage);
182    /// let _builder = storage.into_builder();
183    /// ```
184    #[must_use]
185    pub fn into_builder(self: Pin<&mut Self>) -> DefaultBuilder<InternalArrayStorage<'_, IT>> {
186        let len = self.len.replace(0);
187        debug_assert!(len <= CAP);
188        let (storage, heap) = {
189            // SAFETY: we're never moving the data out of mut_ref, we're just copying the
190            // mut pointers.
191            // InternalArrayStorage lives for 'a, same as our pinned pointer
192            // during this time it's safe to rely on pin guarantee
193            let mut_ref = unsafe { Pin::get_unchecked_mut(self) };
194            (
195                mut_ref.storage.as_mut_ptr().cast::<PeekIter<IT>>(),
196                mut_ref.heap.as_mut_ptr().cast::<*mut PeekIter<IT>>(),
197            )
198        };
199        for i in 0..len {
200            // SAFETY: storage pointer is valid for adding up to CAP (>= len), heap - for writitng
201            //         up to CAP (>= len).
202            //         self is pinned up to 'a, so we are relying on pin guarantee by constructing
203            //         InternalArrayStorage valid for 'a
204            unsafe {
205                heap.add(i).write(storage.add(i));
206            }
207        }
208        InternalArrayStorage {
209            heap,
210            len,
211            _p: PhantomData,
212        }
213        .into_builder()
214    }
215
216    /// Constructs a [`MergeIter`] from this storage with default parameters.
217    ///
218    /// Equivalent to calling <code>[Self::into_builder()].[build()](crate::merge_iter::Builder::build)</code>
219    #[must_use]
220    pub fn build(self: Pin<&mut Self>) -> DefaultMergeIter<InternalArrayStorage<'_, IT>>
221    where
222        IT::Item: Ord,
223    {
224        self.into_builder().build()
225    }
226}
227
228impl<const CAP: usize, IT, Item> FromIterator<Item> for ArrayStorage<CAP, IT>
229where
230    IT: Iterator,
231    Item: IntoIterator<IntoIter = IT>,
232{
233    fn from_iter<T: IntoIterator<Item = Item>>(iter: T) -> Self {
234        let mut res = Self::new();
235        res.extend(iter);
236        res
237    }
238}
239
240impl<const CAP: usize, IT: Iterator, A> Extend<A> for ArrayStorage<CAP, IT>
241where
242    A: IntoIterator<IntoIter = IT>,
243{
244    fn extend<T: IntoIterator<Item = A>>(&mut self, iter: T) {
245        for item in iter {
246            self.push(item);
247        }
248    }
249}
250
251impl<const CAP: usize, IT: Iterator> Default for ArrayStorage<CAP, IT> {
252    #[inline]
253    fn default() -> Self {
254        Self::new()
255    }
256}
257
258impl<const CAP: usize, IT: Iterator> Drop for ArrayStorage<CAP, IT> {
259    fn drop(&mut self) {
260        // We are dropping potentially pinned value.
261        // As soon as the [`build`] method is called
262        // (first method that relies on pinning guarantees)
263        // the self.len is set to 0 and never modified until pinned borrow ends,
264        // so this loop is a noop and we're not violating any guarantees of pin.
265        for i in 0..self.len.replace(0) {
266            // SAFETY: up to self.len items are initialized, the pointers were not given
267            // to Heap that could've invalidated some stored items.
268            // self.len is replaced by 0, so there's no possibility of double-free,
269            // only memory leak if the item drop code panics
270            unsafe {
271                self.storage[i].assume_init_drop();
272            }
273        }
274    }
275}
276
277/// Internal representation of the [`ArrayStorage`] that's actually used as the
278/// [`MergeIter`](crate::MergeIter)'s [`Storage`](crate::internal::BaseStorage) backend.
279pub struct InternalArrayStorage<'a, IT: Iterator> {
280    heap: *mut *mut PeekIter<IT>,
281    len: usize,
282    // represents us holding the pinned ArrayStorage, capacity is irrelevant,
283    // this is only for lifetime management
284    _p: PhantomData<Pin<&'a mut ArrayStorage<1, IT>>>,
285}
286
287unsafe impl<IT: Iterator> BaseStorage for InternalArrayStorage<'_, IT> {
288    type IT = IT;
289
290    #[inline]
291    fn heap(&self) -> *mut *mut PeekIter<IT> {
292        self.heap
293    }
294
295    #[inline]
296    fn len(&self) -> usize {
297        self.len
298    }
299
300    #[inline]
301    unsafe fn set_len(&mut self, new_len: usize) {
302        self.len = new_len;
303    }
304}
305
306impl<IT: Iterator> Debug for InternalArrayStorage<'_, IT>
307where
308    PeekIter<<Self as BaseStorage>::IT>: Debug,
309{
310    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
311        f.debug_struct("InternalArrayStorage")
312            .field("len", &self.len)
313            .field("storage", &debug_formatter(self))
314            .finish_non_exhaustive()
315    }
316}
317
318impl<IT: Iterator> Drop for InternalArrayStorage<'_, IT> {
319    fn drop(&mut self) {
320        crate::storage::StorageOps::clear(self);
321        // The storage itself is owned by ArrayStorage and will be deallocated by it
322    }
323}
324
325// SAFETY: InternalArrayStorage is just a reference to pinned ArrayStorage.
326// It's safe for them to be send and sync, if the `Pin<&'a mut ArrayStorage<IT>>` is send and sync
327// respectively
328unsafe impl<'a, IT> Send for InternalArrayStorage<'a, IT>
329where
330    IT: Iterator,
331    Pin<&'a mut ArrayStorage<1, IT>>: Send,
332{
333}
334
335// SAFETY: see above.
336unsafe impl<'a, IT> Sync for InternalArrayStorage<'a, IT>
337where
338    IT: Iterator,
339    Pin<&'a mut ArrayStorage<1, IT>>: Sync,
340{
341}
342
343#[cfg(test)]
344mod tests {
345    use super::*;
346
347    #[test]
348    fn capacity_overflow() {
349        let mut s: ArrayStorage<1, _> = ArrayStorage::default();
350        assert_eq!(s.len(), 0);
351        assert!(s.is_empty());
352        s.push([1, 2, 3]);
353        assert_eq!(s.len(), 1);
354        assert!(!s.is_empty());
355        assert!(matches!(s.try_push([4, 5, 6]), Err(ArrayCapacityOverflow)));
356    }
357}
iter_merge/storage/array.rs

iter_merge/storage/
array.rs
