Make it possible to extend the size of the heap

Add a max_size parameter to the heap.
Allocations exceeding the size of the heap will fail,
but one can extend the size of the heap with the call to extend.

That allows heap to be more memory efficient, consuming more memory only
when necessary, eg. os allocator might map additional virtual memory when
the allocation fails.

Author: rafalmiel

src/hole.rs

@@ -40,6 +40,38 @@ impl HoleList {
         }
     }
 
+    /// Places a hole of a given size at a given address after a given hole
+    unsafe fn put_hole(hole: &mut Hole, addr: usize, size: usize) {
+        mem::forget(mem::replace(&mut *(addr as *mut Hole),
+                                 Hole {
+                                     size: size,
+                                     next: None,
+                                 }));
+        hole.next = Some(Unique::new(addr as *mut Hole));
+    }
+
+    /// Extends the size of the list by increasing the size of the last hole
+    /// or appending new hole at the end of the list
+    pub fn extend(&mut self, addr: usize, size: usize) {
+        let mut last: &mut Hole = &mut self.first;
+
+        // Find the last hole in the list
+        while !last.next.is_none() {
+            last =  move_helper(last).next_unwrap();
+        }
+
+        if last.size == 0 || (last as *const _ as usize) + last.size != addr {
+            // Create new hole at the end of the list if the list is empty
+            // or if after last hole there is some allocated memory
+            unsafe {
+                HoleList::put_hole(last, addr, size);
+            }
+        } else {
+            // Last hole reaches the end of the heap so just increase its size
+            last.size += size;
+        }
+    }
+
     /// Searches the list for a big enough hole. A hole is big enough if it can hold an allocation
     /// of `size` bytes with the given `align`. If such a hole is found in the list, a block of the
     /// required size is allocated from it. Then the start address of that block is returned.

src/lib.rs

@@ -17,6 +17,7 @@ mod test;
 pub struct Heap {
     bottom: usize,
     size: usize,
+    max_size: usize,
     holes: HoleList,
 }
 
@@ -26,6 +27,7 @@ impl Heap {
         Heap {
             bottom: 0,
             size: 0,
+            max_size: 0,
             holes: HoleList::empty(),
         }
     }
@@ -39,6 +41,7 @@ impl Heap {
     pub unsafe fn init(&mut self, heap_bottom: usize, heap_size: usize) {
         self.bottom = heap_bottom;
         self.size = heap_size;
+        self.max_size = heap_size;
         self.holes = HoleList::new(heap_bottom, heap_size);
     }
 
@@ -50,6 +53,22 @@ impl Heap {
         Heap {
             bottom: heap_bottom,
             size: heap_size,
+            max_size: heap_size,
+            holes: HoleList::new(heap_bottom, heap_size),
+        }
+    }
+
+    /// Creates a new heap with the given `bottom` and `size` and `max_size`. The bottom address must be valid
+    /// and the memory in the `[heap_bottom, heap_bottom + heap_size)` range must not be used for
+    /// anything else. This function is unsafe because it can cause undefined behavior if the
+    /// given address is invalid.
+    /// Heap size may be extended up to the maximum size with the call to extent
+    pub unsafe fn new_max(heap_bottom: usize, heap_size: usize, heap_max_size: usize) -> Heap {
+        assert!(heap_size <= heap_max_size);
+        Heap {
+            bottom: heap_bottom,
+            size: heap_size,
+            max_size: heap_max_size,
             holes: HoleList::new(heap_bottom, heap_size),
         }
     }
@@ -93,6 +112,34 @@ impl Heap {
     pub fn size(&self) -> usize {
         self.size
     }
+
+    /// Return the maximum size of the heap
+    pub fn max_size(&self) -> usize {
+        self.max_size
+    }
+
+    /// Return the top address of the heap
+    pub fn top(&self) -> usize {
+        self.bottom + self.size
+    }
+
+    /// Returns the maximum top address of the heap
+    pub fn max_top(&self) -> usize {
+        self.bottom + self.max_size
+    }
+
+    /// Extends the size of the heap if it won't exceed max_size
+    /// Returns true if the call succeeded
+    pub fn extend(&mut self, by: usize) -> bool {
+        let top = self.top();
+        if self.size + by <= self.max_size {
+            self.holes.extend(top, by);
+            self.size += by;
+            true
+        } else {
+            false
+        }
+    }
 }
 
 /// Align downwards. Returns the greatest x with alignment `align`

src/test.rs

@@ -12,6 +12,18 @@ fn new_heap() -> Heap {
     heap
 }
 
+fn new_max_heap() -> Heap {
+    const HEAP_SIZE: usize = 1024;
+    const HEAP_SIZE_MAX: usize = 2048;
+    let heap_space = Box::into_raw(Box::new([0u8; HEAP_SIZE_MAX]));
+
+    let heap = unsafe { Heap::new_max(heap_space as usize, HEAP_SIZE, HEAP_SIZE_MAX) };
+    assert!(heap.bottom == heap_space as usize);
+    assert!(heap.size == HEAP_SIZE);
+    assert!(heap.max_size == HEAP_SIZE_MAX);
+    heap
+}
+
 #[test]
 fn empty() {
     let mut heap = Heap::empty();
@@ -201,3 +213,45 @@ fn align_from_small_to_big() {
     // try to allocate a 8 byte aligned block
     assert!(heap.allocate_first_fit(8, 8).is_some());
 }
+
+#[test]
+fn extend_empty_heap() {
+    let mut heap = new_max_heap();
+
+    assert!(heap.extend(1024));
+
+    // Try to allocate full heap after extend
+    assert!(heap.allocate_first_fit(2048, 1).is_some());
+}
+
+#[test]
+fn extend_full_heap() {
+    let mut heap = new_max_heap();
+
+    // Allocate full heap, extend and allocate again to the max
+    assert!(heap.allocate_first_fit(1024, 1).is_some());
+    assert!(heap.extend(1024));
+    assert!(heap.allocate_first_fit(1024, 1).is_some());
+}
+
+#[test]
+fn extend_fragmented_heap() {
+    let mut heap = new_max_heap();
+
+    let alloc1 = heap.allocate_first_fit(512, 1);
+    let alloc2 = heap.allocate_first_fit(512, 1);
+
+    assert!(alloc1.is_some());
+    assert!(alloc2.is_some());
+
+    unsafe {
+        // Create a hole at the beginning of the heap
+        heap.deallocate(alloc1.unwrap(), 512, 1);
+    }
+
+    assert!(heap.extend(1024));
+
+    // We got additional 1024 bytes hole at the end of the heap
+    // Try to allocate there
+    assert!(heap.allocate_first_fit(1024, 1).is_some());
+}