Skip to main content

alloc/vec/
splice.rs

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