rten_vecmath/
extend_init.rs

1use std::mem::MaybeUninit;
2
3/// Extend a buffer by incrementally initializing spare capacity.
4///
5/// This is implemented for [`Vec<T>`], where it provides a safe API to
6/// initialize the spare capacity returned by
7/// [`spare_capacity_mut`](Vec::spare_capacity_mut).
8pub trait ExtendInit {
9    /// Element type in the buffer.
10    type Elem;
11
12    /// Extend the buffer by initializing a portion of the buffer's spare
13    /// capacity.
14    ///
15    /// The function `f` is passed the uninitialized portion of the buffer and
16    /// should return the portion that it has initialized. `extend_init` can
17    /// be called many times, until the entire buffer has been initialized.
18    ///
19    /// # Panics
20    ///
21    /// Panics if `f` returns a slice that is not a prefix of the slice that
22    /// was passed to it.
23    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        // Safety: `n_init` elements from the spare capacity have been initialized.
49        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    // Implementation of `MaybeUninit::fill` from nightly Rust.
60    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}