summaryrefslogtreecommitdiffstats
path: root/library/alloc/src/collections/vec_deque/spec_extend.rs
blob: 97ff8b76524035e5356fc2e75545a99b927d20db (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
use crate::alloc::Allocator;
use crate::vec;
use core::iter::{ByRefSized, TrustedLen};
use core::slice;

use super::VecDeque;

// Specialization trait used for VecDeque::extend
pub(super) trait SpecExtend<T, I> {
    fn spec_extend(&mut self, iter: I);
}

impl<T, I, A: Allocator> SpecExtend<T, I> for VecDeque<T, A>
where
    I: Iterator<Item = T>,
{
    default fn spec_extend(&mut self, mut iter: I) {
        // This function should be the moral equivalent of:
        //
        //      for item in iter {
        //          self.push_back(item);
        //      }
        while let Some(element) = iter.next() {
            if self.len() == self.capacity() {
                let (lower, _) = iter.size_hint();
                self.reserve(lower.saturating_add(1));
            }

            let head = self.head;
            self.head = self.wrap_add(self.head, 1);
            unsafe {
                self.buffer_write(head, element);
            }
        }
    }
}

impl<T, I, A: Allocator> SpecExtend<T, I> for VecDeque<T, A>
where
    I: TrustedLen<Item = T>,
{
    default fn spec_extend(&mut self, mut iter: I) {
        // This is the case for a TrustedLen iterator.
        let (low, high) = iter.size_hint();
        if let Some(additional) = high {
            debug_assert_eq!(
                low,
                additional,
                "TrustedLen iterator's size hint is not exact: {:?}",
                (low, high)
            );
            self.reserve(additional);

            struct WrapAddOnDrop<'a, T, A: Allocator> {
                vec_deque: &'a mut VecDeque<T, A>,
                written: usize,
            }

            impl<'a, T, A: Allocator> Drop for WrapAddOnDrop<'a, T, A> {
                fn drop(&mut self) {
                    self.vec_deque.head =
                        self.vec_deque.wrap_add(self.vec_deque.head, self.written);
                }
            }

            let mut wrapper = WrapAddOnDrop { vec_deque: self, written: 0 };

            let head_room = wrapper.vec_deque.cap() - wrapper.vec_deque.head;
            unsafe {
                wrapper.vec_deque.write_iter(
                    wrapper.vec_deque.head,
                    ByRefSized(&mut iter).take(head_room),
                    &mut wrapper.written,
                );

                if additional > head_room {
                    wrapper.vec_deque.write_iter(0, iter, &mut wrapper.written);
                }
            }

            debug_assert_eq!(
                additional, wrapper.written,
                "The number of items written to VecDeque doesn't match the TrustedLen size hint"
            );
        } else {
            // Per TrustedLen contract a `None` upper bound means that the iterator length
            // truly exceeds usize::MAX, which would eventually lead to a capacity overflow anyway.
            // Since the other branch already panics eagerly (via `reserve()`) we do the same here.
            // This avoids additional codegen for a fallback code path which would eventually
            // panic anyway.
            panic!("capacity overflow");
        }
    }
}

impl<T, A: Allocator> SpecExtend<T, vec::IntoIter<T>> for VecDeque<T, A> {
    fn spec_extend(&mut self, mut iterator: vec::IntoIter<T>) {
        let slice = iterator.as_slice();
        self.reserve(slice.len());

        unsafe {
            self.copy_slice(self.head, slice);
            self.head = self.wrap_add(self.head, slice.len());
        }
        iterator.forget_remaining_elements();
    }
}

impl<'a, T: 'a, I, A: Allocator> SpecExtend<&'a T, I> for VecDeque<T, A>
where
    I: Iterator<Item = &'a T>,
    T: Copy,
{
    default fn spec_extend(&mut self, iterator: I) {
        self.spec_extend(iterator.copied())
    }
}

impl<'a, T: 'a, A: Allocator> SpecExtend<&'a T, slice::Iter<'a, T>> for VecDeque<T, A>
where
    T: Copy,
{
    fn spec_extend(&mut self, iterator: slice::Iter<'a, T>) {
        let slice = iterator.as_slice();
        self.reserve(slice.len());

        unsafe {
            self.copy_slice(self.head, slice);
            self.head = self.wrap_add(self.head, slice.len());
        }
    }
}