alloc/vec/
spec_extend_with.rs

1use core::clone::TrivialClone;
2use core::{mem, ptr};
3
4use crate::alloc::Allocator;
5use crate::vec::{SetLenOnDrop, Vec};
6
7pub(super) trait SpecExtendWith<T> {
8    /// Extend the vector by `n` clones of value.
9    fn spec_extend_with(&mut self, n: usize, value: T);
10}
11
12impl<T: Clone, A: Allocator> SpecExtendWith<T> for Vec<T, A> {
13    default fn spec_extend_with(&mut self, n: usize, value: T) {
14        self.reserve(n);
15
16        // SAFETY: Since `reserve` guarantees `self.len() + n` <= `capacity`, this
17        // pointer is always within bounds, and valid for writes for the next `n` items.
18        let mut ptr = unsafe { self.as_mut_ptr().add(self.len()) };
19        // Use SetLenOnDrop to work around bug where compiler
20        // might not realize the store through `ptr` through self.set_len()
21        // don't alias.
22        let mut local_len = SetLenOnDrop::new(&mut self.len);
23
24        // Write all elements except the last one
25        for _ in 1..n {
26            // three things have been taken into consideration:
27            // - `T::clone` can be provided by the user, and can panic.
28            // - `T` may implement Drop.
29            // - Since we already have a `value: T`, we can clone `n - 1` times.
30
31            // This loop runs at most `n - 1` times, for all slots except the last one.
32            // As per the guarantee of `reserve`, increments to `ptr` within this range won't wrap around.
33
34            // if clone() panics and starts unwinding;
35            // - the write has not been done yet
36            // - the length will be set by SetLenOnDrop, to the amount of initialized elements
37            // - later, when the Vec is dropped, only initialized elements are dropped.
38            // SAFETY: According to the diagram above, this write is within bounds.
39            unsafe { ptr::write(ptr, value.clone()) };
40            // SAFETY: According to the diagram above, this increment is within bounds of the
41            // allocation and won't wrap around
42            ptr = unsafe { ptr.add(1) };
43            // Increment the length of SetLenOnDrop in every step in case clone() panics to drop the
44            // already-initialized elements.
45            local_len.increment_len(1);
46        }
47
48        if n > 0 {
49            // We can write the last element directly without cloning needlessly
50
51            // SAFETY: According to the diagram above, the last element is not yet written,
52            // so this write is within bounds.
53            unsafe { ptr::write(ptr, value) };
54            local_len.increment_len(1);
55        }
56
57        // len set by scope guard
58    }
59}
60
61impl<T: TrivialClone, A: Allocator> SpecExtendWith<T> for Vec<T, A> {
62    fn spec_extend_with(&mut self, n: usize, value: T) {
63        // The `n == 0` check is already being inserted by the loop canonicalization pass,
64        // so this costs nothing.
65        // https://llvm.org/docs/LoopTerminology.html#more-canonical-loops
66        if n == 0 {
67            // three things:
68            // `self.reserve(0)` does not need to be run, since reserving space for 0 elements
69            // is a no-op. since the `value` argument did not have anywhere to be moved to, it is
70            // dropped here. with this branch present here, we can elide the write through self.len,
71            // since `self.len += 0` would be a no-op.
72            return;
73        }
74
75        self.reserve(n);
76
77        // SAFETY: Since `reserve` guarantees `self.len() + n` <= `capacity`, this
78        // pointer is always within bounds, and valid for writes for the next `n` items.
79        let ptr = unsafe { self.as_mut_ptr().add(self.len()) };
80
81        let mut i = 0;
82        // Write all the elements by copying `value`.
83        // We use an incrementing index here instead of incrementing pointer because the former
84        // has been observed to have less register pressure in the loop initialization.
85        // https://llvm.org/docs/Vectorizers.html#inductions
86        while i < n {
87            // SAFETY: `TrivialClone` indicates that copying `value` is equivalent to
88            // calling `Clone::clone` for `T`. As per the safety comment for `ptr`, this pointer
89            // is within bounds, won't wrap around while `i < n`, and valid for writes. The values
90            // do not overlap since `value` is stored as a function argument, and the target is
91            // in the Vec.
92            unsafe { ptr.add(i).copy_from_nonoverlapping(&value, 1) };
93
94            // Because above is infallible, we don't need to worry about setting the len
95            // in this loop, and instead we just increment the loop induction variable.
96            i += 1;
97        }
98
99        // Above, `n` items have been infallibly copied into the vector
100        self.len += n;
101
102        // The reason we did not treat the last position in the vector specially by doing
103        // `n - 1` copies and semantically moving `value` into it is that we are forced to do `n`
104        // copies of the value no matter what, and it's better to do it in the above loop than
105        // by manually writing since it can be vectorized.
106        // Then, `value` was never semantically moved into the vec, only copied from,
107        // so we forget it here to prevent the Drop.
108        mem::forget(value);
109    }
110}
alloc/vec/spec_extend_with.rs

alloc/vec/
spec_extend_with.rs
