allocator_api2/vec/
splice.rs

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