small_vec2/
splice.rs

1use std::{ptr, slice};
2
3use crate::Drain;
4
5/// A splicing iterator for `SmallVec`.
6///
7/// This struct is created by [`splice`][crate::SmallVec::splice].
8/// See its documentation for more.
9///
10/// # Example
11///
12/// ```
13/// # use small_vec2::{SmallVec, Splice};
14/// let mut v: SmallVec<_, 3> = SmallVec::from(vec![0, 1, 2]);
15/// let new = [7, 8];
16/// let iter: Splice<_, 3> = v.splice(1.., new.iter().cloned());
17/// ```
18#[derive(Debug)]
19pub struct Splice<'a, I: Iterator + 'a, const N: usize> {
20    pub(crate) drain: Drain<'a, I::Item, N>,
21    pub(crate) replace_with: I,
22}
23
24impl<'a, I: Iterator + 'a, const N: usize> Splice<'a, I, N> {
25    pub fn new(drain: Drain<'a, I::Item, N>, replace_with: I) -> Self {
26        Self {
27            drain,
28            replace_with,
29        }
30    }
31}
32
33impl<I: Iterator, const N: usize> Iterator for Splice<'_, I, N> {
34    type Item = I::Item;
35
36    fn next(&mut self) -> Option<Self::Item> {
37        self.drain.next()
38    }
39
40    fn size_hint(&self) -> (usize, Option<usize>) {
41        self.drain.size_hint()
42    }
43}
44
45impl<I: Iterator, const N: usize> DoubleEndedIterator for Splice<'_, I, N> {
46    fn next_back(&mut self) -> Option<Self::Item> {
47        self.drain.next_back()
48    }
49}
50
51impl<I: Iterator, const N: usize> ExactSizeIterator for Splice<'_, I, N> {}
52
53impl<I: Iterator, const N: usize> Drop for Splice<'_, I, N> {
54    fn drop(&mut self) {
55        self.drain.by_ref().for_each(drop);
56
57        unsafe {
58            if self.drain.tail_len == 0 {
59                //这块未测试
60                self.drain.vec.as_mut().extend(self.replace_with.by_ref());
61                return;
62            }
63
64            // First fill the range left by drain().
65            let fill = self.drain.fill(&mut self.replace_with);
66            if !fill {
67                return;
68            }
69
70            // There may be more elements. Use the lower bound as an estimate.
71            // FIXME: Is the upper bound a better guess? Or something else?
72
73            let (lower_bound, _upper_bound) = self.replace_with.size_hint();
74            if lower_bound > 0 {
75                self.drain.move_tail(lower_bound);
76                let fill = self.drain.fill(&mut self.replace_with);
77                if !fill {
78                    return;
79                }
80            }
81            // Collect any remaining elements.
82            // This is a zero-length vector which does not allocate if `lower_bound` was exact.
83            let mut collected = self
84                .replace_with
85                .by_ref()
86                .collect::<Vec<I::Item>>()
87                .into_iter();
88            // Now we have an exact count.
89            if collected.len() > 0 {
90                self.drain.move_tail(collected.len());
91                let filled = self.drain.fill(&mut collected);
92                debug_assert!(filled);
93                debug_assert_eq!(collected.len(), 0);
94            }
95        }
96        // Let `Drain::drop` move the tail back if necessary and restore `vec.len`.
97    }
98}
99/// Private helper methods for [`Splice::drop`]
100impl<T, const N: usize> Drain<'_, T, N> {
101    /// The range from `self.vec.len` to `self.tail_start` contains elements
102    /// that have been moved out.
103    /// Fill that range as much as possible with new elements from the `replace_with` iterator.
104    /// Returns `true` if we filled the entire range. (`replace_with.next()` didn’t return `None`.)
105    pub(crate) unsafe fn fill<I: Iterator<Item = T>>(&mut self, replace_with: &mut I) -> bool {
106        let vec = unsafe { self.vec.as_mut() };
107
108        let range_start = vec.len;
109        let range_end = self.tail_start;
110
111        let data = unsafe { vec.as_mut_ptr().add(range_start) };
112        let range_len = range_end - range_start;
113        let range_slice = unsafe { slice::from_raw_parts_mut(data, range_len) };
114
115        for place in range_slice {
116            if let Some(new_item) = replace_with.next() {
117                unsafe { ptr::write(place, new_item) };
118                vec.len += 1;
119            } else {
120                return false;
121            }
122        }
123        true
124    }
125
126    /// Makes room for inserting more elements before the tail.
127    pub(crate) unsafe fn move_tail(&mut self, additional: usize) {
128        let small_vec = unsafe { self.vec.as_mut() };
129        let len = self.tail_start + self.tail_len;
130
131        let src = unsafe { small_vec.as_ptr().add(self.tail_start) };
132        let src_value = unsafe { ptr::read(src) };
133
134        let a = small_vec.spilled();
135
136        small_vec.reserve_exact(additional + len);
137
138        let new_tail_start = self.tail_start + additional;
139        let b = small_vec.spilled();
140        if a == b {
141            unsafe {
142                let src = small_vec.as_ptr().add(self.tail_start);
143                let dst = small_vec.as_mut_ptr().add(new_tail_start);
144                ptr::copy(src, dst, self.tail_len);
145            }
146        } else {
147            unsafe {
148                let src = &src_value as *const T;
149                let dst = small_vec.as_mut_ptr().add(new_tail_start);
150                ptr::copy(src, dst, self.tail_len);
151            }
152        }
153        self.tail_start = new_tail_start;
154    }
155}