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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135

use core::cell::UnsafeCell;
use core::ptr::NonNull;

use alloc::boxed::Box;
use alloc::vec::Vec;

use musli::{Allocator, Buf};

/// Buffer used in combination with an [`Allocator`].
pub struct System {
    internal: UnsafeCell<Internal>,
}

impl System {
    /// Construct a new allocator.
    #[inline]
    pub const fn new() -> Self {
        Self {
            internal: UnsafeCell::new(Internal { head: None }),
        }
    }
}

impl Default for System {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}

impl Allocator for System {
    type Buf<'this> = SystemBuf<'this> where Self: 'this;

    #[inline(always)]
    fn alloc(&self) -> Option<Self::Buf<'_>> {
        Some(SystemBuf {
            region: Internal::alloc(&self.internal),
            internal: &self.internal,
        })
    }
}

impl Drop for System {
    fn drop(&mut self) {
        let internal = unsafe { &mut *self.internal.get() };

        while let Some(mut head) = internal.head.take() {
            // SAFETY: This collection has exclusive access to any heads it
            // contain.
            unsafe {
                internal.head = head.as_mut().next.take();
                drop(Box::from_raw(head.as_ptr()));
            }
        }
    }
}

/// A vector-backed allocation.
pub struct SystemBuf<'a> {
    region: &'a mut Region,
    internal: &'a UnsafeCell<Internal>,
}

impl<'a> Buf for SystemBuf<'a> {
    #[inline]
    fn write(&mut self, bytes: &[u8]) -> bool {
        self.region.data.extend_from_slice(bytes);
        true
    }

    #[inline(always)]
    fn len(&self) -> usize {
        self.region.data.len()
    }

    #[inline(always)]
    fn as_slice(&self) -> &[u8] {
        &self.region.data
    }
}

impl<'a> Drop for SystemBuf<'a> {
    fn drop(&mut self) {
        Internal::free(self.internal, self.region);
    }
}

/// An allocated region.
#[repr(C)]
struct Region {
    data: Vec<u8>,
    // Pointer to the next free region.
    next: Option<NonNull<Region>>,
}

/// Internals of the allocator.
struct Internal {
    // Regions of re-usable allocations we can hand out.
    head: Option<NonNull<Region>>,
}

impl Internal {
    /// Allocate a new region.
    ///
    /// Note that this will return a leaked memory region, so the unbound
    /// lifetime is intentional.
    fn alloc<'a>(this: &UnsafeCell<Self>) -> &'a mut Region {
        // SAFETY: We take care to only access internals in a single-threaded
        // mutable fashion.
        let internal = unsafe { &mut *this.get() };

        if let Some(mut head) = internal.head.take() {
            // SAFETY: This collection has exclusive access to any heads it contain.
            unsafe {
                let head = head.as_mut();
                internal.head = head.next.take();
                head
            }
        } else {
            Box::leak(Box::new(Region {
                data: Vec::new(),
                next: None,
            }))
        }
    }

    fn free(this: &UnsafeCell<Self>, region: &mut Region) {
        unsafe {
            let this = &mut *this.get();
            region.data.clear();
            region.next = this.head;
            this.head = Some(NonNull::from(region));
        }
    }
}