From 4450c44fbcf71b40ab5aaea885c13b181145b64c Mon Sep 17 00:00:00 2001
From: Sebastien Boeuf <sebastien.boeuf@intel.com>
Date: Thu, 23 Sep 2021 10:21:46 +0200
Subject: [PATCH] virtio-devices: mem: Create a MemoryRangeTable from
 BlocksState

This is going to be useful to let virtio-mem report the list of ranges
that are currently plugged, so that both snapshot/restore and migration
will copy only what is needed.

Signed-off-by: Sebastien Boeuf <sebastien.boeuf@intel.com>
---
 virtio-devices/src/mem.rs | 21 +++++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/virtio-devices/src/mem.rs b/virtio-devices/src/mem.rs
index 8f28bd1e4..dee35b0e3 100644
--- a/virtio-devices/src/mem.rs
+++ b/virtio-devices/src/mem.rs
@@ -40,6 +40,7 @@ use vm_memory::{
     Address, ByteValued, Bytes, GuestAddress, GuestAddressSpace, GuestMemoryAtomic,
     GuestMemoryError, GuestMemoryRegion,
 };
+use vm_migration::protocol::MemoryRangeTable;
 use vm_migration::{
     Migratable, MigratableError, Pausable, Snapshot, Snapshottable, Transportable, VersionMapped,
 };
@@ -426,6 +427,26 @@ impl BlocksState {
     fn inner(&self) -> &Vec<bool> {
         &self.bitmap
     }
+
+    pub fn memory_ranges(&self, start_addr: u64, plugged: bool) -> MemoryRangeTable {
+        let mut bitmap: Vec<u64> = Vec::new();
+        let mut i = 0;
+        for (j, bit) in self.bitmap.iter().enumerate() {
+            if j % 64 == 0 {
+                bitmap.push(0);
+
+                if j != 0 {
+                    i += 1;
+                }
+            }
+
+            if *bit == plugged {
+                bitmap[i] |= 1 << (j % 64);
+            }
+        }
+
+        MemoryRangeTable::from_bitmap(bitmap, start_addr, VIRTIO_MEM_DEFAULT_BLOCK_SIZE)
+    }
 }
 
 struct MemEpollHandler {