heap_arr 0.4.0

Allocate fixed-size arrays directly on the heap, bypassing stack allocation limits
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

use core::mem::{self, MaybeUninit};

extern crate alloc;

pub struct Guard<'a, T, const N: usize> {
    array: &'a mut [MaybeUninit<T>; N],
    initialized: usize,
}

impl<'a, T, const N: usize> Guard<'a, T, N> {
    pub fn new(array: &'a mut [MaybeUninit<T>; N]) -> Self {
        Self {
            array,
            initialized: 0,
        }
    }

    /// # Safety
    /// The caller must ensure not to call this method more than `N` times
    pub unsafe fn push_unchecked(&mut self, value: T) {
        debug_assert!(self.initialized < N);

        unsafe { self.array.get_unchecked_mut(self.initialized).write(value) };
        self.initialized += 1;
    }

    /// # Safety
    /// The caller must ensure that `self.initialized == N`, so that
    /// the [`Self::push_unchecked`] method has been called exactly `N` times
    pub unsafe fn finalize(self) {
        debug_assert_eq!(self.initialized, N);

        mem::forget(self);
    }
}

impl<T, const N: usize> Drop for Guard<'_, T, N> {
    fn drop(&mut self) {
        // Safety: We are dropping all initialized entries
        for i in 0..self.initialized {
            unsafe { self.array.get_unchecked_mut(i).assume_init_drop() };
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Guard;
    use core::sync::atomic::{AtomicUsize, Ordering};

    #[test]
    fn init_tracking() {
        let mut array = crate::uninit::<u32, 5>();
        let mut guard = Guard::new(&mut array);

        assert_eq!(0, guard.initialized);
        unsafe { guard.push_unchecked(10) };
        assert_eq!(1, guard.initialized);
    }

    #[test]
    fn full_init() {
        let mut array = crate::uninit::<u32, 3>();
        let mut guard = Guard::new(&mut array);

        unsafe {
            guard.push_unchecked(1);
            guard.push_unchecked(2);
            guard.push_unchecked(3);
            guard.finalize();
        }

        assert_eq!(*unsafe { crate::assume_init(array) }, [1, 2, 3]);
    }

    #[test]
    fn partial_drop_cleanup() {
        struct DropCounter {
            counter: &'static AtomicUsize,
        }

        impl Drop for DropCounter {
            fn drop(&mut self) {
                self.counter.fetch_add(1, Ordering::SeqCst);
            }
        }

        static DROPS: AtomicUsize = AtomicUsize::new(0);
        {
            let mut array = crate::uninit::<DropCounter, 5>();
            let mut guard = Guard::new(&mut array);

            unsafe {
                guard.push_unchecked(DropCounter { counter: &DROPS });
                guard.push_unchecked(DropCounter { counter: &DROPS });
            }
            // Guard drops here, should clean up 2 elements
        }

        assert_eq!(DROPS.load(Ordering::SeqCst), 2);
    }
}