From 2f8e0cfb16b3b306e560b26b9f4e1ee3e0185e92 Mon Sep 17 00:00:00 2001
From: bigsaltyfishes <bigsaltyfishes@gmail.com>
Date: Fri, 20 Sep 2024 22:15:42 +0800
Subject: [PATCH] Split into smaller blocks if the order of block size is
 larger than max order

Signed-off-by: bigsaltyfishes <bigsaltyfishes@gmail.com>
---
 src/lib.rs  |  8 +++++++-
 src/test.rs | 15 +++++++++++++++
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index e42887e..2362acf 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -88,7 +88,13 @@ impl<const ORDER: usize> Heap<ORDER> {
 
         while current_start + size_of::<usize>() <= end {
             let lowbit = current_start & (!current_start + 1);
-            let size = min(lowbit, prev_power_of_two(end - current_start));
+            let mut size = min(lowbit, prev_power_of_two(end - current_start));
+            
+            // If the order of size is larger than the max order,
+            // split it into smaller blocks.
+            if size.trailing_zeros() as usize > (ORDER - 1) {
+                size = 1 << (ORDER - 1);
+            }
             total += size;
 
             self.free_list[size.trailing_zeros() as usize].push(current_start as *mut usize);
diff --git a/src/test.rs b/src/test.rs
index 7f2e0f1..d25ac23 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -60,6 +60,21 @@ fn test_heap_add() {
     assert!(addr.is_ok());
 }
 
+#[test]
+fn test_heap_add_large() {
+    // Max size of block is 2^7 == 128 bytes
+    let mut heap = Heap::<8>::new();
+    assert!(heap.alloc(Layout::from_size_align(1, 1).unwrap()).is_err());
+
+    // 512 bytes of space
+    let space: [u8; 512] = [0; 512];
+    unsafe {
+        heap.add_to_heap(space.as_ptr() as usize, space.as_ptr().add(100) as usize);
+    }
+    let addr = heap.alloc(Layout::from_size_align(1, 1).unwrap());
+    assert!(addr.is_ok());
+}
+
 #[test]
 fn test_heap_oom() {
     let mut heap = Heap::<32>::new();