pvmfw: Move MemorySlices to crate::memory

Note: No functional change intended.

Test: m pvmfw_bin
Bug: 377276983
Change-Id: I443a20a0436b912a9335ddca510cf1f73737e035
diff --git a/guest/pvmfw/src/entry.rs b/guest/pvmfw/src/entry.rs
index 945ad6a..dc6b5fa 100644
--- a/guest/pvmfw/src/entry.rs
+++ b/guest/pvmfw/src/entry.rs
@@ -15,14 +15,11 @@
 //! Low-level entry and exit points of pvmfw.
 
 use crate::config;
-use crate::fdt;
 use crate::memory;
 use core::arch::asm;
 use core::mem::{drop, size_of};
-use core::num::NonZeroUsize;
 use core::ops::Range;
 use core::slice;
-use log::debug;
 use log::error;
 use log::info;
 use log::warn;
@@ -30,7 +27,7 @@
 use vmbase::util::RangeExt as _;
 use vmbase::{
     configure_heap, console_writeln,
-    hyp::{get_mem_sharer, get_mmio_guard},
+    hyp::get_mmio_guard,
     layout::{self, crosvm, UART_PAGE_ADDR},
     main,
     memory::{min_dcache_line_size, MemoryTracker, MEMORY, SIZE_128KB, SIZE_4KB},
@@ -94,112 +91,6 @@
     // if we reach this point and return, vmbase::entry::rust_entry() will call power::shutdown().
 }
 
-struct MemorySlices<'a> {
-    fdt: &'a mut libfdt::Fdt,
-    kernel: &'a [u8],
-    ramdisk: Option<&'a [u8]>,
-}
-
-impl<'a> MemorySlices<'a> {
-    fn new(
-        fdt: usize,
-        kernel: usize,
-        kernel_size: usize,
-        vm_dtbo: Option<&mut [u8]>,
-        vm_ref_dt: Option<&[u8]>,
-    ) -> Result<Self, RebootReason> {
-        let fdt_size = NonZeroUsize::new(crosvm::FDT_MAX_SIZE).unwrap();
-        // TODO - Only map the FDT as read-only, until we modify it right before jump_to_payload()
-        // e.g. by generating a DTBO for a template DT in main() and, on return, re-map DT as RW,
-        // overwrite with the template DT and apply the DTBO.
-        let range = MEMORY.lock().as_mut().unwrap().alloc_mut(fdt, fdt_size).map_err(|e| {
-            error!("Failed to allocate the FDT range: {e}");
-            RebootReason::InternalError
-        })?;
-
-        // SAFETY: The tracker validated the range to be in main memory, mapped, and not overlap.
-        let fdt = unsafe { slice::from_raw_parts_mut(range.start as *mut u8, range.len()) };
-
-        let info = fdt::sanitize_device_tree(fdt, vm_dtbo, vm_ref_dt)?;
-        let fdt = libfdt::Fdt::from_mut_slice(fdt).map_err(|e| {
-            error!("Failed to load sanitized FDT: {e}");
-            RebootReason::InvalidFdt
-        })?;
-        debug!("Fdt passed validation!");
-
-        let memory_range = info.memory_range;
-        debug!("Resizing MemoryTracker to range {memory_range:#x?}");
-        MEMORY.lock().as_mut().unwrap().shrink(&memory_range).map_err(|e| {
-            error!("Failed to use memory range value from DT: {memory_range:#x?}: {e}");
-            RebootReason::InvalidFdt
-        })?;
-
-        if let Some(mem_sharer) = get_mem_sharer() {
-            let granule = mem_sharer.granule().map_err(|e| {
-                error!("Failed to get memory protection granule: {e}");
-                RebootReason::InternalError
-            })?;
-            MEMORY.lock().as_mut().unwrap().init_dynamic_shared_pool(granule).map_err(|e| {
-                error!("Failed to initialize dynamically shared pool: {e}");
-                RebootReason::InternalError
-            })?;
-        } else {
-            let range = info.swiotlb_info.fixed_range().ok_or_else(|| {
-                error!("Pre-shared pool range not specified in swiotlb node");
-                RebootReason::InvalidFdt
-            })?;
-
-            MEMORY.lock().as_mut().unwrap().init_static_shared_pool(range).map_err(|e| {
-                error!("Failed to initialize pre-shared pool {e}");
-                RebootReason::InvalidFdt
-            })?;
-        }
-
-        let kernel_range = if let Some(r) = info.kernel_range {
-            MEMORY.lock().as_mut().unwrap().alloc_range(&r).map_err(|e| {
-                error!("Failed to obtain the kernel range with DT range: {e}");
-                RebootReason::InternalError
-            })?
-        } else if cfg!(feature = "legacy") {
-            warn!("Failed to find the kernel range in the DT; falling back to legacy ABI");
-
-            let kernel_size = NonZeroUsize::new(kernel_size).ok_or_else(|| {
-                error!("Invalid kernel size: {kernel_size:#x}");
-                RebootReason::InvalidPayload
-            })?;
-
-            MEMORY.lock().as_mut().unwrap().alloc(kernel, kernel_size).map_err(|e| {
-                error!("Failed to obtain the kernel range with legacy range: {e}");
-                RebootReason::InternalError
-            })?
-        } else {
-            error!("Failed to locate the kernel from the DT");
-            return Err(RebootReason::InvalidPayload);
-        };
-
-        let kernel = kernel_range.start as *const u8;
-        // SAFETY: The tracker validated the range to be in main memory, mapped, and not overlap.
-        let kernel = unsafe { slice::from_raw_parts(kernel, kernel_range.len()) };
-
-        let ramdisk = if let Some(r) = info.initrd_range {
-            debug!("Located ramdisk at {r:?}");
-            let r = MEMORY.lock().as_mut().unwrap().alloc_range(&r).map_err(|e| {
-                error!("Failed to obtain the initrd range: {e}");
-                RebootReason::InvalidRamdisk
-            })?;
-
-            // SAFETY: The region was validated by memory to be in main memory, mapped, and
-            // not overlap.
-            Some(unsafe { slice::from_raw_parts(r.start as *const u8, r.len()) })
-        } else {
-            info!("Couldn't locate the ramdisk from the device tree");
-            None
-        };
-
-        Ok(Self { fdt, kernel, ramdisk })
-    }
-}
-
 /// Sets up the environment for main() and wraps its result for start().
 ///
 /// Provide the abstractions necessary for start() to abort the pVM boot and for main() to run with
@@ -240,7 +131,7 @@
         Some(memory::appended_payload_range()),
     ));
 
-    let slices = MemorySlices::new(
+    let slices = memory::MemorySlices::new(
         fdt,
         payload,
         payload_size,
diff --git a/guest/pvmfw/src/memory.rs b/guest/pvmfw/src/memory.rs
index 8d12b57..f49d79b 100644
--- a/guest/pvmfw/src/memory.rs
+++ b/guest/pvmfw/src/memory.rs
@@ -14,15 +14,23 @@
 
 //! Low-level allocation and tracking of main memory.
 
+use crate::entry::RebootReason;
+use crate::fdt;
 use crate::helpers::PVMFW_PAGE_SIZE;
 use aarch64_paging::paging::VirtualAddress;
 use aarch64_paging::MapError;
+use core::num::NonZeroUsize;
 use core::ops::Range;
 use core::result;
+use core::slice;
+use log::debug;
 use log::error;
+use log::info;
+use log::warn;
 use vmbase::{
-    layout,
-    memory::{PageTable, SIZE_2MB, SIZE_4KB},
+    hyp::get_mem_sharer,
+    layout::{self, crosvm},
+    memory::{PageTable, MEMORY, SIZE_2MB, SIZE_4KB},
     util::align_up,
 };
 
@@ -57,3 +65,109 @@
     }
     Ok(page_table)
 }
+
+pub(crate) struct MemorySlices<'a> {
+    pub fdt: &'a mut libfdt::Fdt,
+    pub kernel: &'a [u8],
+    pub ramdisk: Option<&'a [u8]>,
+}
+
+impl<'a> MemorySlices<'a> {
+    pub fn new(
+        fdt: usize,
+        kernel: usize,
+        kernel_size: usize,
+        vm_dtbo: Option<&mut [u8]>,
+        vm_ref_dt: Option<&[u8]>,
+    ) -> Result<Self, RebootReason> {
+        let fdt_size = NonZeroUsize::new(crosvm::FDT_MAX_SIZE).unwrap();
+        // TODO - Only map the FDT as read-only, until we modify it right before jump_to_payload()
+        // e.g. by generating a DTBO for a template DT in main() and, on return, re-map DT as RW,
+        // overwrite with the template DT and apply the DTBO.
+        let range = MEMORY.lock().as_mut().unwrap().alloc_mut(fdt, fdt_size).map_err(|e| {
+            error!("Failed to allocate the FDT range: {e}");
+            RebootReason::InternalError
+        })?;
+
+        // SAFETY: The tracker validated the range to be in main memory, mapped, and not overlap.
+        let fdt = unsafe { slice::from_raw_parts_mut(range.start as *mut u8, range.len()) };
+
+        let info = fdt::sanitize_device_tree(fdt, vm_dtbo, vm_ref_dt)?;
+        let fdt = libfdt::Fdt::from_mut_slice(fdt).map_err(|e| {
+            error!("Failed to load sanitized FDT: {e}");
+            RebootReason::InvalidFdt
+        })?;
+        debug!("Fdt passed validation!");
+
+        let memory_range = info.memory_range;
+        debug!("Resizing MemoryTracker to range {memory_range:#x?}");
+        MEMORY.lock().as_mut().unwrap().shrink(&memory_range).map_err(|e| {
+            error!("Failed to use memory range value from DT: {memory_range:#x?}: {e}");
+            RebootReason::InvalidFdt
+        })?;
+
+        if let Some(mem_sharer) = get_mem_sharer() {
+            let granule = mem_sharer.granule().map_err(|e| {
+                error!("Failed to get memory protection granule: {e}");
+                RebootReason::InternalError
+            })?;
+            MEMORY.lock().as_mut().unwrap().init_dynamic_shared_pool(granule).map_err(|e| {
+                error!("Failed to initialize dynamically shared pool: {e}");
+                RebootReason::InternalError
+            })?;
+        } else {
+            let range = info.swiotlb_info.fixed_range().ok_or_else(|| {
+                error!("Pre-shared pool range not specified in swiotlb node");
+                RebootReason::InvalidFdt
+            })?;
+
+            MEMORY.lock().as_mut().unwrap().init_static_shared_pool(range).map_err(|e| {
+                error!("Failed to initialize pre-shared pool {e}");
+                RebootReason::InvalidFdt
+            })?;
+        }
+
+        let kernel_range = if let Some(r) = info.kernel_range {
+            MEMORY.lock().as_mut().unwrap().alloc_range(&r).map_err(|e| {
+                error!("Failed to obtain the kernel range with DT range: {e}");
+                RebootReason::InternalError
+            })?
+        } else if cfg!(feature = "legacy") {
+            warn!("Failed to find the kernel range in the DT; falling back to legacy ABI");
+
+            let kernel_size = NonZeroUsize::new(kernel_size).ok_or_else(|| {
+                error!("Invalid kernel size: {kernel_size:#x}");
+                RebootReason::InvalidPayload
+            })?;
+
+            MEMORY.lock().as_mut().unwrap().alloc(kernel, kernel_size).map_err(|e| {
+                error!("Failed to obtain the kernel range with legacy range: {e}");
+                RebootReason::InternalError
+            })?
+        } else {
+            error!("Failed to locate the kernel from the DT");
+            return Err(RebootReason::InvalidPayload);
+        };
+
+        let kernel = kernel_range.start as *const u8;
+        // SAFETY: The tracker validated the range to be in main memory, mapped, and not overlap.
+        let kernel = unsafe { slice::from_raw_parts(kernel, kernel_range.len()) };
+
+        let ramdisk = if let Some(r) = info.initrd_range {
+            debug!("Located ramdisk at {r:?}");
+            let r = MEMORY.lock().as_mut().unwrap().alloc_range(&r).map_err(|e| {
+                error!("Failed to obtain the initrd range: {e}");
+                RebootReason::InvalidRamdisk
+            })?;
+
+            // SAFETY: The region was validated by memory to be in main memory, mapped, and
+            // not overlap.
+            Some(unsafe { slice::from_raw_parts(r.start as *const u8, r.len()) })
+        } else {
+            info!("Couldn't locate the ramdisk from the device tree");
+            None
+        };
+
+        Ok(Self { fdt, kernel, ramdisk })
+    }
+}