[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/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)) {