silly_alloc/bucket/
mod.rs

1/*!
2Bucket allocators.
3
4A bucket allocator defines multiple buckets where each item has the same size. The number and granularity of the buckets can be tuned to what is typical allocation behavior of the app at hand. In contrast to bump allocators, bucket allocators can also free memory.
5
6# Examples
7
8You can create a custom bucket allocator by writing a (pseudo) struct that contains the parameters for all buckets:
9
10```rust
11use silly_alloc::bucket_allocator;
12
13#[bucket_allocator]
14struct MyBucketAllocator {
15    vec2: Bucket<SlotSize<2>, NumSlots<128>, Align<2>>,
16    overflow: Bucket<SlotSize<64>, NumSlots<64>, Align<64>>
17}
18```
19
20Note that these types and generics are not really in use. They are merely there for an idiomatic and syntactically plausible way to provide the parameters. The macro rewrites this struct definition to another one using different types.
21
22The new bucket allocator can then be instantiated and used as a global allocator as per usual:
23
24```rust
25# use silly_alloc::bucket_allocator;
26#
27# #[bucket_allocator]
28# struct MyBucketAllocator {
29#     vec2: Bucket<SlotSize<2>, NumSlots<128>, Align<2>>,
30#     overflow: Bucket<SlotSize<64>, NumSlots<64>, Align<64>>
31# }
32#[global_allocator]
33static ALLOCATOR: MyBucketAllocator = MyBucketAllocator::new();
34```
35
36Buckets are checked for the best fit in order of specification. Full buckets are skipped.
37
38# Technical details
39
40A bucket is defined by three parameters:
41
42- The size of an item in the bucket
43- The number of items that fit in the bucket
44- An optional alignment constraint
45
46The speed of bucket allocators stems from the fact that all items in the bucket are the same size, and as such a simple bit mask is enough to track if a slot is in use or not. For simplicity, 32 slots are grouped into one segment, where a single `u32` is used to hold that bitmask. A bucket, as a consequence, is an array of segments. This also implies that `NumSlots<N>` will be rounded up to the next multiple of 32.
47*/
48
49use core::{
50    fmt::{Debug, Formatter},
51    marker::PhantomData,
52    mem::{size_of, MaybeUninit},
53};
54
55use bytemuck::Zeroable;
56
57// TODO: Implement thread-safe segments
58// #[cfg(target_feature = "feature")]
59#[derive(Clone, Copy)]
60struct SegmentHeader([u32; NUM_U32_PER_HEADER]);
61impl SegmentHeader {
62    const fn new() -> Self {
63        SegmentHeader([0u32; NUM_U32_PER_HEADER])
64    }
65
66    fn first_free_slot_idx(&self) -> Option<usize> {
67        for header in self.0 {
68            let clo = header.leading_ones();
69            if clo == 32 {
70                continue;
71            }
72            return Some(usize::try_from(clo).unwrap());
73        }
74        None
75    }
76
77    fn slot_to_idx(slot_idx: usize) -> (usize, usize) {
78        (slot_idx >> 5, slot_idx % 32)
79    }
80
81    fn set_slot(&mut self, slot_idx: usize) {
82        let (arr_idx, bit_idx) = Self::slot_to_idx(slot_idx);
83        self.0[arr_idx] |= 1 << (31 - bit_idx);
84    }
85
86    fn unset_slot(&mut self, slot_idx: usize) {
87        let (arr_idx, bit_idx) = Self::slot_to_idx(slot_idx);
88        self.0[arr_idx] &= !(1 << (31 - bit_idx));
89    }
90}
91
92unsafe impl Zeroable for SegmentHeader {}
93
94impl Debug for SegmentHeader {
95    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
96        f.write_str("SegmentHeader {")?;
97        for header in self.0 {
98            f.write_fmt(format_args!("{:08b}", (header >> 24) & 0xff))?;
99            f.write_str("_")?;
100            f.write_fmt(format_args!("{:08b}", (header >> 16) & 0xff))?;
101            f.write_str("_")?;
102            f.write_fmt(format_args!("{:08b}", (header >> 8) & 0xff))?;
103            f.write_str("_")?;
104            f.write_fmt(format_args!("{:08b}", (header >> 0) & 0xff))?;
105        }
106        f.write_str("}")?;
107        Ok(())
108    }
109}
110
111const NUM_U32_PER_HEADER: usize = 1;
112pub const NUM_SLOTS_PER_SEGMENT: usize = NUM_U32_PER_HEADER * size_of::<u32>() * 8;
113pub const SEGMENT_HEADER_SIZE: usize = NUM_U32_PER_HEADER * size_of::<u32>();
114
115pub trait Slot: Copy + Default {
116    fn get(&self) -> *const u8;
117    fn size() -> usize;
118}
119
120#[derive(Clone, Copy)]
121pub struct Segment<S: Slot> {
122    header: SegmentHeader,
123    slots: [S; NUM_SLOTS_PER_SEGMENT],
124}
125
126impl<S: Slot> Segment<S> {
127    fn new() -> Self {
128        Segment {
129            header: SegmentHeader::new(),
130            slots: [S::default(); NUM_SLOTS_PER_SEGMENT],
131        }
132    }
133
134    fn get_slot(&self, slot_idx: usize) -> *const u8 {
135        self.slots[slot_idx].get()
136    }
137}
138
139impl<S: Slot> Default for Segment<S> {
140    fn default() -> Self {
141        Self::new()
142    }
143}
144
145unsafe impl<S: Slot + Zeroable> Zeroable for Segment<S> {}
146
147impl<S: Slot> Debug for Segment<S> {
148    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
149        f.debug_struct("Segment")
150            .field("header", &self.header)
151            .finish()?;
152        Ok(())
153    }
154}
155
156macro_rules! align_type {
157    ($name:ident, $n:expr) => {
158        #[derive(Debug, Clone, Copy)]
159        #[repr(C, align($n))]
160        pub struct $name<const N: usize>([u8; N]);
161        impl<const N: usize> $name<N> {
162            pub const fn new() -> Self {
163                $name([0u8; N])
164            }
165        }
166
167        impl<const N: usize> Default for $name<N> {
168            fn default() -> Self {
169                Self::new()
170            }
171        }
172
173        impl<const N: usize> Slot for $name<N> {
174            fn get(&self) -> *const u8 {
175                self.0.as_ptr()
176            }
177
178            fn size() -> usize {
179                N
180            }
181        }
182
183        unsafe impl<const N: usize> Zeroable for $name<N> {}
184    };
185}
186
187align_type!(SlotWithAlign1, 1);
188align_type!(SlotWithAlign2, 2);
189align_type!(SlotWithAlign4, 4);
190align_type!(SlotWithAlign8, 8);
191align_type!(SlotWithAlign16, 16);
192align_type!(SlotWithAlign32, 32);
193align_type!(SlotWithAlign64, 64);
194align_type!(SlotWithAlign128, 128);
195align_type!(SlotWithAlign256, 256);
196align_type!(SlotWithAlign512, 512);
197
198#[derive(Debug, Clone, Copy)]
199pub struct BucketImpl<S: Slot, const N: usize> {
200    // Deeply saddened that I had to introduce a bool flag here. There just currently is no way initializing this array in a const way, so I have to defer it to runtime and track when it has been done. This flag can be removed when at least one of these options is available in stable rust:
201    // - const fn can be in traits
202    // - core::ptr::write_bytes is const
203    // - core::mem::MaybeUninit.zeroed() is const (although that probably relies on previous point)
204    is_init: bool,
205    segments: MaybeUninit<[Segment<S>; N]>,
206}
207
208impl<S: Slot, const NUM_SEGMENTS: usize> BucketImpl<S, NUM_SEGMENTS> {
209    pub const fn new() -> Self {
210        Self {
211            is_init: false,
212            segments: MaybeUninit::uninit(),
213        }
214        // unsafe { MaybeUninit::zeroed().assume_init() }
215    }
216
217    pub fn ensure_init(&mut self) {
218        if self.is_init {
219            return;
220        }
221        unsafe {
222            core::ptr::write_bytes(self.segments.as_mut_ptr(), 0u8, 1);
223        }
224        self.is_init = true
225    }
226
227    fn get_segments(&self) -> &[Segment<S>; NUM_SEGMENTS] {
228        assert!(self.is_init);
229        unsafe { self.segments.assume_init_ref() }
230    }
231
232    fn get_segments_mut(&mut self) -> &mut [Segment<S>; NUM_SEGMENTS] {
233        assert!(self.is_init);
234        unsafe { self.segments.assume_init_mut() }
235    }
236
237    pub fn claim_first_available_slot(&mut self) -> Option<*const u8> {
238        for seg in self.get_segments_mut().iter_mut() {
239            let Some(slot_idx) = seg.header.first_free_slot_idx() else {continue};
240            seg.header.set_slot(slot_idx);
241            return Some(seg.get_slot(slot_idx));
242        }
243        None
244    }
245
246    fn global_to_local(&self, slot_idx: usize) -> (usize, usize) {
247        let seg_idx = slot_idx / NUM_SLOTS_PER_SEGMENT;
248        let slot_idx = slot_idx % NUM_SLOTS_PER_SEGMENT;
249        (seg_idx, slot_idx)
250    }
251
252    pub fn get_slot(&self, slot_idx: usize) -> *const u8 {
253        let (seg_idx, slot_idx) = self.global_to_local(slot_idx);
254        self.get_segments()[seg_idx].get_slot(slot_idx)
255    }
256
257    pub fn set_slot(&mut self, slot_idx: usize) {
258        let (seg_idx, slot_idx) = self.global_to_local(slot_idx);
259        self.get_segments_mut()[seg_idx].header.set_slot(slot_idx);
260    }
261
262    pub fn unset_slot(&mut self, slot_idx: usize) {
263        let (seg_idx, slot_idx) = self.global_to_local(slot_idx);
264        self.get_segments_mut()[seg_idx].header.unset_slot(slot_idx);
265    }
266
267    pub fn slot_idx_for_ptr(&self, ptr: *const u8) -> Option<usize> {
268        let seg_stride = size_of::<Segment<S>>();
269        let slot_stride = size_of::<S>();
270
271        let start = self.segments.as_ptr() as *const u8;
272        let offset = unsafe { ptr.offset_from(start) };
273        // Conversion to usize will only succeed for positive numbers.
274        // If it's negative, ptr is in previous segment.
275        let offset = usize::try_from(offset).ok()?;
276        let seg_idx = offset / seg_stride;
277        let offset = offset % seg_stride;
278        let slot_idx = offset / slot_stride;
279        if seg_idx >= NUM_SEGMENTS || slot_idx > NUM_SLOTS_PER_SEGMENT {
280            return None;
281        }
282
283        Some(seg_idx * NUM_SLOTS_PER_SEGMENT + slot_idx)
284    }
285}
286
287unsafe impl<S: Slot, const N: usize> Zeroable for BucketImpl<S, N> {}
288
289impl<S: Slot, const N: usize> Default for BucketImpl<S, N> {
290    fn default() -> Self {
291        BucketImpl::<S, N>::new()
292    }
293}
294
295pub struct SlotSize<const N: usize>;
296pub struct NumSlots<const N: usize>;
297pub struct Align<const N: usize>;
298
299pub struct Bucket<S, N, A = Align<1>>(PhantomData<S>, PhantomData<N>, PhantomData<A>);
300
301#[cfg(test)]
302mod test {
303    use super::*;
304
305    use core::alloc::{GlobalAlloc, Layout};
306
307    use anyhow::Result;
308
309    use silly_alloc_macros::bucket_allocator;
310
311    #[bucket_allocator]
312    struct MyBucketAllocator {
313        vec2: Bucket<SlotSize<2>, NumSlots<32>, Align<2>>,
314        vec4: Bucket<SlotSize<4>, NumSlots<32>, Align<4>>,
315        vec8: Bucket<SlotSize<8>, NumSlots<32>, Align<8>>,
316    }
317
318    #[test]
319    fn next_in_bucket() -> Result<()> {
320        let b = MyBucketAllocator::new();
321        unsafe {
322            let ptr1 = b.alloc(Layout::from_size_align(2, 1)?);
323            let ptr2 = b.alloc(Layout::from_size_align(2, 1)?);
324            assert!(!ptr1.is_null());
325            assert!(!ptr2.is_null());
326            assert_eq!(ptr1.offset(2), ptr2);
327        }
328        Ok(())
329    }
330
331    #[test]
332    fn full_bucket() -> Result<()> {
333        #[bucket_allocator]
334        struct MyBucketAllocator {
335            vec2: Bucket<SlotSize<2>, NumSlots<32>, Align<2>>,
336        }
337        unsafe {
338            let b = MyBucketAllocator::new();
339            let l = Layout::from_size_align(2, 1)?;
340            for _ in 0..32 {
341                b.alloc(l);
342            }
343            let ptr = b.alloc(l);
344            assert!(ptr.is_null());
345        }
346        Ok(())
347    }
348
349    #[test]
350    fn reuse() -> Result<()> {
351        let b = MyBucketAllocator::new();
352        unsafe {
353            let layout = Layout::from_size_align(2, 1)?;
354            let ptr1 = b.alloc(layout.clone());
355            let ptr2 = b.alloc(layout.clone());
356            let ptr3 = b.alloc(layout.clone());
357            assert_eq!(ptr1.offset(2), ptr2);
358            assert_eq!(ptr2.offset(2), ptr3);
359            b.dealloc(ptr2, layout);
360            let ptr4 = b.alloc(layout.clone());
361            assert_eq!(ptr2, ptr4);
362        }
363        Ok(())
364    }
365
366    #[test]
367    fn bucket_overflow() -> Result<()> {
368        let b = MyBucketAllocator::new();
369        unsafe {
370            let layout = Layout::from_size_align(2, 1)?;
371            // Fill 2 byte bucket
372            for _ in 0..32 {
373                b.alloc(layout.clone());
374            }
375            let ptr1 = b.alloc(Layout::from_size_align(4, 1)?);
376            let ptr2 = b.alloc(layout.clone());
377            assert_eq!(ptr1.offset(4), ptr2);
378        }
379        Ok(())
380    }
381
382    #[test]
383    fn alignment() -> Result<()> {
384        let mut b = MyBucketAllocator::new();
385        unsafe {
386            let layout = Layout::from_size_align(2, 8)?;
387            let ptr1 = b.alloc(layout);
388            // Alignment requirement should force the allocation into the last bucket desipte its size
389            assert!(ptr1 >= &mut b.2 as *mut _ as *mut u8);
390        }
391        Ok(())
392    }
393
394    #[test]
395    fn alignment_fail() -> Result<()> {
396        let b = MyBucketAllocator::new();
397        unsafe {
398            let layout = Layout::from_size_align(2, 32)?;
399            let ptr1 = b.alloc(layout);
400            assert!(ptr1.is_null());
401        }
402        Ok(())
403    }
404
405    #[test]
406    fn first_alloc_in_late_bucket() -> Result<()> {
407        unsafe {
408            // Rust doesn't guarantee field order in structs, so I don't know whether the 2byte bucket or the 8byte bucket comes first. So I am gonna test both, as they both have to work anyway.
409            let layouts = vec![
410                Layout::from_size_align(2, 1)?,
411                Layout::from_size_align(8, 1)?,
412            ];
413            for layout in layouts {
414                let b = MyBucketAllocator::new();
415                let _ptr1 = b.alloc(layout.clone());
416                let ptr2 = b.alloc(layout.clone());
417                let _ptr3 = b.alloc(layout.clone());
418                b.dealloc(ptr2, layout);
419                let ptr4 = b.alloc(layout.clone());
420                assert_eq!(ptr2, ptr4);
421            }
422        }
423        Ok(())
424    }
425
426    #[test]
427    fn small_buckets() -> Result<()> {
428        #[bucket_allocator]
429        struct MyBucketAllocator {
430            vec2: Bucket<SlotSize<2>, NumSlots<1>, Align<2>>,
431        }
432        unsafe {
433            let b = MyBucketAllocator::new();
434            let l = Layout::from_size_align(2, 2)?;
435            let ptr1 = b.alloc(l);
436            assert!(!ptr1.is_null());
437        }
438        Ok(())
439    }
440
441    #[test]
442    fn unsorted_buckets() -> Result<()> {
443        #[bucket_allocator(sort_buckets = true)]
444        struct MyBucketAllocator {
445            vec8: Bucket<SlotSize<8>, NumSlots<32>, Align<8>>,
446            vec2: Bucket<SlotSize<2>, NumSlots<32>, Align<8>>,
447        }
448
449        unsafe {
450            let b = MyBucketAllocator::new();
451            let ptr1 = b.alloc(Layout::from_size_align(2, 2)?);
452            let ptr2 = b.alloc(Layout::from_size_align(8, 8)?);
453            // Assert the two allocations are in different buckets
454            assert!(ptr1.offset_from(ptr2).abs() > 8);
455        }
456        Ok(())
457    }
458}
silly_alloc/bucket/mod.rs

silly_alloc/bucket/
mod.rs
