From 855b651c2b7dead276b258146b94b5ee3c353437 Mon Sep 17 00:00:00 2001
From: Andreas Molzer <andreas.molzer@gmx.de>
Date: Wed, 6 Oct 2021 14:09:45 +0200
Subject: [PATCH] Add Heap constructor, utilize in testing

---
 src/lib.rs  | 14 ++++++++++++--
 src/test.rs | 17 ++++++++++-------
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 0bb5c37..b59ac1e 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -17,6 +17,7 @@ use core::alloc::GlobalAlloc;
 use core::alloc::Layout;
 #[cfg(feature = "alloc_ref")]
 use core::alloc::{AllocError, Allocator};
+use core::mem::MaybeUninit;
 #[cfg(feature = "use_spin")]
 use core::ops::Deref;
 use core::ptr::NonNull;
@@ -89,9 +90,9 @@ impl Heap {
     /// This method panics if the heap is already initialized.
     pub fn init_from_slice(&mut self, mem: &'static mut [MaybeUninit<u8>]) {
         assert!(self.bottom == 0, "The heap has already been initialized.");
-        let size = mem.size();
+        let size = mem.len();
         let address = mem.as_ptr() as usize;
-        // Safety: All initialization requires the bottom address to be valid, which implies it
+        // 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
@@ -116,6 +117,15 @@ impl Heap {
         }
     }
 
+    /// Crates a new heap from a slice of raw memory.
+    pub fn with_memory(mem: &'static mut [MaybeUninit<u8>]) -> Heap {
+        let size = mem.len();
+        let address = mem.as_ptr() as usize;
+        // SAFETY: The given address and size is valid according to the safety invariants of the
+        // mutable reference handed to us by the caller.
+        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
diff --git a/src/test.rs b/src/test.rs
index 74cfbd9..822be12 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -1,14 +1,15 @@
 use super::*;
 use core::alloc::Layout;
-use std::mem::{align_of, size_of};
+use std::mem::{align_of, size_of, MaybeUninit};
 use std::prelude::v1::*;
 
 fn new_heap() -> Heap {
     const HEAP_SIZE: usize = 1000;
-    let heap_space = Box::into_raw(Box::new([0u8; HEAP_SIZE]));
+    let heap_space = Box::leak(Box::new([MaybeUninit::uninit(); HEAP_SIZE]));
+    let assumed_location = heap_space.as_ptr() as usize;
 
-    let heap = unsafe { Heap::new(heap_space as usize, HEAP_SIZE) };
-    assert!(heap.bottom == heap_space as usize);
+    let heap = Heap::with_memory(heap_space);
+    assert!(heap.bottom == assumed_location);
     assert!(heap.size == HEAP_SIZE);
     heap
 }
@@ -16,10 +17,12 @@ fn new_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_space = Box::leak(Box::new([MaybeUninit::<u8>::uninit(); HEAP_SIZE_MAX]));
+    let assumed_location = heap_space.as_ptr() as usize;
 
-    let heap = unsafe { Heap::new(heap_space as usize, HEAP_SIZE) };
-    assert!(heap.bottom == heap_space as usize);
+    // Unsafe so that we have provenance over the whole allocation.
+    let heap = unsafe { Heap::new(assumed_location, HEAP_SIZE) };
+    assert!(heap.bottom == assumed_location);
     assert!(heap.size == HEAP_SIZE);
     heap
 }