spec_extend.rs - source
Description: Source of the Rust file `library/alloc/src/collections/vec_deque/spec_extend.rs`.

#
alloc/collections/vec_deque/
spec_extend.rs

    1use core::iter::{Copied, Rev, TrustedLen};
    2use core::slice;
    3
    4use super::{Drain, VecDeque};
    5use crate::alloc::Allocator;
    6#[cfg(not(test))]
    7use crate::vec;
    8
    9// Specialization trait used for VecDeque::extend
    10pub(super) trait SpecExtend<T, I> {
    11    fn spec_extend(&mut self, iter: I);
    12}
    13
    14impl<T, I, A: Allocator> SpecExtend<T, I> for VecDeque<T, A>
    15where
    16    I: Iterator<Item = T>,
    17{
    18    default fn spec_extend(&mut self, mut iter: I) {
    19        // This function should be the moral equivalent of:
    20        //
    21        // for item in iter {
    22        //     self.push_back(item);
    23        // }
    24
    25        while let Some(element) = iter.next() {
    26            let (lower, _) = iter.size_hint();
    27            self.reserve(lower.saturating_add(1));
    28
    29            // SAFETY: We just reserved space for at least one element.
    30            unsafe { self.push_unchecked(element) };
    31
    32            // Inner loop to avoid repeatedly calling `reserve`.
    33            while self.len < self.capacity() {
    34                let Some(element) = iter.next() else {
    35                    return;
    36                };
    37                // SAFETY: The loop condition guarantees that `self.len() < self.capacity()`.
    38                unsafe { self.push_unchecked(element) };
    39            }
    40        }
    41    }
    42}
    43
    44impl<T, I, A: Allocator> SpecExtend<T, I> for VecDeque<T, A>
    45where
    46    I: TrustedLen<Item = T>,
    47{
    48    default fn spec_extend(&mut self, iter: I) {
    49        // This is the case for a TrustedLen iterator.
    50        let (low, high) = iter.size_hint();
    51        if let Some(additional) = high {
    52            debug_assert_eq!(
    53                low,
    54                additional,
    55                "TrustedLen iterator's size hint is not exact: {:?}",
    56                (low, high)
    57            );
    58            self.reserve(additional);
    59
    60            let written = unsafe {
    61                self.write_iter_wrapping(self.to_physical_idx(self.len), iter, additional)
    62            };
    63
    64            debug_assert_eq!(
    65                additional, written,
    66                "The number of items written to VecDeque doesn't match the TrustedLen size hint"
    67            );
    68        } else {
    69            // Per TrustedLen contract a `None` upper bound means that the iterator length
    70            // truly exceeds usize::MAX, which would eventually lead to a capacity overflow anyway.
    71            // Since the other branch already panics eagerly (via `reserve()`) we do the same here.
    72            // This avoids additional codegen for a fallback code path which would eventually
    73            // panic anyway.
    74            panic!("capacity overflow");
    75        }
    76    }
    77}
    78
    79#[cfg(not(test))]
    80impl<T, A1: Allocator, A2: Allocator> SpecExtend<T, vec::IntoIter<T, A2>> for VecDeque<T, A1> {
    81    fn spec_extend(&mut self, mut iterator: vec::IntoIter<T, A2>) {
    82        let slice = iterator.as_slice();
    83        self.reserve(slice.len());
    84
    85        unsafe {
    86            self.copy_slice(self.to_physical_idx(self.len), slice);
    87            self.len += slice.len();
    88        }
    89        iterator.forget_remaining_elements();
    90    }
    91}
    92
    93impl<'a, T: 'a, I, A: Allocator> SpecExtend<&'a T, I> for VecDeque<T, A>
    94where
    95    I: Iterator<Item = &'a T>,
    96    T: Copy,
    97{
    98    default fn spec_extend(&mut self, iterator: I) {
    99        self.spec_extend(iterator.copied())
    100    }
    101}
    102
    103impl<'a, T: 'a, A: Allocator> SpecExtend<&'a T, slice::Iter<'a, T>> for VecDeque<T, A>
    104where
    105    T: Copy,
    106{
    107    fn spec_extend(&mut self, iterator: slice::Iter<'a, T>) {
    108        let slice = iterator.as_slice();
    109        self.reserve(slice.len());
    110
    111        unsafe {
    112            self.copy_slice(self.to_physical_idx(self.len), slice);
    113            self.len += slice.len();
    114        }
    115    }
    116}
    117
    118// Specialization trait used for VecDeque::extend_front
    119pub(super) trait SpecExtendFront<T, I> {
    120    #[track_caller]
    121    fn spec_extend_front(&mut self, iter: I);
    122}
    123
    124impl<T, I, A: Allocator> SpecExtendFront<T, I> for VecDeque<T, A>
    125where
    126    I: Iterator<Item = T>,
    127{
    128    #[track_caller]
    129    default fn spec_extend_front(&mut self, mut iter: I) {
    130        // This function should be the moral equivalent of:
    131        //
    132        // for item in iter {
    133        //     self.push_front(item);
    134        // }
    135
    136        while let Some(element) = iter.next() {
    137            let (lower, _) = iter.size_hint();
    138            self.reserve(lower.saturating_add(1));
    139
    140            // SAFETY: We just reserved space for at least one element.
    141            unsafe { self.push_front_unchecked(element) };
    142
    143            // Inner loop to avoid repeatedly calling `reserve`.
    144            while self.len < self.capacity() {
    145                let Some(element) = iter.next() else {
    146                    return;
    147                };
    148                // SAFETY: The loop condition guarantees that `self.len() < self.capacity()`.
    149                unsafe { self.push_front_unchecked(element) };
    150            }
    151        }
    152    }
    153}
    154
    155#[cfg(not(test))]
    156impl<T, A1: Allocator, A2: Allocator> SpecExtendFront<T, vec::IntoIter<T, A2>> for VecDeque<T, A1> {
    157    #[track_caller]
    158    fn spec_extend_front(&mut self, mut iterator: vec::IntoIter<T, A2>) {
    159        let slice = iterator.as_slice();
    160        self.reserve(slice.len());
    161        // SAFETY: `slice.len()` space was just reserved and elements in the slice are forgotten after this call
    162        unsafe { prepend_reversed(self, slice) };
    163        iterator.forget_remaining_elements();
    164    }
    165}
    166
    167#[cfg(not(test))]
    168impl<T, A1: Allocator, A2: Allocator> SpecExtendFront<T, Rev<vec::IntoIter<T, A2>>>
    169    for VecDeque<T, A1>
    170{
    171    #[track_caller]
    172    fn spec_extend_front(&mut self, iterator: Rev<vec::IntoIter<T, A2>>) {
    173        let mut iterator = iterator.into_inner();
    174        let slice = iterator.as_slice();
    175        self.reserve(slice.len());
    176        // SAFETY: `slice.len()` space was just reserved and elements in the slice are forgotten after this call
    177        unsafe { prepend(self, slice) };
    178        iterator.forget_remaining_elements();
    179    }
    180}
    181
    182impl<'a, T, A: Allocator> SpecExtendFront<T, Copied<slice::Iter<'a, T>>> for VecDeque<T, A>
    183where
    184    Copied<slice::Iter<'a, T>>: Iterator<Item = T>,
    185{
    186    #[track_caller]
    187    fn spec_extend_front(&mut self, iter: Copied<slice::Iter<'a, T>>) {
    188        let slice = iter.into_inner().as_slice();
    189        self.reserve(slice.len());
    190        // SAFETY: `slice.len()` space was just reserved and T is Copy because Copied<slice::Iter<'a, T>> is Iterator
    191        unsafe { prepend_reversed(self, slice) };
    192    }
    193}
    194
    195impl<'a, T, A: Allocator> SpecExtendFront<T, Rev<Copied<slice::Iter<'a, T>>>> for VecDeque<T, A>
    196where
    197    Rev<Copied<slice::Iter<'a, T>>>: Iterator<Item = T>,
    198{
    199    #[track_caller]
    200    fn spec_extend_front(&mut self, iter: Rev<Copied<slice::Iter<'a, T>>>) {
    201        let slice = iter.into_inner().into_inner().as_slice();
    202        self.reserve(slice.len());
    203        // SAFETY: `slice.len()` space was just reserved and T is Copy because Rev<Copied<slice::Iter<'a, T>>> is Iterator
    204        unsafe { prepend(self, slice) };
    205    }
    206}
    207
    208impl<'a, T, A1: Allocator, A2: Allocator> SpecExtendFront<T, Drain<'a, T, A2>> for VecDeque<T, A1> {
    209    #[track_caller]
    210    fn spec_extend_front(&mut self, mut iter: Drain<'a, T, A2>) {
    211        if iter.remaining == 0 {
    212            return;
    213        }
    214
    215        self.reserve(iter.remaining);
    216        unsafe {
    217            // SAFETY: iter.remaining != 0.
    218            let (left, right) = iter.as_slices();
    219            // SAFETY:
    220            // - `iter.remaining` space was reserved, `iter.remaining == left.len() + right.len()`.
    221            // - The elements in `left` and `right` are forgotten after these calls.
    222            prepend_reversed(self, &*left);
    223            prepend_reversed(self, &*right);
    224        }
    225
    226        iter.idx += iter.remaining;
    227        iter.remaining = 0;
    228    }
    229}
    230
    231impl<'a, T, A1: Allocator, A2: Allocator> SpecExtendFront<T, Rev<Drain<'a, T, A2>>>
    232    for VecDeque<T, A1>
    233{
    234    #[track_caller]
    235    fn spec_extend_front(&mut self, iter: Rev<Drain<'a, T, A2>>) {
    236        let mut iter = iter.into_inner();
    237
    238        if iter.remaining == 0 {
    239            return;
    240        }
    241
    242        self.reserve(iter.remaining);
    243        unsafe {
    244            // SAFETY: iter.remaining != 0.
    245            let (left, right) = iter.as_slices();
    246            // SAFETY:
    247            // - `iter.remaining` space was reserved, `iter.remaining == left.len() + right.len()`.
    248            // - The elements in `left` and `right` are forgotten after these calls.
    249            prepend(self, &*right);
    250            prepend(self, &*left);
    251        }
    252
    253        iter.idx += iter.remaining;
    254        iter.remaining = 0;
    255    }
    256}
    257
    258/// Prepends elements of `slice` to `deque` using a copy.
    259///
    260/// # Safety
    261///
    262/// - `deque` must have space for `slice.len()` new elements.
    263/// - Elements of `slice` will be copied into the deque, make sure to forget the elements if `T` is not `Copy`.
    264unsafe fn prepend<T, A: Allocator>(deque: &mut VecDeque<T, A>, slice: &[T]) {
    265    unsafe {
    266        deque.head = deque.wrap_sub(deque.head, slice.len());
    267        deque.copy_slice(deque.head, slice);
    268        deque.len += slice.len();
    269    }
    270}
    271
    272/// Prepends elements of `slice` to `deque` in reverse order using a copy.
    273///
    274/// # Safety
    275///
    276/// - `deque` must have space for `slice.len()` new elements.
    277/// - Elements of `slice` will be copied into the deque, make sure to forget the elements if `T` is not `Copy`.
    278unsafe fn prepend_reversed<T, A: Allocator>(deque: &mut VecDeque<T, A>, slice: &[T]) {
    279    unsafe {
    280        deque.head = deque.wrap_sub(deque.head, slice.len());
    281        deque.copy_slice_reversed(deque.head, slice);
    282        deque.len += slice.len();
    283    }
    284}