splice.rs - source
Description: Source of the Rust file `library/alloc/src/vec/splice.rs`.

#
alloc/vec/
splice.rs

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