tiny_str/
drain.rs

1//! [drain](TinyString::drain) implementation
2
3use core::iter::FusedIterator;
4use core::marker::PhantomData;
5use core::mem::ManuallyDrop;
6use core::ops::{Range, RangeBounds};
7use core::ptr::{self, NonNull};
8
9use alloc::{slice, str};
10
11use crate::TinyString;
12
13/// A draining iterator for [TinyString]
14///
15/// This struct is created by the [TinyString::drain] method
16pub struct Drain<'a, const N: usize> {
17    string: NonNull<TinyString<N>>,
18    remaining_start: usize,
19    remaining_len: usize,
20    tail_start: usize,
21    tail_len: usize,
22    _marker: PhantomData<&'a mut TinyString<N>>,
23}
24
25impl<const N: usize> Drain<'_, N> {
26
27    /// Returns the remaining slice
28    pub const fn as_str(&mut self) -> &str {
29        if self.remaining_len == 0 {
30            return ""
31        }
32        unsafe {
33            let str = self.string.as_ref().as_ptr().add(self.remaining_start);
34            let slice = slice::from_raw_parts(str, self.remaining_len);
35            str::from_utf8_unchecked(slice)
36        }
37    }
38
39    /// Consumes this [Drain], preserving any un-yielded characters
40    ///
41    /// # Example
42    /// ```
43    /// use tiny_str::TinyString;
44    ///
45    /// let mut tv = TinyString::<10>::from("abcdefghijk");
46    ///
47    /// let mut drain = tv.drain(3..7);
48    ///
49    /// assert_eq!(drain.next(), Some('d'));
50    /// assert_eq!(drain.next(), Some('e'));
51    ///
52    /// drain.keep_rest();
53    ///
54    /// assert_eq!(tv.as_str(), "abcfghijk");
55    /// ```
56    pub fn keep_rest(self) {
57        let mut slf = ManuallyDrop::new(self);
58
59        /* [ HEAD ] [ yieled  | remaining | yielded_back ] [ TAIL ] */
60        unsafe {
61            let string = slf.string.as_mut();
62            let start = string.len();
63
64            let buf = string.as_mut_ptr();
65
66            let dst = buf.add(start);
67            let src = buf.add(slf.remaining_start);
68            /* First: We move the remaining chunk to the start */
69            ptr::copy(src, dst, slf.remaining_len);
70
71            let src = buf.add(slf.tail_start);
72            let dst = dst.add(slf.remaining_len);
73            /* Now we move the tail */
74            ptr::copy(src, dst, slf.tail_len);
75
76            string.buf.set_len(start + slf.remaining_len + slf.tail_len);
77        }
78    }
79}
80
81impl<const N: usize> Iterator for Drain<'_, N> {
82    type Item = char;
83
84    fn next(&mut self) -> Option<Self::Item> {
85        let c = self.as_str().chars().next()?;
86        self.remaining_start += c.len_utf8();
87        self.remaining_len -= c.len_utf8();
88        Some(c)
89    }
90
91    fn size_hint(&self) -> (usize, Option<usize>) {
92        (self.remaining_len, Some(self.remaining_len))
93    }
94}
95
96impl<const N: usize> DoubleEndedIterator for Drain<'_, N> {
97    fn next_back(&mut self) -> Option<Self::Item> {
98        let c = self.as_str().chars().next_back()?;
99        self.remaining_len -= c.len_utf8();
100        Some(c)
101    }
102}
103
104impl<const N: usize> ExactSizeIterator for Drain<'_, N> {}
105
106impl<const N: usize> FusedIterator for Drain<'_, N> {}
107
108impl<const N: usize> Drop for Drain<'_, N> {
109    fn drop(&mut self) {
110        unsafe {
111            let string = self.string.as_mut();
112            let len = string.len();
113            let ptr = string.as_mut_ptr();
114
115            let src = ptr.add(self.tail_start);
116            let dst = ptr.add(len);
117            ptr::copy(src, dst, self.tail_len);
118
119            string.buf.set_len(len + self.tail_len);
120        }
121    }
122}
123
124impl<const N: usize> TinyString<N> {
125
126    /// Removes the substring indicated by the given range from the string,
127    /// returning a double-ended iterator over the removed substring.
128    ///
129    /// If the iterator is dropped before being fully consumed,
130    /// it drops the remaining removed elements.
131    ///
132    /// # Panics
133    ///
134    /// Panics if the starting point is greater than the end point, if
135    /// the end point is greater than the length of the vector,
136    /// or if any of the ends of the range fall outside a char boundary.
137    ///
138    /// # Leaking
139    ///
140    /// If the returned iterator goes out of scope without being dropped (due to
141    /// [`core::mem::forget`], for example), the vector may have lost characters
142    /// arbitrarily, including characters outside the range.
143    ///
144    /// # Examples
145    ///
146    /// ```
147    /// use tiny_str::TinyString;
148    /// let mut s = TinyString::<10>::from("abcdefghijk");
149    ///
150    /// let mut drain = s.drain(2..=4);
151    /// assert_eq!(drain.next(), Some('c'));
152    /// assert_eq!(drain.next(), Some('d'));
153    /// assert_eq!(drain.next(), Some('e'));
154    /// assert_eq!(drain.next(), None);
155    /// drop(drain);
156    ///
157    /// assert_eq!(s.as_str(), "abfghijk");
158    ///
159    /// // A full range clears the string, like `clear()` does
160    /// s.drain(..);
161    /// assert_eq!(s.as_str(), "");
162    /// ```
163    pub fn drain<R>(&mut self, range: R) -> Drain<N>
164    where
165        R: RangeBounds<usize>
166    {
167        let len = self.len();
168        let Range { start, end } = self.slice_range(range, len);
169        unsafe {
170            self.buf.set_len(start);
171
172            let string = NonNull::new_unchecked(self as *mut _);
173            Drain {
174                tail_len: len - end,
175                tail_start: end,
176                remaining_start: start,
177                remaining_len: end - start,
178                string,
179                _marker: PhantomData,
180            }
181        }
182    }
183}