From a7fa3a0c866fb48a06b7446125f31e6becb2198d Mon Sep 17 00:00:00 2001
From: Philipp Schuster <philipp.schuster@cyberus-technology.de>
Date: Wed, 12 Nov 2025 11:18:35 +0100
Subject: [PATCH] vm-migration: better naming + unittests

Signed-off-by: Philipp Schuster <philipp.schuster@cyberus-technology.de>
On-behalf-of: SAP philipp.schuster@sap.com
---
 virtio-devices/src/mem.rs                     |  2 +-
 .../src/vhost_user/vu_common_ctrl.rs          |  2 +-
 vm-migration/src/protocol.rs                  | 68 +++++++++++++++----
 vmm/src/memory_manager.rs                     |  2 +-
 4 files changed, 56 insertions(+), 18 deletions(-)

diff --git a/virtio-devices/src/mem.rs b/virtio-devices/src/mem.rs
index 3378498b1..51bb41ba0 100644
--- a/virtio-devices/src/mem.rs
+++ b/virtio-devices/src/mem.rs
@@ -393,7 +393,7 @@ impl BlocksState {
 
         // TODO We can avoid creating a new bitmap here, if we switch the code
         // to use Vec<u64> to keep dirty bits and just pass it as is.
-        MemoryRangeTable::from_bitmap(bitmap, start_addr, VIRTIO_MEM_DEFAULT_BLOCK_SIZE)
+        MemoryRangeTable::from_dirty_bitmap(bitmap, start_addr, VIRTIO_MEM_DEFAULT_BLOCK_SIZE)
     }
 }
 
diff --git a/virtio-devices/src/vhost_user/vu_common_ctrl.rs b/virtio-devices/src/vhost_user/vu_common_ctrl.rs
index b61b2aad3..5f81bc247 100644
--- a/virtio-devices/src/vhost_user/vu_common_ctrl.rs
+++ b/virtio-devices/src/vhost_user/vu_common_ctrl.rs
@@ -575,7 +575,7 @@ impl VhostUserHandle {
                 let ptr = region.as_ptr() as *const u64;
                 std::slice::from_raw_parts(ptr, len)
             };
-            Ok(MemoryRangeTable::from_bitmap(
+            Ok(MemoryRangeTable::from_dirty_bitmap(
                 bitmap.iter().copied(),
                 0,
                 4096,
diff --git a/vm-migration/src/protocol.rs b/vm-migration/src/protocol.rs
index 8c320c68e..b66c9a7ad 100644
--- a/vm-migration/src/protocol.rs
+++ b/vm-migration/src/protocol.rs
@@ -207,19 +207,28 @@ impl Response {
 }
 
 #[repr(C)]
-#[derive(Clone, Default, Serialize, Deserialize)]
+#[derive(Clone, Default, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub struct MemoryRange {
     pub gpa: u64,
     pub length: u64,
 }
 
-impl MemoryRange {
-    /// Turn an iterator over the dirty bitmap into an iterator of dirty ranges.
-    pub fn dirty_ranges(
+#[derive(Clone, Default, Serialize, Deserialize)]
+pub struct MemoryRangeTable {
+    data: Vec<MemoryRange>,
+}
+
+impl MemoryRangeTable {
+    /// Converts an iterator over a dirty bitmap into an iterator of dirty
+    /// [`MemoryRange`]s, merging consecutive dirty pages into contiguous ranges.
+    ///
+    /// A memory page (i.e., a range) is marked dirty when its corresponding bit
+    /// is set.
+    fn dirty_ranges_iter(
         bitmap: impl IntoIterator<Item = u64>,
         start_addr: u64,
         page_size: u64,
-    ) -> impl Iterator<Item = Self> {
+    ) -> impl Iterator<Item = MemoryRange> {
         bitmap
             .into_iter()
             .bit_positions()
@@ -233,26 +242,23 @@ impl MemoryRange {
                     Err((prev, curr))
                 }
             })
-            .map(move |r| Self {
+            .map(move |r| MemoryRange {
                 gpa: start_addr + r.start * page_size,
                 length: (r.end - r.start) * page_size,
             })
     }
-}
 
-#[derive(Clone, Default, Serialize, Deserialize)]
-pub struct MemoryRangeTable {
-    data: Vec<MemoryRange>,
-}
-
-impl MemoryRangeTable {
-    pub fn from_bitmap(
+    /// Creates a new [`MemoryRangeTable`] from a bitmap (represented as
+    /// multiple `u64`) where each bit corresponds to a dirty memory page.
+    ///
+    /// Only dirty ranges are represented in the resulting bitmap.
+    pub fn from_dirty_bitmap(
         bitmap: impl IntoIterator<Item = u64>,
         start_addr: u64,
         page_size: u64,
     ) -> Self {
         Self {
-            data: MemoryRange::dirty_ranges(bitmap, start_addr, page_size).collect(),
+            data: Self::dirty_ranges_iter(bitmap, start_addr, page_size).collect(),
         }
     }
 
@@ -312,3 +318,35 @@ impl MemoryRangeTable {
         Self { data }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::protocol::{MemoryRange, MemoryRangeTable};
+
+    #[test]
+    fn test_memory_range_table_from_dirty_ranges_iter() {
+        let input = [0b1111_1110_1110, 0b1_0000];
+
+        let start_gpa = 0x1000;
+        let page_size = 0x1000;
+
+        let range = MemoryRangeTable::from_dirty_bitmap(input, start_gpa, page_size);
+        assert_eq!(
+            range.regions(),
+            &[
+                MemoryRange {
+                    gpa: start_gpa + page_size,
+                    length: page_size * 3,
+                },
+                MemoryRange {
+                    gpa: start_gpa + 5 * page_size,
+                    length: page_size * 7,
+                },
+                MemoryRange {
+                    gpa: start_gpa + (64 + 4) * page_size,
+                    length: page_size,
+                }
+            ]
+        );
+    }
+}
diff --git a/vmm/src/memory_manager.rs b/vmm/src/memory_manager.rs
index e3b08adf7..db00f4d61 100644
--- a/vmm/src/memory_manager.rs
+++ b/vmm/src/memory_manager.rs
@@ -2605,7 +2605,7 @@ impl Migratable for MemoryManager {
                 .zip(vmm_dirty_bitmap.iter())
                 .map(|(x, y)| x | y);
 
-            let sub_table = MemoryRangeTable::from_bitmap(dirty_bitmap, r.gpa, 4096);
+            let sub_table = MemoryRangeTable::from_dirty_bitmap(dirty_bitmap, r.gpa, 4096);
 
             if sub_table.regions().is_empty() {
                 info!("Dirty Memory Range Table is empty");