From ad9034ed1d5852d7882756bb1cfb29b955911be6 Mon Sep 17 00:00:00 2001
From: Julian Stecklina <julian.stecklina@cyberus-technology.de>
Date: Fri, 7 Nov 2025 14:51:49 +0100
Subject: [PATCH] vm-migration: optimize dirty bitmap scanning

Adding itertools as dependency improves the iteration code in the
following significantly.

With this change, we don't need a copy of the vector. Just something
that can be coerced into an iterator. We also use the bit position
iterator to make the code somewhat clearer. The new code is much
faster, because it will not iterate over every bit, just each 1 bit in
the input.

The next commit will complete this optimization and have some concrete
numbers.

On-behalf-of: SAP julian.stecklina@sap.com
Signed-off-by: Julian Stecklina <julian.stecklina@cyberus-technology.de>
---
 Cargo.lock                   | 16 ++++++++++
 Cargo.toml                   |  1 +
 vm-migration/Cargo.toml      |  1 +
 vm-migration/src/protocol.rs | 61 +++++++++++++++++++++---------------
 4 files changed, 54 insertions(+), 25 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 118b60cea..feb65da67 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -597,6 +597,12 @@ dependencies = [
  "windows-sys 0.61.0",
 ]
 
+[[package]]
+name = "either"
+version = "1.15.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719"
+
 [[package]]
 name = "endi"
 version = "1.1.0"
@@ -1031,6 +1037,15 @@ version = "1.70.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf"
 
+[[package]]
+name = "itertools"
+version = "0.14.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285"
+dependencies = [
+ "either",
+]
+
 [[package]]
 name = "itoa"
 version = "1.0.15"
@@ -2418,6 +2433,7 @@ name = "vm-migration"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "itertools",
  "serde",
  "serde_json",
  "thiserror 2.0.17",
diff --git a/Cargo.toml b/Cargo.toml
index 4b74245de..e1d2f158b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -145,6 +145,7 @@ dirs = "6.0.0"
 env_logger = "0.11.8"
 epoll = "4.4.0"
 flume = "0.11.1"
+itertools = "0.14.0"
 libc = "0.2.177"
 log = "0.4.28"
 signal-hook = "0.3.18"
diff --git a/vm-migration/Cargo.toml b/vm-migration/Cargo.toml
index 7a8c9337b..69d57076f 100644
--- a/vm-migration/Cargo.toml
+++ b/vm-migration/Cargo.toml
@@ -6,6 +6,7 @@ version = "0.1.0"
 
 [dependencies]
 anyhow = { workspace = true }
+itertools = { workspace = true }
 serde = { workspace = true, features = ["derive", "rc"] }
 serde_json = { workspace = true }
 thiserror = { workspace = true }
diff --git a/vm-migration/src/protocol.rs b/vm-migration/src/protocol.rs
index 2ed782ae3..8c320c68e 100644
--- a/vm-migration/src/protocol.rs
+++ b/vm-migration/src/protocol.rs
@@ -5,10 +5,12 @@
 
 use std::io::{Read, Write};
 
+use itertools::Itertools;
 use serde::{Deserialize, Serialize};
 use vm_memory::ByteValued;
 
 use crate::MigratableError;
+use crate::bitpos_iterator::BitposIteratorExt;
 
 // Migration protocol
 // 1: Source establishes communication with destination (file socket or TCP connection.)
@@ -211,38 +213,47 @@ pub struct MemoryRange {
     pub length: u64,
 }
 
+impl MemoryRange {
+    /// Turn an iterator over the dirty bitmap into an iterator of dirty ranges.
+    pub fn dirty_ranges(
+        bitmap: impl IntoIterator<Item = u64>,
+        start_addr: u64,
+        page_size: u64,
+    ) -> impl Iterator<Item = Self> {
+        bitmap
+            .into_iter()
+            .bit_positions()
+            // Turn them into single-element ranges for coalesce.
+            .map(|b| b..(b + 1))
+            // Merge adjacent ranges.
+            .coalesce(|prev, curr| {
+                if prev.end == curr.start {
+                    Ok(prev.start..curr.end)
+                } else {
+                    Err((prev, curr))
+                }
+            })
+            .map(move |r| Self {
+                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(bitmap: Vec<u64>, start_addr: u64, page_size: u64) -> Self {
-        let mut table = MemoryRangeTable::default();
-        let mut entry: Option<MemoryRange> = None;
-        for (i, block) in bitmap.iter().enumerate() {
-            for j in 0..64 {
-                let is_page_dirty = ((block >> j) & 1u64) != 0u64;
-                let page_offset = ((i * 64) + j) as u64 * page_size;
-                if is_page_dirty {
-                    if let Some(entry) = &mut entry {
-                        entry.length += page_size;
-                    } else {
-                        entry = Some(MemoryRange {
-                            gpa: start_addr + page_offset,
-                            length: page_size,
-                        });
-                    }
-                } else if let Some(entry) = entry.take() {
-                    table.push(entry);
-                }
-            }
+    pub fn from_bitmap(
+        bitmap: impl IntoIterator<Item = u64>,
+        start_addr: u64,
+        page_size: u64,
+    ) -> Self {
+        Self {
+            data: MemoryRange::dirty_ranges(bitmap, start_addr, page_size).collect(),
         }
-        if let Some(entry) = entry.take() {
-            table.push(entry);
-        }
-
-        table
     }
 
     pub fn regions(&self) -> &[MemoryRange] {