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

//! Smart pointer guard.

use alloc::alloc;
use core::{
    alloc::{AllocError, Layout},
    mem,
    ptr::{slice_from_raw_parts, slice_from_raw_parts_mut, NonNull},
};

/// Smart pointer guard.
pub struct AllocGuard {
    ptr: NonNull<u8>,
    layout: Layout,
}

impl AllocGuard {
    /// Create a new smart pointer.
    pub fn new(layout: Layout) -> Result<Self, AllocError> {
        unsafe {
            NonNull::new(alloc::alloc(layout))
                .map_or(Err(AllocError), |ptr| Ok(Self::from_ptr(ptr, layout)))
        }
    }

    /// Get the pointer inside.
    pub fn ptr(&self) -> NonNull<u8> {
        self.ptr.clone()
    }

    /// Extract the slice.
    pub fn as_slice(&self) -> &[u8] {
        unsafe { &*slice_from_raw_parts(self.ptr.as_ptr() as *const _, self.layout.size()) }
    }

    /// Extract the mutable slice.
    pub fn as_mut_slice(&self) -> &mut [u8] {
        unsafe { &mut *slice_from_raw_parts_mut(self.ptr.as_ptr(), self.layout.size()) }
    }

    /// Drop and don't deallocate the pointer.
    ///
    /// Often called when the ownership of the pointer is transfered.
    pub fn consume(self) {
        mem::forget(self)
    }

    /// Create a new smart pointer from raw pointer.
    ///
    /// SAFETY: You must guarantee that
    /// - `ptr` and `layout `are valid.
    /// - [consume] is called before anyone uses `ptr`.
    pub unsafe fn from_ptr(ptr: NonNull<u8>, layout: Layout) -> Self {
        Self { ptr, layout }
    }
}

impl Drop for AllocGuard {
    fn drop(&mut self) {
        unsafe { alloc::dealloc(self.ptr.as_ptr(), self.layout) }
    }
}

unsafe impl Sync for AllocGuard {}
unsafe impl Send for AllocGuard {}