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
103

use {
    crate::{Memory, MultipleStorage, PinningStorage, SharedMutabilityStorage, Storage},
    core::{
        alloc::{AllocError, Allocator, Layout},
        mem::MaybeUninit,
        ptr::NonNull,
    },
};

/// A storage that stores objects via an [`Allocator`].
pub struct AllocStorage<A: Allocator> {
    alloc: A,
}

impl<A: Allocator> AllocStorage<A> {
    pub const fn new(alloc: A) -> Self {
        Self { alloc }
    }
}

unsafe impl<A: Allocator> Storage for AllocStorage<A> {
    type Handle = AllocHandle;

    fn allocate(&mut self, layout: Layout) -> Result<Self::Handle, AllocError> {
        let (ptr, _meta) = self.alloc.allocate(layout)?.to_raw_parts();
        Ok(AllocHandle::new(ptr))
    }

    unsafe fn deallocate(&mut self, handle: Self::Handle, layout: Layout) {
        self.alloc.deallocate(handle.pointer.cast(), layout)
    }

    unsafe fn resolve(&self, handle: Self::Handle, layout: Layout) -> &Memory {
        NonNull::from_raw_parts(handle.pointer, layout.size()).as_ref()
    }

    unsafe fn resolve_mut(&mut self, handle: Self::Handle, layout: Layout) -> &mut Memory {
        NonNull::from_raw_parts(handle.pointer, layout.size()).as_mut()
    }

    unsafe fn grow(
        &mut self,
        handle: Self::Handle,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<Self::Handle, AllocError> {
        let (ptr, _meta) = self
            .alloc
            .grow(handle.pointer.cast(), old_layout, new_layout)?
            .to_raw_parts();
        Ok(AllocHandle::new(ptr))
    }

    unsafe fn shrink(
        &mut self,
        handle: Self::Handle,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<Self::Handle, AllocError> {
        let (ptr, _meta) = self
            .alloc
            .shrink(handle.pointer.cast(), old_layout, new_layout)?
            .to_raw_parts();
        Ok(AllocHandle::new(ptr))
    }
}

unsafe impl<A: Allocator> MultipleStorage for AllocStorage<A> {
    unsafe fn resolve_many_mut<const N: usize>(
        &mut self,
        handles: [(Self::Handle, Layout); N],
    ) -> [&mut Memory; N] {
        let mut ptrs: [MaybeUninit<&mut Memory>; N] = MaybeUninit::uninit().assume_init();
        for (ptr, (handle, layout)) in ptrs.iter_mut().zip(handles) {
            ptr.write(self.resolve_raw(handle, layout));
        }
        MaybeUninit::array_assume_init(ptrs)
    }
}

unsafe impl<A: Allocator> SharedMutabilityStorage for AllocStorage<A> {
    unsafe fn resolve_raw(&self, handle: Self::Handle, layout: Layout) -> &mut Memory {
        NonNull::from_raw_parts(handle.pointer, layout.size()).as_mut()
    }
}

unsafe impl<A: Allocator> PinningStorage for AllocStorage<A> {}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct AllocHandle {
    pointer: NonNull<()>,
}

unsafe impl Send for AllocHandle {}
unsafe impl Sync for AllocHandle {}

impl AllocHandle {
    const fn new<T: ?Sized>(pointer: NonNull<T>) -> Self {
        Self {
            pointer: pointer.cast(),
        }
    }
}