From c80fc8cf2008ccf6d725c447d07e68764b425540 Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Mon, 25 Aug 2025 15:56:21 -0700
Subject: [PATCH 01/10] port simple bump alloc

---
 openhcl/openhcl_boot/src/bump_alloc.rs        | 128 ++++++++++++++++++
 openhcl/openhcl_boot/src/host_params/dt.rs    |  23 ++++
 .../src/host_params/shim_params.rs            |  10 ++
 openhcl/openhcl_boot/src/main.rs              |   3 +
 vm/loader/loader_defs/src/shim.rs             |   4 +
 vm/loader/src/paravisor.rs                    |  39 ++++++
 6 files changed, 207 insertions(+)
 create mode 100644 openhcl/openhcl_boot/src/bump_alloc.rs

diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/bump_alloc.rs
new file mode 100644
index 0000000000..24a8cfb716
--- /dev/null
+++ b/openhcl/openhcl_boot/src/bump_alloc.rs
@@ -0,0 +1,128 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+//! Simple bump allocator using https://os.phil-opp.com/allocator-designs/ as a
+//! reference and starting point.
+//!
+//! Note that we only allow allocations in a small window for supporting
+//! mesh_protobuf. Any other attempts to allocate will result in a panic.
+
+use crate::boot_logger::log;
+use crate::single_threaded::SingleThreaded;
+use core::alloc::GlobalAlloc;
+use core::alloc::Layout;
+use core::cell::RefCell;
+use memory_range::MemoryRange;
+
+#[global_allocator]
+pub static ALLOCATOR: BumpAllocator = BumpAllocator::new();
+
+#[derive(Debug)]
+pub struct Inner {
+    mem: MemoryRange,
+    next: usize,
+    allow_alloc: bool,
+    alloc_count: usize,
+}
+
+pub struct BumpAllocator {
+    inner: SingleThreaded<RefCell<Inner>>,
+}
+
+/// Align downwards. Returns the greatest x with alignment `align`
+/// so that x <= addr. The alignment must be a power of 2.
+pub fn align_down(addr: usize, align: usize) -> usize {
+    if align.is_power_of_two() {
+        addr & !(align - 1)
+    } else if align == 0 {
+        addr
+    } else {
+        panic!("`align` must be a power of 2");
+    }
+}
+
+/// Align upwards. Returns the smallest x with alignment `align`
+/// so that x >= addr. The alignment must be a power of 2.
+pub fn align_up(addr: usize, align: usize) -> usize {
+    align_down(addr + align - 1, align)
+}
+
+impl BumpAllocator {
+    pub const fn new() -> Self {
+        BumpAllocator {
+            inner: SingleThreaded(RefCell::new(Inner {
+                mem: MemoryRange::EMPTY,
+                next: 0,
+                allow_alloc: false,
+                alloc_count: 0,
+            })),
+        }
+    }
+
+    /// Initialize the bump allocator with the specified memory range.
+    ///
+    /// # Safety
+    /// The caller must guarantee that the memory range is both valid to
+    /// access via the current pagetable identity map, and that it is unused.
+    pub unsafe fn init(&self, mem: MemoryRange) {
+        let mut inner = self.inner.borrow_mut();
+        assert_eq!(
+            inner.mem,
+            MemoryRange::EMPTY,
+            "bump allocator memory range previously set {}",
+            inner.mem
+        );
+
+        inner.mem = mem;
+        inner.next = mem.start() as usize;
+    }
+
+    pub fn enable_alloc(&self) {
+        let mut inner = self.inner.borrow_mut();
+        inner.allow_alloc = true;
+    }
+
+    pub fn disable_alloc(&self) {
+        let mut inner = self.inner.borrow_mut();
+        inner.allow_alloc = false;
+    }
+
+    pub fn log_stats(&self) {
+        let inner = self.inner.borrow();
+        log!(
+            "Bump allocator: allocated {} bytes in {} allocations ({} bytes free)",
+            inner.next - inner.mem.start() as usize,
+            inner.alloc_count,
+            inner.mem.end() as usize - inner.next
+        );
+    }
+}
+
+// SAFETY: The allocator points to a valid identity mapped memory range via
+// init, which is unused by anything else.
+unsafe impl GlobalAlloc for BumpAllocator {
+    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+        let mut inner = self.inner.borrow_mut();
+
+        if !inner.allow_alloc {
+            panic!("allocations are not allowed");
+        }
+
+        let alloc_start: usize = align_up(inner.next, layout.align());
+
+        let alloc_end = match alloc_start.checked_add(layout.size()) {
+            Some(end) => end,
+            None => return core::ptr::null_mut(),
+        };
+
+        if alloc_end > inner.mem.end() as usize {
+            core::ptr::null_mut() // out of memory
+        } else {
+            inner.next = alloc_end;
+            inner.alloc_count += 1;
+            alloc_start as *mut u8
+        }
+    }
+
+    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
+}
diff --git a/openhcl/openhcl_boot/src/host_params/dt.rs b/openhcl/openhcl_boot/src/host_params/dt.rs
index b36ce2821e..b9d0043557 100644
--- a/openhcl/openhcl_boot/src/host_params/dt.rs
+++ b/openhcl/openhcl_boot/src/host_params/dt.rs
@@ -383,6 +383,29 @@ impl PartitionInfo {
         storage.vmbus_vtl2 = parsed.vmbus_vtl2.clone().ok_or(DtError::Vtl2Vmbus)?;
         storage.vmbus_vtl0 = parsed.vmbus_vtl0.clone().ok_or(DtError::Vtl0Vmbus)?;
 
+        // Initialize the temporary heap.
+        //
+        // This is only to be enabled for mesh decode.
+        //
+        // SAFETY: The heap range is reserved at file build time, and is
+        // guaranteed to be unused by anything else.
+        unsafe {
+            crate::bump_alloc::ALLOCATOR.init(params.heap);
+        }
+
+        // TODO: test using heap, as no mesh decode yet.
+        #[cfg(debug_assertions)]
+        {
+            use alloc::boxed::Box;
+            crate::bump_alloc::ALLOCATOR.enable_alloc();
+
+            let box_int = Box::new(42);
+            log!("box int {box_int}");
+            drop(box_int);
+            crate::bump_alloc::ALLOCATOR.disable_alloc();
+            crate::bump_alloc::ALLOCATOR.log_stats();
+        }
+
         // The host is responsible for allocating MMIO ranges for non-isolated
         // guests when it also provides the ram VTL2 should use.
         //
diff --git a/openhcl/openhcl_boot/src/host_params/shim_params.rs b/openhcl/openhcl_boot/src/host_params/shim_params.rs
index 1dff9ad5a4..ae75aee119 100644
--- a/openhcl/openhcl_boot/src/host_params/shim_params.rs
+++ b/openhcl/openhcl_boot/src/host_params/shim_params.rs
@@ -105,6 +105,8 @@ pub struct ShimParams {
     pub page_tables: Option<MemoryRange>,
     /// Log buffer region used by the shim.
     pub log_buffer: MemoryRange,
+    /// Memory to be used for the heap.
+    pub heap: MemoryRange,
 }
 
 impl ShimParams {
@@ -135,6 +137,8 @@ impl ShimParams {
             page_tables_size,
             log_buffer_start,
             log_buffer_size,
+            heap_start_offset,
+            heap_size,
         } = raw;
 
         let isolation_type = get_isolation_type(supported_isolation_type);
@@ -158,6 +162,11 @@ impl ShimParams {
             MemoryRange::new(base..base + log_buffer_size)
         };
 
+        let heap = {
+            let base = shim_base_address.wrapping_add_signed(heap_start_offset);
+            MemoryRange::new(base..base + heap_size)
+        };
+
         Self {
             kernel_entry_address: shim_base_address.wrapping_add_signed(kernel_entry_offset),
             cmdline_base: shim_base_address.wrapping_add_signed(cmdline_offset),
@@ -182,6 +191,7 @@ impl ShimParams {
             bounce_buffer,
             page_tables,
             log_buffer,
+            heap,
         }
     }
 
diff --git a/openhcl/openhcl_boot/src/main.rs b/openhcl/openhcl_boot/src/main.rs
index 33da2d90a8..b562396720 100644
--- a/openhcl/openhcl_boot/src/main.rs
+++ b/openhcl/openhcl_boot/src/main.rs
@@ -11,6 +11,7 @@
 
 mod arch;
 mod boot_logger;
+mod bump_alloc;
 mod cmdline;
 mod dt;
 mod host_params;
@@ -20,6 +21,8 @@ mod rt;
 mod sidecar;
 mod single_threaded;
 
+extern crate alloc;
+
 use crate::arch::setup_vtl2_memory;
 use crate::arch::setup_vtl2_vp;
 #[cfg(target_arch = "x86_64")]
diff --git a/vm/loader/loader_defs/src/shim.rs b/vm/loader/loader_defs/src/shim.rs
index ac8464f645..ebc3dc1800 100644
--- a/vm/loader/loader_defs/src/shim.rs
+++ b/vm/loader/loader_defs/src/shim.rs
@@ -61,6 +61,10 @@ pub struct ShimParamsRaw {
     pub log_buffer_start: i64,
     /// The size of the persisted bootshim log buffer.
     pub log_buffer_size: u64,
+    /// The offset to the start of the bootshim heap.
+    pub heap_start_offset: i64,
+    /// The size of the bootshim heap.
+    pub heap_size: u64,
 }
 
 open_enum! {
diff --git a/vm/loader/src/paravisor.rs b/vm/loader/src/paravisor.rs
index f74a029e69..bd6982db8f 100644
--- a/vm/loader/src/paravisor.rs
+++ b/vm/loader/src/paravisor.rs
@@ -136,6 +136,7 @@ where
     // free space
     //
     // page tables
+    // 16 pages reserved for bootshim heap
     // 8K bootshim logs
     // IGVM parameters
     // reserved vtl2 ranges
@@ -350,6 +351,23 @@ where
         &[],
     )?;
 
+    // Reserve 16 pages for a bootshim heap. This is only used to parse the
+    // protobuf payload from the previous instance in a servicing boot.
+    //
+    // Import these pages as it greatly simplifies the early startup code in the
+    // bootshim for isolated guests. This allows the bootshim to use these pages
+    // early on without extra acceptance calls.
+    let heap_start = offset;
+    let heap_size = 16 * HV_PAGE_SIZE;
+    importer.import_pages(
+        heap_start / HV_PAGE_SIZE,
+        heap_size / HV_PAGE_SIZE,
+        "ohcl-boot-shim-heap",
+        BootPageAcceptance::Exclusive,
+        &[],
+    )?;
+    offset += heap_size;
+
     // The end of memory used by the loader, excluding pagetables.
     let end_of_underhill_mem = offset;
 
@@ -508,6 +526,8 @@ where
         page_tables_size: page_table_region_size,
         log_buffer_start: calculate_shim_offset(bootshim_log_start),
         log_buffer_size: bootshim_log_size,
+        heap_start_offset: calculate_shim_offset(heap_start),
+        heap_size,
     };
 
     tracing::debug!(boot_params_base, "shim gpa");
@@ -1062,6 +1082,23 @@ where
         &[],
     )?;
 
+    // Reserve 16 pages for a bootshim heap. This is only used to parse the
+    // protobuf payload from the previous instance in a servicing boot.
+    //
+    // Import these pages as it greatly simplifies the early startup code in the
+    // bootshim for isolated guests. This allows the bootshim to use these pages
+    // early on without extra acceptance calls.
+    let heap_start = next_addr;
+    let heap_size = 16 * HV_PAGE_SIZE;
+    importer.import_pages(
+        heap_start / HV_PAGE_SIZE,
+        heap_size / HV_PAGE_SIZE,
+        "ohcl-boot-shim-heap",
+        BootPageAcceptance::Exclusive,
+        &[],
+    )?;
+    next_addr += heap_size;
+
     // The end of memory used by the loader, excluding pagetables.
     let end_of_underhill_mem = next_addr;
 
@@ -1115,6 +1152,8 @@ where
         page_tables_size: 0,
         log_buffer_start: calculate_shim_offset(bootshim_log_start),
         log_buffer_size: bootshim_log_size,
+        heap_start_offset: calculate_shim_offset(heap_start),
+        heap_size,
     };
 
     importer

From 075a40318c7af68bf41104443e15bcc6571c1c3a Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Thu, 11 Sep 2025 11:00:07 -0700
Subject: [PATCH 02/10] only set global allocator when running under minimal_rt
 to allow unittests

---
 openhcl/openhcl_boot/src/bump_alloc.rs     |  2 +-
 openhcl/openhcl_boot/src/host_params/dt.rs | 49 ++++++++++++----------
 2 files changed, 28 insertions(+), 23 deletions(-)

diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/bump_alloc.rs
index 24a8cfb716..bd39ad114e 100644
--- a/openhcl/openhcl_boot/src/bump_alloc.rs
+++ b/openhcl/openhcl_boot/src/bump_alloc.rs
@@ -14,7 +14,7 @@ use core::alloc::Layout;
 use core::cell::RefCell;
 use memory_range::MemoryRange;
 
-#[global_allocator]
+#[cfg_attr(minimal_rt, global_allocator)]
 pub static ALLOCATOR: BumpAllocator = BumpAllocator::new();
 
 #[derive(Debug)]
diff --git a/openhcl/openhcl_boot/src/host_params/dt.rs b/openhcl/openhcl_boot/src/host_params/dt.rs
index b9d0043557..f237d84027 100644
--- a/openhcl/openhcl_boot/src/host_params/dt.rs
+++ b/openhcl/openhcl_boot/src/host_params/dt.rs
@@ -311,6 +311,32 @@ fn parse_host_vtl2_ram(
     vtl2_ram
 }
 
+#[cfg_attr(minimal_rt, allow(dead_code))]
+fn init_heap(params: &ShimParams) {
+    // Initialize the temporary heap.
+    //
+    // This is only to be enabled for mesh decode.
+    //
+    // SAFETY: The heap range is reserved at file build time, and is
+    // guaranteed to be unused by anything else.
+    unsafe {
+        crate::bump_alloc::ALLOCATOR.init(params.heap);
+    }
+
+    // TODO: test using heap, as no mesh decode yet.
+    #[cfg(debug_assertions)]
+    {
+        use alloc::boxed::Box;
+        crate::bump_alloc::ALLOCATOR.enable_alloc();
+
+        let box_int = Box::new(42);
+        log!("box int {box_int}");
+        drop(box_int);
+        crate::bump_alloc::ALLOCATOR.disable_alloc();
+        crate::bump_alloc::ALLOCATOR.log_stats();
+    }
+}
+
 impl PartitionInfo {
     // Read the IGVM provided DT for the vtl2 partition info.
     pub fn read_from_dt<'a>(
@@ -383,28 +409,7 @@ impl PartitionInfo {
         storage.vmbus_vtl2 = parsed.vmbus_vtl2.clone().ok_or(DtError::Vtl2Vmbus)?;
         storage.vmbus_vtl0 = parsed.vmbus_vtl0.clone().ok_or(DtError::Vtl0Vmbus)?;
 
-        // Initialize the temporary heap.
-        //
-        // This is only to be enabled for mesh decode.
-        //
-        // SAFETY: The heap range is reserved at file build time, and is
-        // guaranteed to be unused by anything else.
-        unsafe {
-            crate::bump_alloc::ALLOCATOR.init(params.heap);
-        }
-
-        // TODO: test using heap, as no mesh decode yet.
-        #[cfg(debug_assertions)]
-        {
-            use alloc::boxed::Box;
-            crate::bump_alloc::ALLOCATOR.enable_alloc();
-
-            let box_int = Box::new(42);
-            log!("box int {box_int}");
-            drop(box_int);
-            crate::bump_alloc::ALLOCATOR.disable_alloc();
-            crate::bump_alloc::ALLOCATOR.log_stats();
-        }
+        init_heap(params);
 
         // The host is responsible for allocating MMIO ranges for non-isolated
         // guests when it also provides the ram VTL2 should use.

From 60267faddae742a4da26b14ac50b1c6afa61f96f Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Fri, 12 Sep 2025 10:50:26 -0700
Subject: [PATCH 03/10] allocator bug debugging

---
 openhcl/openhcl_boot/src/bump_alloc.rs | 71 +++++++++++++++++++++-----
 1 file changed, 58 insertions(+), 13 deletions(-)

diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/bump_alloc.rs
index bd39ad114e..002116501f 100644
--- a/openhcl/openhcl_boot/src/bump_alloc.rs
+++ b/openhcl/openhcl_boot/src/bump_alloc.rs
@@ -29,22 +29,10 @@ pub struct BumpAllocator {
     inner: SingleThreaded<RefCell<Inner>>,
 }
 
-/// Align downwards. Returns the greatest x with alignment `align`
-/// so that x <= addr. The alignment must be a power of 2.
-pub fn align_down(addr: usize, align: usize) -> usize {
-    if align.is_power_of_two() {
-        addr & !(align - 1)
-    } else if align == 0 {
-        addr
-    } else {
-        panic!("`align` must be a power of 2");
-    }
-}
-
 /// Align upwards. Returns the smallest x with alignment `align`
 /// so that x >= addr. The alignment must be a power of 2.
 pub fn align_up(addr: usize, align: usize) -> usize {
-    align_down(addr + align - 1, align)
+    (addr + align - 1) & !(align - 1)
 }
 
 impl BumpAllocator {
@@ -115,6 +103,14 @@ unsafe impl GlobalAlloc for BumpAllocator {
             None => return core::ptr::null_mut(),
         };
 
+        log!(
+            "bump_alloc: allocating {} bytes with alignment {} at {:#x}, alloc_end {:#x}",
+            layout.size(),
+            layout.align(),
+            alloc_start,
+            alloc_end,
+        );
+
         if alloc_end > inner.mem.end() as usize {
             core::ptr::null_mut() // out of memory
         } else {
@@ -126,3 +122,52 @@ unsafe impl GlobalAlloc for BumpAllocator {
 
     unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_align_up() {
+        assert_eq!(align_up(0x1000, 0x1000), 0x1000);
+        assert_eq!(align_up(0x1001, 0x1000), 0x2000);
+        assert_eq!(align_up(0x1FFF, 0x1000), 0x2000);
+        assert_eq!(align_up(0x2000, 0x1000), 0x2000);
+
+        assert_eq!(align_up(0x1003, 4), 0x1004);
+        assert_eq!(align_up(0x1003, 8), 0x1008);
+        assert_eq!(align_up(0x1003, 16), 0x1010);
+    }
+
+    #[test]
+    fn test_alloc() {
+        let allocator = BumpAllocator::new();
+        // create a new page aligned box of memory with 20 pages.
+        let mut buffer = Box::new([0; 0x1000 * 20]);
+        let base_addr = buffer.as_ptr() as usize;
+        // align up base_addr to the next page.
+        let base_addr = align_up(base_addr as usize, 0x1000) as u64;
+        assert_eq!(base_addr & 0xFFF, 0); // ensure page aligned
+
+        let buffer_range = MemoryRange::new(base_addr..(base_addr + 10 * 0x1000));
+
+        unsafe {
+            allocator.init(buffer_range);
+        }
+        allocator.enable_alloc();
+
+        unsafe {
+            let ptr1 = allocator.alloc(Layout::from_size_align(100, 8).unwrap());
+            *ptr1 = 42;
+            assert_eq!(*ptr1, 42);
+
+            let ptr2 = allocator.alloc(Layout::from_size_align(200, 16).unwrap());
+            *ptr2 = 55;
+            assert_eq!(*ptr2, 55);
+
+            let ptr3 = allocator.alloc(Layout::from_size_align(300, 32).unwrap());
+            *ptr3 = 77;
+            assert_eq!(*ptr3, 77);
+        }
+    }
+}

From 3cad39d09ef89e7fe1c0f1d50ea3c419d4010159 Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Fri, 12 Sep 2025 12:51:34 -0700
Subject: [PATCH 04/10] move allocator to use pointers, not raw u64

---
 openhcl/openhcl_boot/src/bump_alloc.rs | 75 ++++++++++++--------------
 1 file changed, 35 insertions(+), 40 deletions(-)

diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/bump_alloc.rs
index 002116501f..2c62323096 100644
--- a/openhcl/openhcl_boot/src/bump_alloc.rs
+++ b/openhcl/openhcl_boot/src/bump_alloc.rs
@@ -19,8 +19,9 @@ pub static ALLOCATOR: BumpAllocator = BumpAllocator::new();
 
 #[derive(Debug)]
 pub struct Inner {
-    mem: MemoryRange,
-    next: usize,
+    start: *mut u8,
+    next: *mut u8,
+    end: *mut u8,
     allow_alloc: bool,
     alloc_count: usize,
 }
@@ -29,18 +30,13 @@ pub struct BumpAllocator {
     inner: SingleThreaded<RefCell<Inner>>,
 }
 
-/// Align upwards. Returns the smallest x with alignment `align`
-/// so that x >= addr. The alignment must be a power of 2.
-pub fn align_up(addr: usize, align: usize) -> usize {
-    (addr + align - 1) & !(align - 1)
-}
-
 impl BumpAllocator {
     pub const fn new() -> Self {
         BumpAllocator {
             inner: SingleThreaded(RefCell::new(Inner {
-                mem: MemoryRange::EMPTY,
-                next: 0,
+                start: core::ptr::null_mut(),
+                next: core::ptr::null_mut(),
+                end: core::ptr::null_mut(),
                 allow_alloc: false,
                 alloc_count: 0,
             })),
@@ -50,19 +46,20 @@ impl BumpAllocator {
     /// Initialize the bump allocator with the specified memory range.
     ///
     /// # Safety
+    ///
     /// The caller must guarantee that the memory range is both valid to
     /// access via the current pagetable identity map, and that it is unused.
     pub unsafe fn init(&self, mem: MemoryRange) {
         let mut inner = self.inner.borrow_mut();
-        assert_eq!(
-            inner.mem,
-            MemoryRange::EMPTY,
-            "bump allocator memory range previously set {}",
-            inner.mem
+        assert!(
+            inner.start.is_null(),
+            "bump allocator memory range previously set {:#x?}",
+            inner.start
         );
 
-        inner.mem = mem;
-        inner.next = mem.start() as usize;
+        inner.start = mem.start() as *mut u8;
+        inner.next = mem.start() as *mut u8;
+        inner.end = mem.end() as *mut u8;
     }
 
     pub fn enable_alloc(&self) {
@@ -77,11 +74,15 @@ impl BumpAllocator {
 
     pub fn log_stats(&self) {
         let inner = self.inner.borrow();
+
+        // FIXME: unsafe calcs
+        let allocated = unsafe { inner.next.offset_from(inner.start) };
+        let free = unsafe { inner.end.offset_from(inner.next) };
         log!(
             "Bump allocator: allocated {} bytes in {} allocations ({} bytes free)",
-            inner.next - inner.mem.start() as usize,
+            allocated,
             inner.alloc_count,
-            inner.mem.end() as usize - inner.next
+            free
         );
     }
 }
@@ -96,49 +97,43 @@ unsafe impl GlobalAlloc for BumpAllocator {
             panic!("allocations are not allowed");
         }
 
-        let alloc_start: usize = align_up(inner.next, layout.align());
+        // FIXME: verify math and wraparounds
+        let align_offset = inner.next.align_offset(layout.align());
+        let alloc_start = inner.next.wrapping_add(align_offset);
+        let alloc_end = alloc_start.wrapping_add(layout.size());
 
-        let alloc_end = match alloc_start.checked_add(layout.size()) {
-            Some(end) => end,
-            None => return core::ptr::null_mut(),
-        };
+        if alloc_end < alloc_start {
+            // overflow
+            return core::ptr::null_mut();
+        }
 
         log!(
-            "bump_alloc: allocating {} bytes with alignment {} at {:#x}, alloc_end {:#x}",
+            "bump_alloc: allocating {} bytes with alignment {} at with offset {} alloc_start {:#x?}, alloc_end {:#x?}",
             layout.size(),
             layout.align(),
+            align_offset,
             alloc_start,
             alloc_end,
         );
 
-        if alloc_end > inner.mem.end() as usize {
+        if alloc_end > inner.end {
             core::ptr::null_mut() // out of memory
         } else {
             inner.next = alloc_end;
             inner.alloc_count += 1;
-            alloc_start as *mut u8
+            alloc_start
         }
     }
 
-    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {}
+    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        log!("dealloc called on {:#x?} of size {}", ptr, layout.size());
+    }
 }
 
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    #[test]
-    fn test_align_up() {
-        assert_eq!(align_up(0x1000, 0x1000), 0x1000);
-        assert_eq!(align_up(0x1001, 0x1000), 0x2000);
-        assert_eq!(align_up(0x1FFF, 0x1000), 0x2000);
-        assert_eq!(align_up(0x2000, 0x1000), 0x2000);
-
-        assert_eq!(align_up(0x1003, 4), 0x1004);
-        assert_eq!(align_up(0x1003, 8), 0x1008);
-        assert_eq!(align_up(0x1003, 16), 0x1010);
-    }
-
     #[test]
     fn test_alloc() {
         let allocator = BumpAllocator::new();

From 07bb47fa25499610fbe84a68ebb236105937da26 Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Fri, 12 Sep 2025 12:51:34 -0700
Subject: [PATCH 05/10] add nightly miri tests

---
 openhcl/openhcl_boot/Cargo.toml        |  3 ++
 openhcl/openhcl_boot/src/bump_alloc.rs | 59 ++++++++++++++++++++------
 openhcl/openhcl_boot/src/main.rs       |  1 +
 3 files changed, 50 insertions(+), 13 deletions(-)

diff --git a/openhcl/openhcl_boot/Cargo.toml b/openhcl/openhcl_boot/Cargo.toml
index bd9feb399c..effb7aabeb 100644
--- a/openhcl/openhcl_boot/Cargo.toml
+++ b/openhcl/openhcl_boot/Cargo.toml
@@ -6,6 +6,9 @@ name = "openhcl_boot"
 edition.workspace = true
 rust-version.workspace = true
 
+[features]
+nightly = []
+
 [dependencies]
 aarch64defs.workspace = true
 minimal_rt.workspace = true
diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/bump_alloc.rs
index 2c62323096..06f314bccc 100644
--- a/openhcl/openhcl_boot/src/bump_alloc.rs
+++ b/openhcl/openhcl_boot/src/bump_alloc.rs
@@ -130,25 +130,45 @@ unsafe impl GlobalAlloc for BumpAllocator {
     }
 }
 
+#[cfg(feature = "nightly")]
+unsafe impl core::alloc::Allocator for BumpAllocator {
+    fn allocate(&self, layout: Layout) -> Result<core::ptr::NonNull<[u8]>, std::alloc::AllocError> {
+        let ptr = unsafe { self.alloc(layout) };
+        if ptr.is_null() {
+            Err(std::alloc::AllocError)
+        } else {
+            unsafe {
+                Ok(core::ptr::NonNull::slice_from_raw_parts(
+                    core::ptr::NonNull::new_unchecked(ptr),
+                    layout.size(),
+                ))
+            }
+        }
+    }
+
+    unsafe fn deallocate(&self, ptr: core::ptr::NonNull<u8>, layout: Layout) {
+        log!("deallocate called on {:#x?} of size {}", ptr, layout.size());
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
 
+    #[cfg(feature = "nightly")]
     #[test]
     fn test_alloc() {
-        let allocator = BumpAllocator::new();
-        // create a new page aligned box of memory with 20 pages.
-        let mut buffer = Box::new([0; 0x1000 * 20]);
-        let base_addr = buffer.as_ptr() as usize;
-        // align up base_addr to the next page.
-        let base_addr = align_up(base_addr as usize, 0x1000) as u64;
-        assert_eq!(base_addr & 0xFFF, 0); // ensure page aligned
-
-        let buffer_range = MemoryRange::new(base_addr..(base_addr + 10 * 0x1000));
-
-        unsafe {
-            allocator.init(buffer_range);
-        }
+        let buffer: Box<[u8]> = Box::new([0; 0x1000 * 20]);
+        let addr = Box::into_raw(buffer) as *mut u8;
+        let allocator = BumpAllocator {
+            inner: SingleThreaded(RefCell::new(Inner {
+                start: addr,
+                next: addr,
+                end: unsafe { addr.add(0x1000 * 20) },
+                allow_alloc: false,
+                alloc_count: 0,
+            })),
+        };
         allocator.enable_alloc();
 
         unsafe {
@@ -163,6 +183,19 @@ mod tests {
             let ptr3 = allocator.alloc(Layout::from_size_align(300, 32).unwrap());
             *ptr3 = 77;
             assert_eq!(*ptr3, 77);
+
+            let mut vec: Vec<u8, BumpAllocator> = Vec::new_in(allocator);
+
+            // Push 4096 bytes, which should force a vec realloc.
+            for i in 0..4096 {
+                vec.push(i as u8);
+            }
+
+            // force an explicit resize to 10000 bytes
+            vec.resize(10000, 0);
         }
+
+        // Recreate the box, then drop it so miri is satisifed.
+        let _buf = unsafe { Box::from_raw(core::ptr::slice_from_raw_parts_mut(addr, 0x1000 * 20)) };
     }
 }
diff --git a/openhcl/openhcl_boot/src/main.rs b/openhcl/openhcl_boot/src/main.rs
index b562396720..ffdca6b2e0 100644
--- a/openhcl/openhcl_boot/src/main.rs
+++ b/openhcl/openhcl_boot/src/main.rs
@@ -8,6 +8,7 @@
 #![cfg_attr(minimal_rt, no_std, no_main)]
 // UNSAFETY: Interacting with low level hardware and bootloader primitives.
 #![expect(unsafe_code)]
+#![cfg_attr(feature = "nightly", feature(allocator_api))]
 
 mod arch;
 mod boot_logger;

From bf8140cfbb516cfef8d992b8cce1dfdde95dec5a Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Fri, 12 Sep 2025 15:03:22 -0700
Subject: [PATCH 06/10] remove nightly feature and use cfg, safety and cleanup

---
 openhcl/openhcl_boot/Cargo.toml        |  3 -
 openhcl/openhcl_boot/build.rs          |  3 +
 openhcl/openhcl_boot/src/bump_alloc.rs | 84 ++++++++++++++++++++------
 openhcl/openhcl_boot/src/main.rs       |  7 ++-
 4 files changed, 75 insertions(+), 22 deletions(-)

diff --git a/openhcl/openhcl_boot/Cargo.toml b/openhcl/openhcl_boot/Cargo.toml
index effb7aabeb..bd9feb399c 100644
--- a/openhcl/openhcl_boot/Cargo.toml
+++ b/openhcl/openhcl_boot/Cargo.toml
@@ -6,9 +6,6 @@ name = "openhcl_boot"
 edition.workspace = true
 rust-version.workspace = true
 
-[features]
-nightly = []
-
 [dependencies]
 aarch64defs.workspace = true
 minimal_rt.workspace = true
diff --git a/openhcl/openhcl_boot/build.rs b/openhcl/openhcl_boot/build.rs
index 98fa0f2482..35eba97e94 100644
--- a/openhcl/openhcl_boot/build.rs
+++ b/openhcl/openhcl_boot/build.rs
@@ -4,5 +4,8 @@
 #![expect(missing_docs)]
 
 fn main() {
+    // Allow a cfg of nightly to avoid using a feature, see main.rs.
+    println!("cargo:rustc-check-cfg=cfg(nightly)");
+
     minimal_rt_build::init();
 }
diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/bump_alloc.rs
index 06f314bccc..3783373427 100644
--- a/openhcl/openhcl_boot/src/bump_alloc.rs
+++ b/openhcl/openhcl_boot/src/bump_alloc.rs
@@ -1,8 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-//! Simple bump allocator using https://os.phil-opp.com/allocator-designs/ as a
-//! reference and starting point.
+//! A simple bump allocator that can be used in the bootloader.
 //!
 //! Note that we only allow allocations in a small window for supporting
 //! mesh_protobuf. Any other attempts to allocate will result in a panic.
@@ -14,15 +13,28 @@ use core::alloc::Layout;
 use core::cell::RefCell;
 use memory_range::MemoryRange;
 
+// Only enable the bump allocator when compiling with minimal_rt, as otherwise
+// it will override the global allocator in unit tests which is not what we
+// want.
 #[cfg_attr(minimal_rt, global_allocator)]
 pub static ALLOCATOR: BumpAllocator = BumpAllocator::new();
 
+#[derive(Debug, PartialEq, Eq)]
+enum State {
+    /// Allocations can be enabled via `enable_alloc`.
+    Allowed,
+    /// Allocations are currently enabled.
+    Enabled,
+    /// Allocations are disabled and cannot be enabled again.
+    Disabled,
+}
+
 #[derive(Debug)]
 pub struct Inner {
     start: *mut u8,
     next: *mut u8,
     end: *mut u8,
-    allow_alloc: bool,
+    allow_alloc: State,
     alloc_count: usize,
 }
 
@@ -37,7 +49,7 @@ impl BumpAllocator {
                 start: core::ptr::null_mut(),
                 next: core::ptr::null_mut(),
                 end: core::ptr::null_mut(),
-                allow_alloc: false,
+                allow_alloc: State::Allowed,
                 alloc_count: 0,
             })),
         }
@@ -62,22 +74,47 @@ impl BumpAllocator {
         inner.end = mem.end() as *mut u8;
     }
 
+    /// Enable allocations. This panics if allocations were ever previously
+    /// allocated.
     pub fn enable_alloc(&self) {
         let mut inner = self.inner.borrow_mut();
-        inner.allow_alloc = true;
+
+        inner.allow_alloc = match inner.allow_alloc {
+            State::Allowed => State::Enabled,
+            State::Enabled => {
+                panic!("allocations are already enabled");
+            }
+            State::Disabled => {
+                panic!("allocations were previously disabled and cannot be re-enabled");
+            }
+        };
     }
 
+    /// Disable allocations. Panis if the allocator was not previously enabled.
     pub fn disable_alloc(&self) {
         let mut inner = self.inner.borrow_mut();
-        inner.allow_alloc = false;
+        inner.allow_alloc = match inner.allow_alloc {
+            State::Allowed => panic!("allocations were never enabled"),
+            State::Enabled => State::Disabled,
+            State::Disabled => {
+                panic!("allocations were previously disabled and cannot be disabled again");
+            }
+        };
     }
 
     pub fn log_stats(&self) {
         let inner = self.inner.borrow();
 
-        // FIXME: unsafe calcs
-        let allocated = unsafe { inner.next.offset_from(inner.start) };
-        let free = unsafe { inner.end.offset_from(inner.next) };
+        // SAFETY: The pointers are within the same original allocation,
+        // specified by init. They are u8 pointers, so there is no alignment
+        // requirement.
+        let (allocated, free) = unsafe {
+            (
+                inner.next.offset_from(inner.start),
+                inner.end.offset_from(inner.next),
+            )
+        };
+
         log!(
             "Bump allocator: allocated {} bytes in {} allocations ({} bytes free)",
             allocated,
@@ -87,23 +124,26 @@ impl BumpAllocator {
     }
 }
 
-// SAFETY: The allocator points to a valid identity mapped memory range via
-// init, which is unused by anything else.
+// SAFETY: The allocator points to a valid identity VA range via the
+// construction at init.
 unsafe impl GlobalAlloc for BumpAllocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
         let mut inner = self.inner.borrow_mut();
 
-        if !inner.allow_alloc {
-            panic!("allocations are not allowed");
+        if inner.allow_alloc != State::Enabled {
+            panic!("allocations are not allowed {:?}", inner.allow_alloc);
         }
 
-        // FIXME: verify math and wraparounds
         let align_offset = inner.next.align_offset(layout.align());
         let alloc_start = inner.next.wrapping_add(align_offset);
         let alloc_end = alloc_start.wrapping_add(layout.size());
 
+        // If end overflowed this allocation is too large. If start overflowed,
+        // end will also overflow.
+        //
+        // Rust `Layout` guarnantees that the size is not larger than `isize`,
+        // so it's not possible to wrap around twice.
         if alloc_end < alloc_start {
-            // overflow
             return core::ptr::null_mut();
         }
 
@@ -128,9 +168,13 @@ unsafe impl GlobalAlloc for BumpAllocator {
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
         log!("dealloc called on {:#x?} of size {}", ptr, layout.size());
     }
+
+    // TODO: consider implementing realloc for the Vec grow case, which is the
+    // main usecase we see. This would mean supporting realloc if the allocation
+    // being realloced was the last one aka the tail.
 }
 
-#[cfg(feature = "nightly")]
+#[cfg(nightly)]
 unsafe impl core::alloc::Allocator for BumpAllocator {
     fn allocate(&self, layout: Layout) -> Result<core::ptr::NonNull<[u8]>, std::alloc::AllocError> {
         let ptr = unsafe { self.alloc(layout) };
@@ -151,11 +195,13 @@ unsafe impl core::alloc::Allocator for BumpAllocator {
     }
 }
 
+#[cfg(nightly)]
 #[cfg(test)]
 mod tests {
     use super::*;
 
-    #[cfg(feature = "nightly")]
+    // NOTE: run these tests with miri via
+    // `RUSTFLAGS="--cfg nightly" cargo +nightly miri test -p openhcl_boot`
     #[test]
     fn test_alloc() {
         let buffer: Box<[u8]> = Box::new([0; 0x1000 * 20]);
@@ -165,7 +211,7 @@ mod tests {
                 start: addr,
                 next: addr,
                 end: unsafe { addr.add(0x1000 * 20) },
-                allow_alloc: false,
+                allow_alloc: State::Allowed,
                 alloc_count: 0,
             })),
         };
@@ -183,7 +229,9 @@ mod tests {
             let ptr3 = allocator.alloc(Layout::from_size_align(300, 32).unwrap());
             *ptr3 = 77;
             assert_eq!(*ptr3, 77);
+        }
 
+        {
             let mut vec: Vec<u8, BumpAllocator> = Vec::new_in(allocator);
 
             // Push 4096 bytes, which should force a vec realloc.
diff --git a/openhcl/openhcl_boot/src/main.rs b/openhcl/openhcl_boot/src/main.rs
index ffdca6b2e0..69b0c06ee5 100644
--- a/openhcl/openhcl_boot/src/main.rs
+++ b/openhcl/openhcl_boot/src/main.rs
@@ -8,7 +8,12 @@
 #![cfg_attr(minimal_rt, no_std, no_main)]
 // UNSAFETY: Interacting with low level hardware and bootloader primitives.
 #![expect(unsafe_code)]
-#![cfg_attr(feature = "nightly", feature(allocator_api))]
+// Allow the allocator api when compiling with `RUSTFLAGS="--cfg nightly"`.
+//
+// Do not use a normal feature, as that shows errors rust-analyzer since most
+// people are using stable. We could remove this once the allocator_api feature
+// is stable.
+#![cfg_attr(nightly, feature(allocator_api))]
 
 mod arch;
 mod boot_logger;

From e551e14db87bb27cde2133af26c72ef623f39fad Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Mon, 22 Sep 2025 10:04:28 -0700
Subject: [PATCH 07/10] cleanup

---
 openhcl/openhcl_boot/src/bump_alloc.rs | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/bump_alloc.rs
index 3783373427..07703977a6 100644
--- a/openhcl/openhcl_boot/src/bump_alloc.rs
+++ b/openhcl/openhcl_boot/src/bump_alloc.rs
@@ -90,7 +90,7 @@ impl BumpAllocator {
         };
     }
 
-    /// Disable allocations. Panis if the allocator was not previously enabled.
+    /// Disable allocations. Panics if the allocator was not previously enabled.
     pub fn disable_alloc(&self) {
         let mut inner = self.inner.borrow_mut();
         inner.allow_alloc = match inner.allow_alloc {
@@ -141,7 +141,7 @@ unsafe impl GlobalAlloc for BumpAllocator {
         // If end overflowed this allocation is too large. If start overflowed,
         // end will also overflow.
         //
-        // Rust `Layout` guarnantees that the size is not larger than `isize`,
+        // Rust `Layout` guarantees that the size is not larger than `isize`,
         // so it's not possible to wrap around twice.
         if alloc_end < alloc_start {
             return core::ptr::null_mut();
@@ -176,10 +176,13 @@ unsafe impl GlobalAlloc for BumpAllocator {
 
 #[cfg(nightly)]
 unsafe impl core::alloc::Allocator for BumpAllocator {
-    fn allocate(&self, layout: Layout) -> Result<core::ptr::NonNull<[u8]>, std::alloc::AllocError> {
+    fn allocate(
+        &self,
+        layout: Layout,
+    ) -> Result<core::ptr::NonNull<[u8]>, core::alloc::AllocError> {
         let ptr = unsafe { self.alloc(layout) };
         if ptr.is_null() {
-            Err(std::alloc::AllocError)
+            Err(core::alloc::AllocError)
         } else {
             unsafe {
                 Ok(core::ptr::NonNull::slice_from_raw_parts(
@@ -243,7 +246,7 @@ mod tests {
             vec.resize(10000, 0);
         }
 
-        // Recreate the box, then drop it so miri is satisifed.
+        // Recreate the box, then drop it so miri is satisfied.
         let _buf = unsafe { Box::from_raw(core::ptr::slice_from_raw_parts_mut(addr, 0x1000 * 20)) };
     }
 }

From a204b40eb29ee7f048d6e069d54eec1ffcf4792e Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Mon, 22 Sep 2025 11:38:29 -0700
Subject: [PATCH 08/10] more comment feedback

---
 openhcl/openhcl_boot/src/bump_alloc.rs | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/bump_alloc.rs
index 07703977a6..2800d06645 100644
--- a/openhcl/openhcl_boot/src/bump_alloc.rs
+++ b/openhcl/openhcl_boot/src/bump_alloc.rs
@@ -75,7 +75,7 @@ impl BumpAllocator {
     }
 
     /// Enable allocations. This panics if allocations were ever previously
-    /// allocated.
+    /// enabled.
     pub fn enable_alloc(&self) {
         let mut inner = self.inner.borrow_mut();
 
@@ -147,14 +147,8 @@ unsafe impl GlobalAlloc for BumpAllocator {
             return core::ptr::null_mut();
         }
 
-        log!(
-            "bump_alloc: allocating {} bytes with alignment {} at with offset {} alloc_start {:#x?}, alloc_end {:#x?}",
-            layout.size(),
-            layout.align(),
-            align_offset,
-            alloc_start,
-            alloc_end,
-        );
+        // TODO: renable allocation tracing when we support tracing levels via
+        // the log crate.
 
         if alloc_end > inner.end {
             core::ptr::null_mut() // out of memory
@@ -165,8 +159,9 @@ unsafe impl GlobalAlloc for BumpAllocator {
         }
     }
 
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        log!("dealloc called on {:#x?} of size {}", ptr, layout.size());
+    unsafe fn dealloc(&self, _ptr: *mut u8, _layout: Layout) {
+        // TODO: renable allocation tracing when we support tracing levels via
+        // the log crate.
     }
 
     // TODO: consider implementing realloc for the Vec grow case, which is the
@@ -175,7 +170,9 @@ unsafe impl GlobalAlloc for BumpAllocator {
 }
 
 #[cfg(nightly)]
-unsafe impl core::alloc::Allocator for BumpAllocator {
+// SAFETY: The allocator points to a valid identity VA range via the
+// construction at init, the same as for `GlobalAlloc`.
+unsafe impl core::alloc::Allocator for &BumpAllocator {
     fn allocate(
         &self,
         layout: Layout,
@@ -235,7 +232,7 @@ mod tests {
         }
 
         {
-            let mut vec: Vec<u8, BumpAllocator> = Vec::new_in(allocator);
+            let mut vec: Vec<u8, &BumpAllocator> = Vec::new_in(&allocator);
 
             // Push 4096 bytes, which should force a vec realloc.
             for i in 0..4096 {
@@ -248,5 +245,7 @@ mod tests {
 
         // Recreate the box, then drop it so miri is satisfied.
         let _buf = unsafe { Box::from_raw(core::ptr::slice_from_raw_parts_mut(addr, 0x1000 * 20)) };
+
+        allocator.log_stats();
     }
 }

From ce495a54004e93bf07a05caffd42885a141bc12d Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Mon, 22 Sep 2025 11:44:01 -0700
Subject: [PATCH 09/10] move allocator to just inside dt

---
 .../src/{ => host_params/dt}/bump_alloc.rs             |  0
 .../openhcl_boot/src/host_params/{dt.rs => dt/mod.rs}  | 10 ++++++----
 openhcl/openhcl_boot/src/main.rs                       | 10 +++++-----
 3 files changed, 11 insertions(+), 9 deletions(-)
 rename openhcl/openhcl_boot/src/{ => host_params/dt}/bump_alloc.rs (100%)
 rename openhcl/openhcl_boot/src/host_params/{dt.rs => dt/mod.rs} (99%)

diff --git a/openhcl/openhcl_boot/src/bump_alloc.rs b/openhcl/openhcl_boot/src/host_params/dt/bump_alloc.rs
similarity index 100%
rename from openhcl/openhcl_boot/src/bump_alloc.rs
rename to openhcl/openhcl_boot/src/host_params/dt/bump_alloc.rs
diff --git a/openhcl/openhcl_boot/src/host_params/dt.rs b/openhcl/openhcl_boot/src/host_params/dt/mod.rs
similarity index 99%
rename from openhcl/openhcl_boot/src/host_params/dt.rs
rename to openhcl/openhcl_boot/src/host_params/dt/mod.rs
index f237d84027..062e52f7a0 100644
--- a/openhcl/openhcl_boot/src/host_params/dt.rs
+++ b/openhcl/openhcl_boot/src/host_params/dt/mod.rs
@@ -34,6 +34,8 @@ use memory_range::subtract_ranges;
 use memory_range::walk_ranges;
 use thiserror::Error;
 
+mod bump_alloc;
+
 /// Errors when reading the host device tree.
 #[derive(Debug, Error)]
 pub enum DtError {
@@ -320,20 +322,20 @@ fn init_heap(params: &ShimParams) {
     // SAFETY: The heap range is reserved at file build time, and is
     // guaranteed to be unused by anything else.
     unsafe {
-        crate::bump_alloc::ALLOCATOR.init(params.heap);
+        bump_alloc::ALLOCATOR.init(params.heap);
     }
 
     // TODO: test using heap, as no mesh decode yet.
     #[cfg(debug_assertions)]
     {
         use alloc::boxed::Box;
-        crate::bump_alloc::ALLOCATOR.enable_alloc();
+        bump_alloc::ALLOCATOR.enable_alloc();
 
         let box_int = Box::new(42);
         log!("box int {box_int}");
         drop(box_int);
-        crate::bump_alloc::ALLOCATOR.disable_alloc();
-        crate::bump_alloc::ALLOCATOR.log_stats();
+        bump_alloc::ALLOCATOR.disable_alloc();
+        bump_alloc::ALLOCATOR.log_stats();
     }
 }
 
diff --git a/openhcl/openhcl_boot/src/main.rs b/openhcl/openhcl_boot/src/main.rs
index 69b0c06ee5..8ac3dd6c93 100644
--- a/openhcl/openhcl_boot/src/main.rs
+++ b/openhcl/openhcl_boot/src/main.rs
@@ -8,16 +8,16 @@
 #![cfg_attr(minimal_rt, no_std, no_main)]
 // UNSAFETY: Interacting with low level hardware and bootloader primitives.
 #![expect(unsafe_code)]
-// Allow the allocator api when compiling with `RUSTFLAGS="--cfg nightly"`.
+// Allow the allocator api when compiling with `RUSTFLAGS="--cfg nightly"`. This
+// is used for some miri tests for testing the bump allocator.
 //
-// Do not use a normal feature, as that shows errors rust-analyzer since most
-// people are using stable. We could remove this once the allocator_api feature
-// is stable.
+// Do not use a normal feature, as that shows errors with rust-analyzer since
+// most people are using stable and enable all features. We could remove this
+// once the allocator_api feature is stable.
 #![cfg_attr(nightly, feature(allocator_api))]
 
 mod arch;
 mod boot_logger;
-mod bump_alloc;
 mod cmdline;
 mod dt;
 mod host_params;

From b9799e668dabb9f002f2eeca76ab75bd71f45c2a Mon Sep 17 00:00:00 2001
From: Chris Oo <cho@microsoft.com>
Date: Thu, 25 Sep 2025 15:03:57 -0700
Subject: [PATCH 10/10] add to miri test exhausted allocation

---
 openhcl/openhcl_boot/src/host_params/dt/bump_alloc.rs | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/openhcl/openhcl_boot/src/host_params/dt/bump_alloc.rs b/openhcl/openhcl_boot/src/host_params/dt/bump_alloc.rs
index 2800d06645..2009e9e5d3 100644
--- a/openhcl/openhcl_boot/src/host_params/dt/bump_alloc.rs
+++ b/openhcl/openhcl_boot/src/host_params/dt/bump_alloc.rs
@@ -243,6 +243,12 @@ mod tests {
             vec.resize(10000, 0);
         }
 
+        // Attempt to allocate a large chunk that is not available.
+        unsafe {
+            let ptr4 = allocator.alloc(Layout::from_size_align(0x1000 * 20, 8).unwrap());
+            assert!(ptr4.is_null());
+        }
+
         // Recreate the box, then drop it so miri is satisfied.
         let _buf = unsafe { Box::from_raw(core::ptr::slice_from_raw_parts_mut(addr, 0x1000 * 20)) };