[pvmfw] Refactor MMIO range check in pvmfw

This cl refactors the MMIO range check and moves the
pvmfw-specific memory region constants out of MemoryTracker. This
simplifies the task of moving MemoryTracker to vmbase later.

Bug: 284462758
Test: m pvmfw_img
Change-Id: Ifb563215c59d7bb9dd4fae28bba68054b48b30d2
diff --git a/pvmfw/src/entry.rs b/pvmfw/src/entry.rs
index 19ba9c3..66e994b 100644
--- a/pvmfw/src/entry.rs
+++ b/pvmfw/src/entry.rs
@@ -33,7 +33,9 @@
 use log::LevelFilter;
 use vmbase::util::RangeExt as _;
 use vmbase::{
-    console, layout, logger, main,
+    console,
+    layout::{self, crosvm},
+    logger, main,
     memory::{min_dcache_line_size, SIZE_2MB, SIZE_4KB},
     power::reboot,
 };
@@ -223,7 +225,11 @@
     let (bcc_slice, debug_policy) = appended.get_entries();
 
     // Up to this point, we were using the built-in static (from .rodata) page tables.
-    MEMORY.lock().replace(MemoryTracker::new(page_table));
+    MEMORY.lock().replace(MemoryTracker::new(
+        page_table,
+        crosvm::MEM_START..memory::MAX_ADDR,
+        crosvm::MMIO_START..crosvm::MMIO_END,
+    ));
 
     let slices = MemorySlices::new(fdt, payload, payload_size)?;
 
diff --git a/pvmfw/src/memory.rs b/pvmfw/src/memory.rs
index 216a621..ffbb1ca 100644
--- a/pvmfw/src/memory.rs
+++ b/pvmfw/src/memory.rs
@@ -37,11 +37,10 @@
 use spin::mutex::SpinMutex;
 use tinyvec::ArrayVec;
 use vmbase::{
-    dsb, isb,
-    layout::{self, crosvm::MEM_START},
+    dsb, isb, layout,
     memory::{
         flush_dirty_range, is_leaf_pte, page_4kb_of, set_dbm_enabled, MemorySharer, PageTable,
-        MMIO_LAZY_MAP_FLAG, SIZE_2MB, SIZE_4KB, SIZE_4MB,
+        MMIO_LAZY_MAP_FLAG, SIZE_2MB, SIZE_4KB,
     },
     tlbi,
     util::{align_up, RangeExt as _},
@@ -77,6 +76,7 @@
     page_table: PageTable,
     regions: ArrayVec<[MemoryRegion; MemoryTracker::CAPACITY]>,
     mmio_regions: ArrayVec<[MemoryRange; MemoryTracker::MMIO_CAPACITY]>,
+    mmio_range: MemoryRange,
 }
 
 /// Errors for MemoryTracker operations.
@@ -147,10 +147,14 @@
 impl MemoryTracker {
     const CAPACITY: usize = 5;
     const MMIO_CAPACITY: usize = 5;
-    const PVMFW_RANGE: MemoryRange = (MEM_START - SIZE_4MB)..MEM_START;
 
     /// Create a new instance from an active page table, covering the maximum RAM size.
-    pub fn new(mut page_table: PageTable) -> Self {
+    pub fn new(mut page_table: PageTable, total: MemoryRange, mmio_range: MemoryRange) -> Self {
+        assert!(
+            !total.overlaps(&mmio_range),
+            "MMIO space should not overlap with the main memory region."
+        );
+
         // Activate dirty state management first, otherwise we may get permission faults immediately
         // after activating the new page table. This has no effect before the new page table is
         // activated because none of the entries in the initial idmap have the DBM flag.
@@ -163,10 +167,11 @@
         debug!("... Success!");
 
         Self {
-            total: MEM_START..MAX_ADDR,
+            total,
             page_table,
             regions: ArrayVec::new(),
             mmio_regions: ArrayVec::new(),
+            mmio_range,
         }
     }
 
@@ -223,8 +228,7 @@
     /// Checks that the given range of addresses is within the MMIO region, and then maps it
     /// appropriately.
     pub fn map_mmio_range(&mut self, range: MemoryRange) -> Result<()> {
-        // MMIO space is below the main memory region.
-        if range.end > self.total.start || range.overlaps(&Self::PVMFW_RANGE) {
+        if !range.is_within(&self.mmio_range) {
             return Err(MemoryTrackerError::OutOfRange);
         }
         if self.mmio_regions.iter().any(|r| range.overlaps(r)) {
diff --git a/vmbase/src/layout/crosvm.rs b/vmbase/src/layout/crosvm.rs
index 66b9e44..fbeb477 100644
--- a/vmbase/src/layout/crosvm.rs
+++ b/vmbase/src/layout/crosvm.rs
@@ -16,5 +16,10 @@
 //!
 //! https://crosvm.dev/book/appendix/memory_layout.html#common-layout
 
+/// The start address of MMIO space.
+pub const MMIO_START: usize = 0x0;
+/// The end address of MMIO space.
+pub const MMIO_END: usize = 0x4000_0000;
+
 /// The start of the system's contiguous "main" memory.
 pub const MEM_START: usize = 0x8000_0000;