copy_stack_vec/vec/array/
drain.rs

1// This file is part of copy-stack-vec.
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4// Crate imports
5use crate::{iter::IntoIter, vec::CopyStackVec};
6
7// Core imports
8use core::{
9    iter::FusedIterator,
10    ops::{Bound, RangeBounds},
11};
12
13/// Owned iterator returned by `CopyStackVec::drain`.
14///
15/// - Holds a mutable borrow of the parent vector for the iterator's lifetime.
16/// - Internally just wraps an `IntoIter` over a temporary `CopyStackVec`
17///   containing the drained elements.
18pub struct Drain<'a, T: Copy, const N: usize> {
19    pub(crate) _parent: &'a mut CopyStackVec<T, N>,
20    pub(crate) iter: IntoIter<T, N>,
21}
22
23impl<'a, T: Copy, const N: usize> Iterator for Drain<'a, T, N> {
24    type Item = T;
25
26    fn next(&mut self) -> Option<T> {
27        self.iter.next()
28    }
29
30    fn size_hint(&self) -> (usize, Option<usize>) {
31        self.iter.size_hint()
32    }
33
34    fn nth(&mut self, n: usize) -> Option<T> {
35        self.iter.nth(n)
36    }
37}
38
39impl<'a, T: Copy, const N: usize> DoubleEndedIterator for Drain<'a, T, N> {
40    fn next_back(&mut self) -> Option<T> {
41        self.iter.next_back()
42    }
43
44    fn nth_back(&mut self, n: usize) -> Option<T> {
45        self.iter.nth_back(n)
46    }
47}
48
49impl<'a, T: Copy, const N: usize> ExactSizeIterator for Drain<'a, T, N> {}
50impl<'a, T: Copy, const N: usize> FusedIterator for Drain<'a, T, N> {}
51
52impl<T: Copy + Default, const N: usize> CopyStackVec<T, N> {
53    /// Drains the specified range of elements and returns them as an iterator.
54    ///
55    /// The elements in `range` are removed immediately and yielded by value.
56    /// The tail of the vector is shifted left to fill the gap.
57    ///
58    /// This follows the same semantics as [`Vec::drain`]:
59    ///
60    /// # Panics
61    ///
62    /// Panics if the range is invalid:
63    /// - `start > end`
64    /// - `end > self.len()`
65    ///
66    /// (Note: `start == end` is allowed and produces an empty iterator.)
67    ///
68    /// # Examples
69    /// ```
70    /// # use copy_stack_vec::CopyStackVec;
71    /// let mut v: CopyStackVec<_, 4> = [1, 2, 3, 4].into();
72    /// let drained: CopyStackVec<_, 4> = v.drain(1..3).collect();
73    /// assert_eq!(drained.as_slice(), &[2, 3]);
74    /// assert_eq!(v.as_slice(), &[1, 4]);
75    /// ```
76    pub fn drain<R>(
77        &mut self,
78        range: R,
79    ) -> impl DoubleEndedIterator<Item = T> + ExactSizeIterator + FusedIterator + '_
80    where
81        R: RangeBounds<usize>,
82    {
83        let len = self.len();
84
85        let start = match range.start_bound() {
86            Bound::Included(&i) => i,
87            Bound::Excluded(&i) => i + 1,
88            Bound::Unbounded => 0,
89        };
90        let end = match range.end_bound() {
91            Bound::Included(&i) => i + 1,
92            Bound::Excluded(&i) => i,
93            Bound::Unbounded => len,
94        };
95
96        if start > end {
97            panic!("drain range start > end: {} > {}", start, end);
98        }
99        if end > len {
100            panic!("drain range end {} exceeds length {}", end, len);
101        }
102
103        if start == end {
104            // Empty drain
105            let tmp: CopyStackVec<T, N> = CopyStackVec::default();
106            return Drain {
107                _parent: self,
108                iter: tmp.into_iter(),
109            };
110        }
111
112        // Copy drained items into temporary vec
113        let slice = &self.as_slice()[start..end];
114        let mut tmp: CopyStackVec<T, N> = CopyStackVec::default();
115        tmp.extend_from_slice(slice).unwrap();
116
117        // Shift tail left
118        let range_len = end - start;
119        let tail_len = len - end;
120        if tail_len > 0 {
121            self.as_mut_slice().copy_within(end..len, start);
122        }
123        self.len = len - range_len;
124
125        Drain {
126            _parent: self,
127            iter: tmp.into_iter(),
128        }
129    }
130}
copy_stack_vec/vec/array/drain.rs

copy_stack_vec/vec/array/
drain.rs
