vmbase: Move heap.rs out of pvmfw

The module isn't pvmfw-specific so move it to vmbase and re-use the now
centralized implementation in Rialto and vmbase_example. This will allow
fully wrapping heap-initialization in vmbase in a future commit.

Deduplicate the dependency on buddy_system_allocator.

Test: TH
Change-Id: If39f46d99a3721001b5784f782577ae0c2a4b89d
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index c9909e6..09453e5 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -16,7 +16,6 @@
     rustlibs: [
         "libaarch64_paging",
         "libbssl_ffi_nostd",
-        "libbuddy_system_allocator",
         "libciborium_nostd",
         "libciborium_io_nostd",
         "libdiced_open_dice_nostd",
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 78383d2..56f9462 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -15,10 +15,8 @@
 //! 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;
 use crate::memory;
 use core::arch::asm;
 use core::mem::{drop, size_of};
@@ -33,7 +31,7 @@
 use log::LevelFilter;
 use vmbase::util::RangeExt as _;
 use vmbase::{
-    console,
+    configure_global_allocator_size, console, heap,
     layout::{self, crosvm},
     logger, main,
     memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_128KB, SIZE_4KB},
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index 8c5d5d5..61af048 100644
--- a/pvmfw/src/main.rs
+++ b/pvmfw/src/main.rs
@@ -28,7 +28,6 @@
 mod exceptions;
 mod fdt;
 mod gpt;
-mod heap;
 mod helpers;
 mod instance;
 mod memory;
@@ -49,6 +48,7 @@
 use pvmfw_avb::Capability;
 use pvmfw_avb::DebugLevel;
 use pvmfw_embedded_key::PUBLIC_KEY;
+use vmbase::heap;
 use vmbase::memory::flush;
 use vmbase::memory::MEMORY;
 use vmbase::virtio::pci;
diff --git a/rialto/Android.bp b/rialto/Android.bp
index 59f8ba2..0a1dda8 100644
--- a/rialto/Android.bp
+++ b/rialto/Android.bp
@@ -10,7 +10,6 @@
     defaults: ["vmbase_ffi_defaults"],
     rustlibs: [
         "libaarch64_paging",
-        "libbuddy_system_allocator",
         "libhyp",
         "libfdtpci",
         "liblibfdt",
diff --git a/rialto/src/main.rs b/rialto/src/main.rs
index 29056f1..f3f29fe 100644
--- a/rialto/src/main.rs
+++ b/rialto/src/main.rs
@@ -23,7 +23,6 @@
 extern crate alloc;
 
 use crate::error::{Error, Result};
-use buddy_system_allocator::LockedHeap;
 use core::num::NonZeroUsize;
 use core::result;
 use core::slice;
@@ -32,28 +31,15 @@
 use libfdt::FdtError;
 use log::{debug, error, info};
 use vmbase::{
+    configure_global_allocator_size,
     fdt::SwiotlbInfo,
+    heap,
     layout::{self, crosvm},
     main,
-    memory::{MemoryTracker, PageTable, MEMORY, PAGE_SIZE},
+    memory::{MemoryTracker, PageTable, MEMORY, PAGE_SIZE, SIZE_64KB},
     power::reboot,
 };
 
-const SZ_1K: usize = 1024;
-const SZ_64K: usize = 64 * SZ_1K;
-
-#[global_allocator]
-static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
-
-static mut HEAP: [u8; SZ_64K] = [0; SZ_64K];
-
-fn init_heap() {
-    // SAFETY: Allocator set to otherwise unused, static memory.
-    unsafe {
-        HEAP_ALLOCATOR.lock().init(&mut HEAP as *mut u8 as usize, HEAP.len());
-    }
-}
-
 fn new_page_table() -> Result<PageTable> {
     let mut page_table = PageTable::default();
 
@@ -163,7 +149,9 @@
 
 /// Entry point for Rialto.
 pub fn main(fdt_addr: u64, _a1: u64, _a2: u64, _a3: u64) {
-    init_heap();
+    // SAFETY - Only called once, from here.
+    unsafe { heap::init() };
+
     let Ok(mmio_guard_supported) = try_init_logger() else {
         // Don't log anything if the logger initialization fails.
         reboot();
@@ -181,3 +169,4 @@
 }
 
 main!(main);
+configure_global_allocator_size!(SIZE_64KB);
diff --git a/vmbase/example/Android.bp b/vmbase/example/Android.bp
index dc9a090..ab8a8a1 100644
--- a/vmbase/example/Android.bp
+++ b/vmbase/example/Android.bp
@@ -10,7 +10,6 @@
     edition: "2021",
     rustlibs: [
         "libaarch64_paging",
-        "libbuddy_system_allocator",
         "libdiced_open_dice_nostd",
         "libfdtpci",
         "liblibfdt",
diff --git a/vmbase/example/src/main.rs b/vmbase/example/src/main.rs
index 1dd8517..d604509 100644
--- a/vmbase/example/src/main.rs
+++ b/vmbase/example/src/main.rs
@@ -30,11 +30,10 @@
 use crate::pci::{check_pci, get_bar_region};
 use aarch64_paging::{idmap::IdMap, paging::Attributes};
 use alloc::{vec, vec::Vec};
-use buddy_system_allocator::LockedHeap;
 use fdtpci::PciInfo;
 use libfdt::Fdt;
 use log::{debug, error, info, trace, warn, LevelFilter};
-use vmbase::{cstr, logger, main};
+use vmbase::{configure_global_allocator_size, cstr, heap, logger, main, memory::SIZE_64KB};
 
 static INITIALISED_DATA: [u32; 4] = [1, 2, 3, 4];
 static mut ZEROED_DATA: [u32; 10] = [0; 10];
@@ -43,12 +42,8 @@
 const ASID: usize = 1;
 const ROOT_LEVEL: usize = 1;
 
-#[global_allocator]
-static HEAP_ALLOCATOR: LockedHeap<32> = LockedHeap::<32>::new();
-
-static mut HEAP: [u8; 65536] = [0; 65536];
-
 main!(main);
+configure_global_allocator_size!(SIZE_64KB);
 
 /// Entry point for VM bootloader.
 pub fn main(arg0: u64, arg1: u64, arg2: u64, arg3: u64) {
@@ -74,9 +69,8 @@
 
     modify_fdt(fdt);
 
-    unsafe {
-        HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
-    }
+    // SAFETY - Only called once, from here.
+    unsafe { heap::init() };
 
     check_alloc();
 
@@ -164,7 +158,6 @@
     unsafe {
         info!("ZEROED_DATA: {:?}", ZEROED_DATA.as_ptr());
         info!("MUTABLE_DATA: {:?}", MUTABLE_DATA.as_ptr());
-        info!("HEAP: {:?}", HEAP.as_ptr());
     }
 
     assert_eq!(INITIALISED_DATA[0], 1);
diff --git a/pvmfw/src/heap.rs b/vmbase/src/heap.rs
similarity index 100%
rename from pvmfw/src/heap.rs
rename to vmbase/src/heap.rs
diff --git a/vmbase/src/lib.rs b/vmbase/src/lib.rs
index 7fc7b20..88bad8b 100644
--- a/vmbase/src/lib.rs
+++ b/vmbase/src/lib.rs
@@ -23,6 +23,7 @@
 pub mod console;
 mod entry;
 pub mod fdt;
+pub mod heap;
 mod hvc;
 pub mod layout;
 mod linker;
diff --git a/vmbase/src/memory/mod.rs b/vmbase/src/memory/mod.rs
index 6bc600d..5e78565 100644
--- a/vmbase/src/memory/mod.rs
+++ b/vmbase/src/memory/mod.rs
@@ -25,5 +25,5 @@
 pub use shared::{alloc_shared, dealloc_shared, MemoryRange, MemoryTracker, MEMORY};
 pub use util::{
     flush, flushed_zeroize, min_dcache_line_size, page_4kb_of, phys_to_virt, virt_to_phys,
-    PAGE_SIZE, SIZE_128KB, SIZE_2MB, SIZE_4KB, SIZE_4MB,
+    PAGE_SIZE, SIZE_128KB, SIZE_2MB, SIZE_4KB, SIZE_4MB, SIZE_64KB,
 };
diff --git a/vmbase/src/memory/util.rs b/vmbase/src/memory/util.rs
index 48007f3..b9ef5c9 100644
--- a/vmbase/src/memory/util.rs
+++ b/vmbase/src/memory/util.rs
@@ -22,6 +22,8 @@
 
 /// The size of a 4KB memory in bytes.
 pub const SIZE_4KB: usize = 4 << 10;
+/// The size of a 64KB memory in bytes.
+pub const SIZE_64KB: usize = 64 << 10;
 /// The size of a 128KB memory in bytes.
 pub const SIZE_128KB: usize = 128 << 10;
 /// The size of a 2MB memory in bytes.