junyu0312 · junyu0312 · Feb 28, 2026 · Feb 28, 2026
diff --git a/crates/vm-core/src/virt.rs b/crates/vm-core/src/virt.rs
@@ -7,7 +7,6 @@ use vm_mm::manager::MemoryAddressSpace;
 use crate::arch::Arch;
 use crate::arch::irq::InterruptController;
 use crate::arch::vcpu::Vcpu;
-use crate::device::mmio::MmioLayout;
 use crate::device::vm_exit::DeviceVmExitHandler;
 use crate::error::Result;
 
@@ -27,9 +26,9 @@ pub trait Virt: Sized {
     fn init_irq(&mut self) -> Result<Arc<dyn InterruptController>>;
     fn init_memory(
         &mut self,
-        mmio_layout: &MmioLayout,
         memory: &mut MemoryAddressSpace<Self::Memory>,
-        memory_size: u64,
+        ram_base: u64,
+        memory_size: usize,
     ) -> Result<()>;
     fn post_init(&mut self) -> Result<()>;
 

diff --git a/crates/vm-core/src/virt/hvp.rs b/crates/vm-core/src/virt/hvp.rs
@@ -11,7 +11,9 @@ use applevisor::vm::GicEnabled;
 use applevisor::vm::VirtualMachine;
 use applevisor::vm::VirtualMachineConfig;
 use applevisor::vm::VirtualMachineInstance;
+use vm_mm::allocator::Allocator;
 use vm_mm::manager::MemoryAddressSpace;
+use vm_mm::region::MemoryRegion;
 
 use crate::arch::Arch;
 use crate::arch::aarch64::AArch64;
@@ -27,7 +29,6 @@ use crate::arch::aarch64::vm_exit::handle_vm_exit;
 use crate::arch::irq::InterruptController;
 use crate::arch::layout::MemoryLayout;
 use crate::arch::vcpu::Vcpu;
-use crate::device::mmio::MmioLayout;
 use crate::device::vm_exit::DeviceVmExitHandler;
 use crate::error::Error;
 use crate::error::Result;
@@ -272,20 +273,19 @@ impl Virt for Hvp {
 
     fn init_memory(
         &mut self,
-        _mmio_layout: &MmioLayout,
-        memory: &mut MemoryAddressSpace<MemoryWrapper>,
-        memory_size: u64,
+        memory_address_space: &mut MemoryAddressSpace<MemoryWrapper>,
+        ram_base: u64,
+        memory_size: usize,
     ) -> Result<()> {
         let allocator = HvpAllocator { vm: &self.vm };
 
-        for region in memory {
-            region.alloc(&allocator)?;
+        let mut memory = allocator.alloc(memory_size, None)?;
+        memory.0.map(ram_base, MemPerms::ReadWriteExec)?;
+        memory_address_space
+            .try_insert(MemoryRegion::new(ram_base, memory_size, memory))
+            .map_err(|_| Error::FailedInitialize("Failed to initialize memory".to_string()))?;
 
-            let memory = region.memory.get_mut().unwrap();
-            memory.0.map(region.gpa, MemPerms::ReadWriteExec)?;
-        }
-
-        self.get_layout_mut().set_ram_size(memory_size)?;
+        self.get_layout_mut().set_ram_size(memory_size as u64)?;
 
         Ok(())
     }

diff --git a/crates/vm-core/src/virt/kvm.rs b/crates/vm-core/src/virt/kvm.rs
@@ -3,15 +3,13 @@ use std::marker::PhantomData;
 use std::sync::Arc;
 use std::sync::Mutex;
 
-use kvm_bindings::*;
 use kvm_ioctls::*;
 use memmap2::MmapMut;
 use vm_mm::allocator::mmap_allocator::MmapAllocator;
 use vm_mm::manager::MemoryAddressSpace;
 
 use crate::arch::Arch;
 use crate::arch::irq::InterruptController;
-use crate::device::mmio::MmioLayout;
 use crate::device::vm_exit::DeviceVmExitHandler;
 use crate::error::Error;
 use crate::error::Result;
@@ -69,28 +67,13 @@ where
 
     fn init_memory(
         &mut self,
-        _mmio_layout: &MmioLayout,
-        memory: &mut MemoryAddressSpace<MmapMut>,
-        _memory_size: u64,
+        _memory: &mut MemoryAddressSpace<MmapMut>,
+        _ram_base: u64,
+        _memory_size: usize,
     ) -> Result<()> {
-        let allocator = MmapAllocator;
-
-        for (slot, region) in memory.into_iter().enumerate() {
-            region.alloc(&allocator)?;
-
-            unsafe {
-                self.vm_fd
-                    .set_user_memory_region(kvm_userspace_memory_region {
-                        slot: slot as u32,
-                        flags: 0,
-                        guest_phys_addr: region.gpa,
-                        memory_size: region.len as u64,
-                        userspace_addr: region.to_hva().unwrap() as u64,
-                    })?;
-            }
-        }
+        let _allocator = MmapAllocator;
 
-        Ok(())
+        todo!()
     }
 
     fn post_init(&mut self) -> Result<()> {

diff --git a/crates/vm-core/src/virt/kvm/vcpu/arch/x86_64.rs b/crates/vm-core/src/virt/kvm/vcpu/arch/x86_64.rs
@@ -4,8 +4,8 @@ use kvm_ioctls::Kvm;
 use crate::arch::x86_64::X86_64;
 use crate::arch::x86_64::vcpu::X86Vcpu;
 use crate::arch::x86_64::vm_exit::VmExitReason;
+use crate::device::mmio::MmioLayout;
 use crate::error::Result;
-use crate::virt::MmioLayout;
 use crate::virt::Vcpu;
 use crate::virt::kvm::vcpu::KvmVcpu;
 

diff --git a/crates/vm-machine/src/vm_builder.rs b/crates/vm-machine/src/vm_builder.rs
@@ -7,9 +7,7 @@ use vm_core::device::device_manager::DeviceManager;
 use vm_core::device::mmio::MmioLayout;
 use vm_core::virt::Virt;
 use vm_device::device::Device;
-use vm_mm::allocator::MemoryContainer;
 use vm_mm::manager::MemoryAddressSpace;
-use vm_mm::region::MemoryRegion;
 
 use crate::device::InitDevice;
 use crate::error::Error;
@@ -53,9 +51,9 @@ impl VmBuilder {
             None
         };
 
-        let mut memory = self.init_mm(layout.get_ram_base())?;
-        virt.init_memory(&mmio_layout, &mut memory, self.memory_size as u64)?;
-        let memory = Arc::new(Mutex::new(memory));
+        let mut memory_regions = MemoryAddressSpace::default();
+        virt.init_memory(&mut memory_regions, layout.get_ram_base(), self.memory_size)?;
+        let memory = Arc::new(Mutex::new(memory_regions));
 
         virt.post_init()?;
 
@@ -73,18 +71,4 @@ impl VmBuilder {
 
         Ok(vm)
     }
-
-    fn init_mm<C>(&self, ram_base: u64) -> Result<MemoryAddressSpace<C>>
-    where
-        C: MemoryContainer,
-    {
-        let memory_region = MemoryRegion::placeholder(ram_base, self.memory_size);
-
-        let mut memory_regions = MemoryAddressSpace::default();
-        memory_regions
-            .try_insert(memory_region)
-            .map_err(|_| Error::InitMemory("Failed to insert memory_region".to_string()))?;
-
-        Ok(memory_regions)
-    }
 }
diff --git a/crates/vm-mm/src/manager.rs b/crates/vm-mm/src/manager.rs
@@ -26,15 +26,6 @@ impl<'a, C> IntoIterator for &'a MemoryAddressSpace<C> {
     }
 }
 
-impl<'a, C> IntoIterator for &'a mut MemoryAddressSpace<C> {
-    type Item = &'a mut MemoryRegion<C>;
-    type IntoIter = btree_map::ValuesMut<'a, u64, MemoryRegion<C>>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.regions.values_mut()
-    }
-}
-
 impl<C> MemoryAddressSpace<C>
 where
     C: MemoryContainer,
@@ -51,7 +42,7 @@ where
 
     pub fn gpa_to_hva(&mut self, gpa: u64) -> Result<*mut u8, Error> {
         let region = self.try_get_region_by_gpa_mut(gpa)?;
-        let hva = region.try_to_hva()?;
+        let hva = region.to_hva();
 
         let offset = gpa - region.gpa;
 
@@ -60,7 +51,7 @@ where
 
     pub fn memset(&mut self, gpa: u64, val: u8, len: usize) -> Result<(), Error> {
         let region = self.try_get_region_by_gpa_mut(gpa)?;
-        let hva = region.try_to_hva()?;
+        let hva = region.to_hva();
         let offset = gpa - region.gpa;
 
         if offset + len as u64 > region.len as u64 {
@@ -74,7 +65,7 @@ where
 
     pub fn copy_from_slice(&mut self, gpa: u64, buf: &[u8], len: usize) -> Result<(), Error> {
         let region = self.try_get_region_by_gpa_mut(gpa)?;
-        let hva = region.try_to_hva()?;
+        let hva = region.to_hva();
         let offset = gpa - region.gpa;
 
         if offset + len as u64 > region.len as u64 {
@@ -115,6 +106,7 @@ where
 mod tests {
     use memmap2::MmapMut;
 
+    use crate::allocator::Allocator;
     use crate::allocator::mmap_allocator::MmapAllocator;
     use crate::manager::MemoryAddressSpace;
     use crate::region::MemoryRegion;
@@ -123,19 +115,21 @@ mod tests {
     fn test_memory_layout() -> anyhow::Result<()> {
         let mut memory_as = MemoryAddressSpace::<MmapMut>::default();
 
+        let allocator = MmapAllocator;
+
         assert!(
             memory_as
-                .try_insert(MemoryRegion::placeholder(0, 10))
+                .try_insert(MemoryRegion::new(0, 10, allocator.alloc(10, None)?))
                 .is_ok()
         );
         assert!(
             memory_as
-                .try_insert(MemoryRegion::placeholder(5, 10))
+                .try_insert(MemoryRegion::new(5, 10, allocator.alloc(10, None)?))
                 .is_err()
         );
         assert!(
             memory_as
-                .try_insert(MemoryRegion::placeholder(10, 10))
+                .try_insert(MemoryRegion::new(10, 10, allocator.alloc(10, None)?))
                 .is_ok()
         );
 
@@ -146,11 +140,13 @@ mod tests {
     fn test_memory_write() -> anyhow::Result<()> {
         const GPA: u64 = 0;
         const LEN: usize = 512;
+
         let mut memory_as = MemoryAddressSpace::<MmapMut>::default();
+        let allocator = MmapAllocator;
+
+        let mut region = MemoryRegion::new(GPA, LEN, allocator.alloc(LEN, None)?);
 
-        let mut region = MemoryRegion::placeholder(GPA, LEN);
-        region.alloc(&MmapAllocator)?;
-        let hva = region.try_to_hva()?;
+        let hva = region.to_hva();
 
         assert!(memory_as.try_insert(region).is_ok());
 

diff --git a/crates/vm-mm/src/region.rs b/crates/vm-mm/src/region.rs
@@ -1,47 +1,20 @@
-use std::cell::OnceCell;
-
-use crate::allocator::Allocator;
 use crate::allocator::MemoryContainer;
-use crate::error::Error;
 
 pub struct MemoryRegion<C> {
     pub gpa: u64,
     pub len: usize,
-    // Placeholder: different platforms may use their own allocator for memory.
-    pub memory: OnceCell<C>,
+    pub memory: C,
 }
 
 impl<C> MemoryRegion<C>
 where
     C: MemoryContainer,
 {
-    pub fn placeholder(gpa: u64, len: usize) -> Self {
-        MemoryRegion {
-            gpa,
-            len,
-            memory: OnceCell::new(),
-        }
-    }
-
-    pub fn alloc<A>(&mut self, allocator: &A) -> Result<(), Error>
-    where
-        A: Allocator<Container = C>,
-    {
-        if self.memory.get().is_some() {
-            return Err(Error::MemoryAlreadyAllocated);
-        }
-
-        let memory = allocator.alloc(self.len, None)?;
-        let _ = self.memory.set(memory); // already checked
-
-        Ok(())
-    }
-
-    pub fn to_hva(&mut self) -> Option<*mut u8> {
-        self.memory.get_mut().map(C::to_hva)
+    pub fn new(gpa: u64, len: usize, memory: C) -> Self {
+        MemoryRegion { gpa, len, memory }
     }
 
-    pub fn try_to_hva(&mut self) -> Result<*mut u8, Error> {
-        self.to_hva().ok_or(Error::MemoryIsNotAllocated)
+    pub fn to_hva(&mut self) -> *mut u8 {
+        self.memory.to_hva()
     }
 }