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        unsafe {
40            ptr::write_bytes(v.as_mut_ptr(), elem as u8, n);
41            v.set_len(n);
42        }
43        v
44    }
45}
46
47impl SpecFromElem for u8 {
48    #[inline]
49    fn from_elem<A: Allocator>(elem: u8, n: usize, alloc: A) -> Vec<u8, A> {
50        if elem == 0 {
51            return Vec { buf: RawVec::with_capacity_zeroed_in(n, alloc), len: n };
52        }
53        let mut v = Vec::with_capacity_in(n, alloc);
54        unsafe {
55            ptr::write_bytes(v.as_mut_ptr(), elem, n);
56            v.set_len(n);
57        }
58        v
59    }
60}
61
62// A better way would be to implement this for all ZSTs which are `Copy` and have trivial `Clone`
63// but the latter cannot be detected currently
64impl SpecFromElem for () {
65    #[inline]
66    fn from_elem<A: Allocator>(_elem: (), n: usize, alloc: A) -> Vec<(), A> {
67        let mut v = Vec::with_capacity_in(n, alloc);
68        // SAFETY: the capacity has just been set to `n`
69        // and `()` is a ZST with trivial `Clone` implementation
70        unsafe {
71            v.set_len(n);
72        }
73        v
74    }
75}