rten_tensor/
assume_init.rs

1use std::mem::MaybeUninit;
2
3/// Trait for converting collections of uninitialized (`MaybeUninit<T>`) values
4/// to collections of corresponding initializes values (`T`).
5///
6/// ## Example
7///
8/// ```
9/// use std::mem::MaybeUninit;
10/// use rten_tensor::AssumeInit;
11///
12/// fn scale_values<'a>(dst: &'a mut [MaybeUninit<f32>], src: &[f32], scale: f32) -> &'a mut [f32] {
13///   for (y, x) in dst.into_iter().zip(src) {
14///     y.write(x * scale);
15///   }
16///   // Safety: All elements have been initialized.
17///   unsafe { dst.assume_init() }
18/// }
19///
20/// let src = [1., 2., 3.];
21/// let mut dst = [MaybeUninit::uninit(); 3];
22/// let scaled = scale_values(&mut dst, &src, 2.);
23/// assert_eq!(scaled, [2., 4., 6.]);
24/// ```
25pub trait AssumeInit {
26    /// The type of the initialized storage.
27    type Output;
28
29    /// Cast `self` to a collection of initialized values.
30    ///
31    /// # Safety
32    ///
33    /// The caller must guarantee that all elements have been initialized.
34    unsafe fn assume_init(self) -> Self::Output;
35}
36
37impl<T> AssumeInit for Vec<MaybeUninit<T>> {
38    type Output = Vec<T>;
39
40    unsafe fn assume_init(mut self) -> Self::Output {
41        let (ptr, len, capacity) = (self.as_mut_ptr(), self.len(), self.capacity());
42
43        // Don't drop self, as that would deallocate.
44        std::mem::forget(self);
45
46        // Safety: We're re-constructing a `Vec` with the same length and
47        // capacity and an element type that has the same size and alignment,
48        // just cast from uninitialized to initialized.
49        unsafe { Vec::from_raw_parts(ptr as *mut T, len, capacity) }
50    }
51}
52
53impl<'a, T> AssumeInit for &'a [MaybeUninit<T>] {
54    type Output = &'a [T];
55
56    unsafe fn assume_init(self) -> Self::Output {
57        unsafe { std::mem::transmute(self) }
58    }
59}
60
61impl<'a, T> AssumeInit for &'a mut [MaybeUninit<T>] {
62    type Output = &'a mut [T];
63
64    unsafe fn assume_init(self) -> Self::Output {
65        unsafe { std::mem::transmute(self) }
66    }
67}
68
69#[cfg(test)]
70mod tests {
71    use std::mem::MaybeUninit;
72
73    use super::AssumeInit;
74
75    #[test]
76    fn test_assume_init_vec() {
77        let mut vec = vec![MaybeUninit::uninit(); 3];
78        vec.reserve(4);
79
80        for x in &mut vec {
81            x.write(2.);
82        }
83
84        let vec = unsafe { vec.assume_init() };
85        assert_eq!(vec.len(), 3);
86        assert_eq!(vec.capacity(), 7);
87        assert_eq!(vec, &[2., 2., 2.]);
88    }
89}
rten_tensor/assume_init.rs

rten_tensor/
assume_init.rs
