Skip to main content

alloc/vec/
spec_from_elem.rs

1use core::ptr;
2
3use super::{IsZero, Vec};
4use crate::alloc::Allocator;
5use crate::raw_vec::RawVec;
6
7// Specialization trait used for Vec::from_elem
8pub(super) trait SpecFromElem: Sized {
9    fn from_elem<A: Allocator>(elem: Self, n: usize, alloc: A) -> Vec<Self, A>;
10}
11
12impl<T: Clone> SpecFromElem for T {
13    default fn from_elem<A: Allocator>(elem: Self, n: usize, alloc: A) -> Vec<Self, A> {
14        let mut v = Vec::with_capacity_in(n, alloc);
15        v.extend_with(n, elem);
16        v
17    }
18}
19
20impl<T: Clone + IsZero> SpecFromElem for T {
21    #[inline]
22    default fn from_elem<A: Allocator>(elem: T, n: usize, alloc: A) -> Vec<T, A> {
23        if elem.is_zero() {
24            return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
25        }
26        let mut v = Vec::with_capacity_in(n, alloc);
27        v.extend_with(n, elem);
28        v
29    }
30}
31
32impl SpecFromElem for i8 {
33    #[inline]
34    fn from_elem<A: Allocator>(elem: i8, n: usize, alloc: A) -> Vec<i8, A> {
35        if elem == 0 {
36            return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
37        }
38        let mut v = Vec::with_capacity_in(n, alloc);
39        // ignore-tidy-undocumented-unsafe
40        unsafe {
41            ptr::write_bytes(v.as_mut_ptr(), elem as u8, n);
42            v.set_len(n);
43        }
44        v
45    }
46}
47
48impl SpecFromElem for u8 {
49    #[inline]
50    fn from_elem<A: Allocator>(elem: u8, n: usize, alloc: A) -> Vec<u8, A> {
51        if elem == 0 {
52            return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
53        }
54        let mut v = Vec::with_capacity_in(n, alloc);
55        // ignore-tidy-undocumented-unsafe
56        unsafe {
57            ptr::write_bytes(v.as_mut_ptr(), elem, n);
58            v.set_len(n);
59        }
60        v
61    }
62}
63
64// A better way would be to implement this for all ZSTs which are `Copy` and have trivial `Clone`
65// but the latter cannot be detected currently
66impl SpecFromElem for () {
67    #[inline]
68    fn from_elem<A: Allocator>(_elem: (), n: usize, alloc: A) -> Vec<(), A> {
69        let mut v = Vec::with_capacity_in(n, alloc);
70        // SAFETY: the capacity has just been set to `n`
71        // and `()` is a ZST with trivial `Clone` implementation
72        unsafe {
73            v.set_len(n);
74        }
75        v
76    }
77}