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