From 5000e3d9747d998029586611b5bfdeebb45b5c5a Mon Sep 17 00:00:00 2001 From: Philip Woolford Date: Mon, 21 Oct 2024 14:13:25 +1030 Subject: [PATCH] Implementation of a segmented heap This widens the implementation of the `Heap` struct into a `SegmentedHeap`, allowing multiple regions to be used for heap allocation. There is not, however, any fitting strategy - all allocations using allocate_first_fit will always use the first memory region, unless it fails. Allocations will cascade down the array of hole lists, failing to allocate when all regions fail to allocate. This also makes the assumption that memory regions are non-overlapping, so construction from a list of pointer+sizes is unsafe and that invariant is left to the consumer to enforce. --- src/hole.rs | 18 ++-- src/lib.rs | 230 ++++++++++++++++++++++++++++++++++++---------------- src/test.rs | 60 +++++++------- 3 files changed, 201 insertions(+), 107 deletions(-) diff --git a/src/hole.rs b/src/hole.rs index baf9d96..669bb93 100644 --- a/src/hole.rs +++ b/src/hole.rs @@ -388,17 +388,21 @@ impl HoleList { // NOTE: We could probably replace this with an `Option` instead of a `Result` in a later // release to remove this clippy warning #[allow(clippy::result_unit_err)] - pub fn allocate_first_fit(&mut self, layout: Layout) -> Result<(NonNull, Layout), ()> { - let aligned_layout = Self::align_layout(layout).map_err(|_| ())?; - let mut cursor = self.cursor().ok_or(())?; + pub fn allocate_first_fit(&mut self, layout: Layout) -> Option<(NonNull, Layout)> { + let aligned_layout = Self::align_layout(layout).ok()?; + let mut cursor = self.cursor()?; loop { match cursor.split_current(aligned_layout) { Ok((ptr, _len)) => { - return Ok((NonNull::new(ptr).ok_or(())?, aligned_layout)); + return Some(( + // SAFETY: This can not be null as it is derived from a NonNull pointer in `split_current` + unsafe { NonNull::new_unchecked(ptr) } + , aligned_layout + )); } Err(curs) => { - cursor = curs.next().ok_or(())?; + cursor = curs.next()?; } } } @@ -419,7 +423,7 @@ impl HoleList { /// The function performs exactly the same layout adjustments as [`allocate_first_fit`] and /// returns the aligned layout. pub unsafe fn deallocate(&mut self, ptr: NonNull, layout: Layout) -> Layout { - let aligned_layout = Self::align_layout(layout).unwrap(); + let aligned_layout = Self::align_layout(layout).expect("This should never error, as the validity was checked during allocation."); deallocate(self, ptr.as_ptr(), aligned_layout.size()); aligned_layout } @@ -689,7 +693,7 @@ pub mod test { #[test] fn cursor() { let mut heap = new_heap(); - let curs = heap.holes.cursor().unwrap(); + let curs = heap.holes[0].cursor().unwrap(); // This is the "dummy" node assert_eq!(curs.previous().size, 0); // This is the "full" heap diff --git a/src/lib.rs b/src/lib.rs index b76e7d1..565bcba 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,12 +28,18 @@ pub mod hole; mod test; /// A fixed size heap backed by a linked list of free memory blocks. -pub struct Heap { +/*pub struct Heap { used: usize, holes: HoleList, +} */ +pub struct SegmentedHeap { + used: usize, + holes: [HoleList; SEGMENTS] } -#[cfg(fuzzing)] +pub type Heap = SegmentedHeap<1>; + +#[cfg(any(test, fuzzing))] impl Heap { pub fn debug(&mut self) { println!( @@ -41,23 +47,161 @@ impl Heap { self.bottom(), self.top(), self.size(), - self.holes.first.size, + self.holes[0].first.size, ); - self.holes.debug(); + self.holes[0].debug(); } } unsafe impl Send for Heap {} -impl Heap { + +impl SegmentedHeap { /// Creates an empty heap. All allocate calls will return `None`. - pub const fn empty() -> Heap { - Heap { + pub const fn empty() -> Self { + Self { used: 0, - holes: HoleList::empty(), + holes: unsafe { core::mem::zeroed() }, + } + } + + /// Allocates a chunk of the given size with the given alignment. Returns a pointer to the + /// beginning of that chunk if it was successful. Else it returns `None`. + /// This function scans the list of free memory blocks and uses the first block that is big + /// enough. The runtime is in O(n) where n is the number of free blocks, but it should be + /// reasonably fast for small allocations. + pub fn allocate_first_fit(&mut self, layout: Layout) -> Option> { + let allocation = self.holes.iter_mut() + // Applies the allocation function to each hole list and returns the first successful allocation + .find_map(|list| list.allocate_first_fit(layout)); + + if let Some((allocation, layout)) = allocation { + self.used += layout.size(); + Some(allocation) + } else { + None + } + } + + /// Frees the given allocation. `ptr` must be a pointer returned + /// by a call to the `allocate_first_fit` function with identical size and alignment. + /// + /// This function walks the list of free memory blocks and inserts the freed block at the + /// correct place. If the freed block is adjacent to another free block, the blocks are merged + /// again. This operation is in `O(n)` since the list needs to be sorted by address. + /// + /// # Safety + /// + /// `ptr` must be a pointer returned by a call to the [`allocate_first_fit`] function with + /// identical layout. Undefined behavior may occur for invalid arguments. Panics when none of the hole lists contain the passed pointer. + pub unsafe fn deallocate(&mut self, ptr: NonNull, layout: Layout) { + for hole_list in self.holes.iter_mut() { + // If the pointer is in the hole list's memory range, deallocate from that list and exit + if ptr.as_ptr() >= hole_list.bottom && ptr.as_ptr() < hole_list.top { + self.used -= hole_list.deallocate(ptr, layout).size(); + return; + } + } + panic!("The pointer for deallocation does not match any of the allocator's memory ranges."); + } + + /// Returns the size of the heap. + /// + /// This is the size the heap is using for allocations, not necessarily the + /// total amount of bytes given to the heap. To determine the exact memory + /// boundaries, use [`bottom`][Self::bottom] and [`top`][Self::top]. + pub fn size(&self) -> usize { + unsafe { + self.holes + .iter() + .map(|hole_list| hole_list.top.offset_from(hole_list.bottom) as usize) + .sum() + } + } + + /// Returns the size of the used part of the heap + pub fn used(&self) -> usize { + self.used + } + + /// Returns the size of the free part of the heap + pub fn free(&self) -> usize { + self.size() - self.used + } + + /// Initializes an empty heap + /// + /// The `heap_bottom` pointer is automatically aligned, so the [`bottom()`][Self::bottom] + /// method might return a pointer that is larger than `heap_bottom` after construction. + /// + /// The given `heap_size`s must be large enough to store the required + /// metadata, otherwise this function will panic. Depending on the + /// alignment of the `hole_addr` pointer, the minimum size is between + /// `2 * size_of::` and `3 * size_of::`. + /// + /// The usable size for allocations will be truncated to the nearest + /// alignment of `align_of::`. Any extra bytes left at the end + /// will be reclaimed once sufficient additional space is given to + /// [`extend`][Heap::extend]. + /// + /// # Safety + /// + /// This function must be called at most once and must only be used on an + /// empty heap. + /// + /// The bottom addresses must be valid and the memory in the + /// `[heap_bottom, heap_bottom + heap_size)` range must not be used for anything else and must not overlap. + /// This function is unsafe because it can cause undefined behavior if the given addresses + /// are invalid. + /// + /// The provided memory range must be valid for the `'static` lifetime. + pub unsafe fn init_from_pointers(&mut self, heaps: [(*mut u8, usize); SEGMENTS]) { + self.used = 0; + for (index, (heap_bottom, heap_size)) in heaps.iter().enumerate() { + self.holes[index] = HoleList::new(*heap_bottom, *heap_size) + } + } + + /// Initialize an empty heap with provided memory. + /// + /// The caller is responsible for procuring a region of raw memory that may be utilized by the + /// allocator. This might be done via any method such as (unsafely) taking a region from the + /// program's memory, from a mutable static, or by allocating and leaking such memory from + /// another allocator. + /// + /// The latter approach may be especially useful if the underlying allocator does not perform + /// deallocation (e.g. a simple bump allocator). Then the overlaid linked-list-allocator can + /// provide memory reclamation. + /// + /// The usable size for allocations will be truncated to the nearest + /// alignment of `align_of::`. Any extra bytes left at the end + /// will be reclaimed once sufficient additional space is given to + /// [`extend`][Heap::extend]. + /// + /// # Panics + /// + /// This method panics if the heap is already initialized. + /// + /// It also panics when the length of the given `mem` slice is not large enough to + /// store the required metadata. Depending on the alignment of the slice, the minimum + /// size is between `2 * size_of::` and `3 * size_of::`. + pub fn init_from_slices(&mut self, mut regions: [&'static mut [MaybeUninit]; SEGMENTS]) { + assert!( + self.holes.iter().all(|h| h.bottom.is_null()), + "The heap has already been initialized." + ); + for (index, region) in regions.iter_mut().enumerate() { + // SAFETY: All initialization requires the bottom address to be valid, which implies it + // must not be 0. Initially the address is 0. The assertion above ensures that no + // initialization had been called before. + // The given address and size is valid according to the safety invariants of the mutable + // reference handed to us by the caller. + self.holes[index] = unsafe {HoleList::new((*region).as_mut_ptr().cast(), region.len())}; } } +} +impl Heap { /// Initializes an empty heap /// /// The `heap_bottom` pointer is automatically aligned, so the [`bottom()`][Self::bottom] @@ -86,7 +230,7 @@ impl Heap { /// The provided memory range must be valid for the `'static` lifetime. pub unsafe fn init(&mut self, heap_bottom: *mut u8, heap_size: usize) { self.used = 0; - self.holes = HoleList::new(heap_bottom, heap_size); + self.holes[0] = HoleList::new(heap_bottom, heap_size); } /// Initialize an empty heap with provided memory. @@ -153,7 +297,7 @@ impl Heap { pub unsafe fn new(heap_bottom: *mut u8, heap_size: usize) -> Heap { Heap { used: 0, - holes: HoleList::new(heap_bottom, heap_size), + holes: [HoleList::new(heap_bottom, heap_size)], } } @@ -170,55 +314,12 @@ impl Heap { unsafe { Self::new(address, size) } } - /// Allocates a chunk of the given size with the given alignment. Returns a pointer to the - /// beginning of that chunk if it was successful. Else it returns `None`. - /// This function scans the list of free memory blocks and uses the first block that is big - /// enough. The runtime is in O(n) where n is the number of free blocks, but it should be - /// reasonably fast for small allocations. - // - // NOTE: We could probably replace this with an `Option` instead of a `Result` in a later - // release to remove this clippy warning - #[allow(clippy::result_unit_err)] - pub fn allocate_first_fit(&mut self, layout: Layout) -> Result, ()> { - match self.holes.allocate_first_fit(layout) { - Ok((ptr, aligned_layout)) => { - self.used += aligned_layout.size(); - Ok(ptr) - } - Err(err) => Err(err), - } - } - - /// Frees the given allocation. `ptr` must be a pointer returned - /// by a call to the `allocate_first_fit` function with identical size and alignment. - /// - /// This function walks the list of free memory blocks and inserts the freed block at the - /// correct place. If the freed block is adjacent to another free block, the blocks are merged - /// again. This operation is in `O(n)` since the list needs to be sorted by address. - /// - /// # Safety - /// - /// `ptr` must be a pointer returned by a call to the [`allocate_first_fit`] function with - /// identical layout. Undefined behavior may occur for invalid arguments. - pub unsafe fn deallocate(&mut self, ptr: NonNull, layout: Layout) { - self.used -= self.holes.deallocate(ptr, layout).size(); - } - /// Returns the bottom address of the heap. /// /// The bottom pointer is automatically aligned, so the returned pointer /// might be larger than the bottom pointer used for initialization. pub fn bottom(&self) -> *mut u8 { - self.holes.bottom - } - - /// Returns the size of the heap. - /// - /// This is the size the heap is using for allocations, not necessarily the - /// total amount of bytes given to the heap. To determine the exact memory - /// boundaries, use [`bottom`][Self::bottom] and [`top`][Self::top]. - pub fn size(&self) -> usize { - unsafe { self.holes.top.offset_from(self.holes.bottom) as usize } + self.holes[0].bottom } /// Return the top address of the heap. @@ -227,17 +328,7 @@ impl Heap { /// until there is enough room for metadata, but it still retains ownership /// over memory from [`bottom`][Self::bottom] to the address returned. pub fn top(&self) -> *mut u8 { - unsafe { self.holes.top.add(self.holes.pending_extend as usize) } - } - - /// Returns the size of the used part of the heap - pub fn used(&self) -> usize { - self.used - } - - /// Returns the size of the free part of the heap - pub fn free(&self) -> usize { - self.size() - self.used + unsafe { self.holes[0].top.add(self.holes[0].pending_extend as usize) } } /// Extends the size of the heap by creating a new hole at the end. @@ -259,7 +350,7 @@ impl Heap { /// by exactly `by` bytes, those bytes are still owned by the Heap for /// later use. pub unsafe fn extend(&mut self, by: usize) { - self.holes.extend(by); + self.holes[0].extend(by); } } @@ -270,8 +361,8 @@ unsafe impl Allocator for LockedHeap { return Ok(NonNull::slice_from_raw_parts(layout.dangling(), 0)); } match self.0.lock().allocate_first_fit(layout) { - Ok(ptr) => Ok(NonNull::slice_from_raw_parts(ptr, layout.size())), - Err(()) => Err(AllocError), + Some(ptr) => Ok(NonNull::slice_from_raw_parts(ptr, layout.size())), + None => Err(AllocError), } } @@ -312,7 +403,7 @@ impl LockedHeap { pub unsafe fn new(heap_bottom: *mut u8, heap_size: usize) -> LockedHeap { LockedHeap(Spinlock::new(Heap { used: 0, - holes: HoleList::new(heap_bottom, heap_size), + holes: [HoleList::new(heap_bottom, heap_size)], })) } } @@ -332,7 +423,6 @@ unsafe impl GlobalAlloc for LockedHeap { self.0 .lock() .allocate_first_fit(layout) - .ok() .map_or(core::ptr::null_mut(), |allocation| allocation.as_ptr()) } diff --git a/src/test.rs b/src/test.rs index 3ff0514..3f74d8e 100644 --- a/src/test.rs +++ b/src/test.rs @@ -117,7 +117,7 @@ fn new_heap_skip(ct: usize) -> OwnedHeap<1000> { fn empty() { let mut heap = Heap::empty(); let layout = Layout::from_size_align(1, 1).unwrap(); - assert!(heap.allocate_first_fit(layout.clone()).is_err()); + assert!(heap.allocate_first_fit(layout.clone()).is_none()); } #[test] @@ -131,7 +131,7 @@ fn oom() { let layout = Layout::from_size_align(heap.size() + 1, align_of::()); let addr = heap.allocate_first_fit(layout.unwrap()); - assert!(addr.is_err()); + assert!(addr.is_none()); // Explicitly unleak the heap allocation unsafe { Chonk::unleak(heap_space_ptr) }; @@ -143,10 +143,10 @@ fn allocate_double_usize() { let size = size_of::() * 2; let layout = Layout::from_size_align(size, align_of::()); let addr = heap.allocate_first_fit(layout.unwrap()); - assert!(addr.is_ok()); + assert!(addr.is_some()); let addr = addr.unwrap().as_ptr(); assert!(addr == heap.bottom()); - let (hole_addr, hole_size) = heap.holes.first_hole().expect("ERROR: no hole left"); + let (hole_addr, hole_size) = heap.holes[0].first_hole().expect("ERROR: no hole left"); assert!(hole_addr == heap.bottom().wrapping_add(size)); assert!(hole_size == heap.size() - size); @@ -168,7 +168,7 @@ fn allocate_and_free_double_usize() { *(x.as_ptr() as *mut (usize, usize)) = (0xdeafdeadbeafbabe, 0xdeafdeadbeafbabe); heap.deallocate(x, layout.clone()); - let real_first = heap.holes.first.next.as_ref().unwrap().as_ref(); + let real_first = heap.holes[0].first.next.as_ref().unwrap().as_ref(); assert_eq!(real_first.size, heap.size()); assert!(real_first.next.is_none()); @@ -279,7 +279,7 @@ fn allocate_many_size_aligns() { let aligned_heap_size = align_down_size(1000, size_of::()); assert_eq!(heap.size(), aligned_heap_size); - heap.holes.debug(); + heap.holes[0].debug(); let max_alloc = Layout::from_size_align(aligned_heap_size, 1).unwrap(); let full = heap.allocate_first_fit(max_alloc).unwrap(); @@ -287,7 +287,7 @@ fn allocate_many_size_aligns() { heap.deallocate(full, max_alloc); } - heap.holes.debug(); + heap.holes[0].debug(); struct Alloc { alloc: NonNull, @@ -310,9 +310,9 @@ fn allocate_many_size_aligns() { let mut allocs = vec![]; let layout = Layout::from_size_align(size, 1 << align).unwrap(); - while let Ok(alloc) = heap.allocate_first_fit(layout) { + while let Some(alloc) = heap.allocate_first_fit(layout) { #[cfg(not(miri))] - heap.holes.debug(); + heap.holes[0].debug(); allocs.push(Alloc { alloc, layout }); } @@ -325,7 +325,7 @@ fn allocate_many_size_aligns() { allocs.drain(..).for_each(|a| unsafe { heap.deallocate(a.alloc, a.layout); #[cfg(not(miri))] - heap.holes.debug(); + heap.holes[0].debug(); }); } 1 => { @@ -333,7 +333,7 @@ fn allocate_many_size_aligns() { allocs.drain(..).rev().for_each(|a| unsafe { heap.deallocate(a.alloc, a.layout); #[cfg(not(miri))] - heap.holes.debug(); + heap.holes[0].debug(); }); } 2 => { @@ -350,12 +350,12 @@ fn allocate_many_size_aligns() { a.drain(..).for_each(|a| unsafe { heap.deallocate(a.alloc, a.layout); #[cfg(not(miri))] - heap.holes.debug(); + heap.holes[0].debug(); }); b.drain(..).for_each(|a| unsafe { heap.deallocate(a.alloc, a.layout); #[cfg(not(miri))] - heap.holes.debug(); + heap.holes[0].debug(); }); } 3 => { @@ -372,12 +372,12 @@ fn allocate_many_size_aligns() { a.drain(..).for_each(|a| unsafe { heap.deallocate(a.alloc, a.layout); #[cfg(not(miri))] - heap.holes.debug(); + heap.holes[0].debug(); }); b.drain(..).for_each(|a| unsafe { heap.deallocate(a.alloc, a.layout); #[cfg(not(miri))] - heap.holes.debug(); + heap.holes[0].debug(); }); } _ => panic!(), @@ -474,7 +474,7 @@ fn allocate_usize() { let layout = Layout::from_size_align(size_of::(), 1).unwrap(); - assert!(heap.allocate_first_fit(layout.clone()).is_ok()); + assert!(heap.allocate_first_fit(layout.clone()).is_some()); } #[test] @@ -491,7 +491,7 @@ fn allocate_usize_in_bigger_block() { } let z = heap.allocate_first_fit(layout_2.clone()); - assert!(z.is_ok()); + assert!(z.is_some()); let z = z.unwrap(); assert_eq!(x, z); @@ -510,9 +510,9 @@ fn align_from_small_to_big() { let layout_2 = Layout::from_size_align(8, 8).unwrap(); // allocate 28 bytes so that the heap end is only 4 byte aligned - assert!(heap.allocate_first_fit(layout_1.clone()).is_ok()); + assert!(heap.allocate_first_fit(layout_1.clone()).is_some()); // try to allocate a 8 byte aligned block - assert!(heap.allocate_first_fit(layout_2.clone()).is_ok()); + assert!(heap.allocate_first_fit(layout_2.clone()).is_some()); } #[test] @@ -525,7 +525,7 @@ fn extend_empty_heap() { // Try to allocate full heap after extend let layout = Layout::from_size_align(2048, 1).unwrap(); - assert!(heap.allocate_first_fit(layout.clone()).is_ok()); + assert!(heap.allocate_first_fit(layout.clone()).is_some()); } #[test] @@ -535,11 +535,11 @@ fn extend_full_heap() { let layout = Layout::from_size_align(1024, 1).unwrap(); // Allocate full heap, extend and allocate again to the max - assert!(heap.allocate_first_fit(layout.clone()).is_ok()); + assert!(heap.allocate_first_fit(layout.clone()).is_some()); unsafe { heap.extend(1024); } - assert!(heap.allocate_first_fit(layout.clone()).is_ok()); + assert!(heap.allocate_first_fit(layout.clone()).is_some()); } #[test] @@ -552,8 +552,8 @@ fn extend_fragmented_heap() { let alloc1 = heap.allocate_first_fit(layout_1.clone()); let alloc2 = heap.allocate_first_fit(layout_1.clone()); - assert!(alloc1.is_ok()); - assert!(alloc2.is_ok()); + assert!(alloc1.is_some()); + assert!(alloc2.is_some()); unsafe { // Create a hole at the beginning of the heap @@ -566,7 +566,7 @@ fn extend_fragmented_heap() { // We got additional 1024 bytes hole at the end of the heap // Try to allocate there - assert!(heap.allocate_first_fit(layout_2.clone()).is_ok()); + assert!(heap.allocate_first_fit(layout_2.clone()).is_some()); } /// Ensures that `Heap::extend` fails for very small sizes. @@ -580,7 +580,7 @@ fn small_heap_extension() { unsafe { let mut heap = Heap::new(HEAP.as_mut_ptr().cast(), 32); heap.extend(1); - assert_eq!(1, heap.holes.pending_extend); + assert_eq!(1, heap.holes[0].pending_extend); } } @@ -592,7 +592,7 @@ fn oddly_sized_heap_extension() { unsafe { let mut heap = Heap::new(HEAP.as_mut_ptr().cast(), 16); heap.extend(17); - assert_eq!(1, heap.holes.pending_extend); + assert_eq!(1, heap.holes[0].pending_extend); assert_eq!(16 + 16, heap.size()); } } @@ -607,11 +607,11 @@ fn extend_odd_size() { static mut HEAP: [u64; 6] = [0; 6]; unsafe { let mut heap = Heap::new(HEAP.as_mut_ptr().cast(), 17); - assert_eq!(1, heap.holes.pending_extend); + assert_eq!(1, heap.holes[0].pending_extend); heap.extend(16); - assert_eq!(1, heap.holes.pending_extend); + assert_eq!(1, heap.holes[0].pending_extend); heap.extend(15); - assert_eq!(0, heap.holes.pending_extend); + assert_eq!(0, heap.holes[0].pending_extend); assert_eq!(17 + 16 + 15, heap.size()); } }