rten_vecmath/
extend_init.rs1use std::mem::MaybeUninit;
2
3pub trait ExtendInit {
9 type Elem;
11
12 fn extend_init<F: Fn(&mut [MaybeUninit<Self::Elem>]) -> &[Self::Elem]>(&mut self, f: F);
24}
25
26impl<T> ExtendInit for Vec<T> {
27 type Elem = T;
28
29 fn extend_init<F: Fn(&mut [MaybeUninit<Self::Elem>]) -> &[Self::Elem]>(&mut self, f: F) {
30 let cap = self.spare_capacity_mut();
31 let cap_ptr = cap.as_ptr();
32 let cap_len = cap.len();
33
34 let initialized = f(cap);
35 assert_eq!(
36 initialized.as_ptr(),
37 cap_ptr as *const T,
38 "returned slice must be a prefix of the input"
39 );
40 assert!(
41 initialized.len() <= cap_len,
42 "initialized slice length {} is longer than input {}",
43 initialized.len(),
44 cap_len
45 );
46 let n_init = initialized.len();
47
48 unsafe { self.set_len(self.len() + n_init) }
50 }
51}
52
53#[cfg(test)]
54mod tests {
55 use std::mem::MaybeUninit;
56
57 use super::ExtendInit;
58
59 fn fill<T: Copy>(xs: &mut [MaybeUninit<T>], value: T) -> &mut [T] {
61 for x in xs.iter_mut() {
62 x.write(value);
63 }
64 unsafe { std::mem::transmute::<&mut [MaybeUninit<T>], &mut [T]>(xs) }
65 }
66
67 #[test]
68 fn test_extend_init() {
69 let mut vec = Vec::with_capacity(7);
70
71 vec.extend_init(|uninit| {
72 assert_eq!(uninit.len(), 7);
73 fill(&mut uninit[..3], 1.)
74 });
75 assert_eq!(vec.len(), 3);
76 assert_eq!(vec, &[1., 1., 1.]);
77
78 vec.extend_init(|uninit| {
79 assert_eq!(uninit.len(), 4);
80 fill(uninit, 2.)
81 });
82 assert_eq!(vec.len(), 7);
83 assert_eq!(vec, &[1., 1., 1., 2., 2., 2., 2.]);
84 }
85}