Merge changes Iae7130f2,Ibe6629a2,If39f46d9

* changes:
  vmbase: Rename configure_global_allocator_size
  vmbase: Initialize heap in rust_entry
  vmbase: Move heap.rs out of pvmfw
diff --git a/pvmfw/Android.bp b/pvmfw/Android.bp
index 6fa23d8..bbe00b5 100644
--- a/pvmfw/Android.bp
+++ b/pvmfw/Android.bp
@@ -15,7 +15,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..6f96fc0 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_heap, console,
     layout::{self, crosvm},
     logger, main,
     memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_128KB, SIZE_4KB},
@@ -63,7 +61,7 @@
 }
 
 main!(start);
-configure_global_allocator_size!(SIZE_128KB);
+configure_heap!(SIZE_128KB);
 
 /// Entry point for pVM firmware.
 pub fn start(fdt_address: u64, payload_start: u64, payload_size: u64, _arg3: u64) {
@@ -71,9 +69,6 @@
     // - can't access non-pvmfw memory (only statically-mapped memory)
     // - can't access MMIO (therefore, no logging)
 
-    // SAFETY - This function should and will only be called once, here.
-    unsafe { heap::init() };
-
     match main_wrapper(fdt_address as usize, payload_start as usize, payload_size as usize) {
         Ok((entry, bcc)) => jump_to_payload(fdt_address, entry.try_into().unwrap(), bcc),
         Err(_) => reboot(), // TODO(b/220071963) propagate the reason back to the host.
diff --git a/pvmfw/src/main.rs b/pvmfw/src/main.rs
index 032d87e..61e2312 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 673f477..9aa4667 100644
--- a/rialto/Android.bp
+++ b/rialto/Android.bp
@@ -9,7 +9,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..ce83624 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,14 @@
 use libfdt::FdtError;
 use log::{debug, error, info};
 use vmbase::{
+    configure_heap,
     fdt::SwiotlbInfo,
     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 +148,6 @@
 
 /// Entry point for Rialto.
 pub fn main(fdt_addr: u64, _a1: u64, _a2: u64, _a3: u64) {
-    init_heap();
     let Ok(mmio_guard_supported) = try_init_logger() else {
         // Don't log anything if the logger initialization fails.
         reboot();
@@ -181,3 +165,4 @@
 }
 
 main!(main);
+configure_heap!(SIZE_64KB);
diff --git a/vmbase/example/Android.bp b/vmbase/example/Android.bp
index 32854c2..ae1a593 100644
--- a/vmbase/example/Android.bp
+++ b/vmbase/example/Android.bp
@@ -9,7 +9,6 @@
     srcs: ["src/main.rs"],
     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..b3b5732 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_heap, cstr, 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_heap!(SIZE_64KB);
 
 /// Entry point for VM bootloader.
 pub fn main(arg0: u64, arg1: u64, arg2: u64, arg3: u64) {
@@ -74,10 +69,6 @@
 
     modify_fdt(fdt);
 
-    unsafe {
-        HEAP_ALLOCATOR.lock().init(HEAP.as_mut_ptr() as usize, HEAP.len());
-    }
-
     check_alloc();
 
     let mut idmap = IdMap::new(ASID, ROOT_LEVEL);
@@ -164,7 +155,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/vmbase/src/entry.rs b/vmbase/src/entry.rs
index 8cdfe77..df0bb7c 100644
--- a/vmbase/src/entry.rs
+++ b/vmbase/src/entry.rs
@@ -14,11 +14,13 @@
 
 //! Rust entry point.
 
-use crate::{console, power::shutdown};
+use crate::{console, heap, power::shutdown};
 
 /// This is the entry point to the Rust code, called from the binary entry point in `entry.S`.
 #[no_mangle]
 extern "C" fn rust_entry(x0: u64, x1: u64, x2: u64, x3: u64) -> ! {
+    // SAFETY - Only called once, from here, and inaccessible to client code.
+    unsafe { heap::init() };
     console::init();
     unsafe {
         main(x0, x1, x2, x3);
diff --git a/pvmfw/src/heap.rs b/vmbase/src/heap.rs
similarity index 97%
rename from pvmfw/src/heap.rs
rename to vmbase/src/heap.rs
index a28a02c..b00ca6f 100644
--- a/pvmfw/src/heap.rs
+++ b/vmbase/src/heap.rs
@@ -29,7 +29,7 @@
 
 /// Configures the size of the global allocator.
 #[macro_export]
-macro_rules! configure_global_allocator_size {
+macro_rules! configure_heap {
     ($len:expr) => {
         static mut __HEAP_ARRAY: [u8; $len] = [0; $len];
         #[export_name = "HEAP"]
@@ -39,7 +39,7 @@
 }
 
 extern "Rust" {
-    /// Slice used by the global allocator, configured using configure_global_allocator_size!().
+    /// Slice used by the global allocator, configured using configure_heap!().
     static mut HEAP: &'static mut [u8];
 }
 
@@ -51,7 +51,7 @@
 /// # Safety
 ///
 /// Must be called no more than once.
-pub unsafe fn init() {
+pub(crate) 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.
     let (start, size) = unsafe { (HEAP.as_mut_ptr() as usize, HEAP.len()) };
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.