diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index 6d34b202b..cafdac26c 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -41,8 +41,8 @@ use vm_memory::{
     GuestMemory, GuestMemoryAtomic, GuestMemoryError, GuestMemoryRegion, GuestUsize, MmapRegion,
 };
 use vm_migration::{
-    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
-    SnapshotDataSection, Snapshottable, Transportable, VersionMapped,
+    protocol::MemoryRange, protocol::MemoryRangeTable, Migratable, MigratableError, Pausable,
+    Snapshot, SnapshotDataSection, Snapshottable, Transportable, VersionMapped,
 };
 
 #[cfg(feature = "acpi")]
@@ -1508,6 +1508,24 @@ impl MemoryManager {
     pub fn memory_zones(&self) -> &MemoryZones {
         &self.memory_zones
     }
+
+    pub fn memory_range_table(&self) -> std::result::Result<MemoryRangeTable, MigratableError> {
+        let mut table = MemoryRangeTable::default();
+
+        for memory_zone in self.memory_zones.values() {
+            for region in memory_zone.regions() {
+                table.push(MemoryRange {
+                    gpa: region.start_addr().raw_value(),
+                    length: region.len() as u64,
+                });
+            }
+            if let Some(virtio_mem_zone) = memory_zone.virtio_mem_zone() {
+                table.extend(virtio_mem_zone.plugged_ranges());
+            }
+        }
+
+        Ok(table)
+    }
 }
 
 #[cfg(feature = "acpi")]
diff --git a/vmm/src/vm.rs b/vmm/src/vm.rs
index 35ccd1a94..360cde873 100644
--- a/vmm/src/vm.rs
+++ b/vmm/src/vm.rs
@@ -66,15 +66,14 @@ use std::{result, str, thread};
 use vm_device::Bus;
 #[cfg(target_arch = "x86_64")]
 use vm_device::BusDevice;
+#[cfg(any(target_arch = "aarch64", feature = "tdx"))]
+use vm_memory::Address;
+use vm_memory::{Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic};
 #[cfg(feature = "tdx")]
-use vm_memory::GuestMemory;
-use vm_memory::{
-    Address, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic, GuestMemoryRegion,
-};
+use vm_memory::{GuestMemory, GuestMemoryRegion};
 use vm_migration::{
-    protocol::{MemoryRange, MemoryRangeTable},
-    Migratable, MigratableError, Pausable, Snapshot, SnapshotDataSection, Snapshottable,
-    Transportable,
+    protocol::MemoryRangeTable, Migratable, MigratableError, Pausable, Snapshot,
+    SnapshotDataSection, Snapshottable, Transportable,
 };
 use vmm_sys_util::eventfd::EventFd;
 use vmm_sys_util::signal::unblock_signal;
@@ -2189,22 +2188,7 @@ impl Vm {
     }
 
     pub fn memory_range_table(&self) -> std::result::Result<MemoryRangeTable, MigratableError> {
-        let mut table = MemoryRangeTable::default();
-        let mm = self.memory_manager.lock().unwrap();
-
-        for memory_zone in mm.memory_zones().values() {
-            for region in memory_zone.regions() {
-                table.push(MemoryRange {
-                    gpa: region.start_addr().raw_value(),
-                    length: region.len() as u64,
-                });
-            }
-            if let Some(virtio_mem_zone) = memory_zone.virtio_mem_zone() {
-                table.extend(virtio_mem_zone.plugged_ranges());
-            }
-        }
-
-        Ok(table)
+        self.memory_manager.lock().unwrap().memory_range_table()
     }
 
     pub fn device_tree(&self) -> Arc<Mutex<DeviceTree>> {
@@ -2670,7 +2654,7 @@ mod tests {
 #[test]
 pub fn test_vm() {
     use hypervisor::VmExit;
-    use vm_memory::{GuestMemory, GuestMemoryRegion};
+    use vm_memory::{Address, GuestMemory, GuestMemoryRegion};
     // This example based on https://lwn.net/Articles/658511/
     let code = [
         0xba, 0xf8, 0x03, /* mov $0x3f8, %dx */