pvmfw: Prepare heap.rs for move

Allow client code to control the size of the HEAP through a macro.

Make the SAFETY comment a compliant docstring.

Test: TH
Change-Id: I386e0c2e070e8c96f90560881005af827607fecc
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 0d2dfda..2582d55 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -15,6 +15,7 @@
 //! Low-level entry and exit points of pvmfw.
 
 use crate::config;
+use crate::configure_global_allocator_size;
 use crate::crypto;
 use crate::fdt;
 use crate::heap;
@@ -36,7 +37,7 @@
     console,
     layout::{self, crosvm},
     logger, main,
-    memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_4KB},
+    memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_128KB, SIZE_4KB},
     power::reboot,
 };
 use zeroize::Zeroize;
@@ -62,6 +63,7 @@
 }
 
 main!(start);
+configure_global_allocator_size!(SIZE_128KB);
 
 /// Entry point for pVM firmware.
 pub fn start(fdt_address: u64, payload_start: u64, payload_size: u64, _arg3: u64) {
diff --git a/pvmfw/src/heap.rs b/pvmfw/src/heap.rs
index 151049e..a28a02c 100644
--- a/pvmfw/src/heap.rs
+++ b/pvmfw/src/heap.rs
@@ -27,14 +27,30 @@
 
 use buddy_system_allocator::LockedHeap;
 
-/// 128 KiB
-const HEAP_SIZE: usize = 0x20000;
-static mut HEAP: [u8; HEAP_SIZE] = [0; HEAP_SIZE];
+/// Configures the size of the global allocator.
+#[macro_export]
+macro_rules! configure_global_allocator_size {
+    ($len:expr) => {
+        static mut __HEAP_ARRAY: [u8; $len] = [0; $len];
+        #[export_name = "HEAP"]
+        // SAFETY - HEAP will only be accessed once as mut, from init().
+        static mut __HEAP: &'static mut [u8] = unsafe { &mut __HEAP_ARRAY };
+    };
+}
+
+extern "Rust" {
+    /// Slice used by the global allocator, configured using configure_global_allocator_size!().
+    static mut HEAP: &'static mut [u8];
+}
 
 #[global_allocator]
 static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
 
-/// SAFETY: Must be called no more than once.
+/// Initialize the global allocator.
+///
+/// # Safety
+///
+/// Must be called no more than once.
 pub unsafe fn init() {
     // SAFETY: Nothing else accesses this memory, and we hand it over to the heap to manage and
     // never touch it again. The heap is locked, so there cannot be any races.