From 4b451b01d9fd7dde1b0dc423c2b49eb87e8e3c27 Mon Sep 17 00:00:00 2001
From: Samuel Ortiz <sameo@linux.intel.com>
Date: Tue, 21 May 2019 17:28:03 +0200
Subject: [PATCH] vm-allocator: Allow for freeing address ranges

We can only free ranges that exactly map an already allocated one, i.e.
this is not a range resizing.

Signed-off-by: Samuel Ortiz <sameo@linux.intel.com>
---
 vm-allocator/src/address.rs | 61 +++++++++++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/vm-allocator/src/address.rs b/vm-allocator/src/address.rs
index f0d1f01df..05b42cd3b 100644
--- a/vm-allocator/src/address.rs
+++ b/vm-allocator/src/address.rs
@@ -181,6 +181,16 @@ impl AddressAllocator {
 
         Some(new_addr)
     }
+
+    /// Free an already allocated address range.
+    /// We can only free a range if it matches exactly an already allocated range.
+    pub fn free(&mut self, address: GuestAddress, size: GuestUsize) {
+        if let Some(&range_size) = self.ranges.get(&address) {
+            if size == range_size {
+                self.ranges.remove(&address);
+            }
+        }
+    }
 }
 
 #[cfg(test)]
@@ -290,4 +300,55 @@ mod tests {
             Some(GuestAddress(0x1b00))
         );
     }
+
+    #[test]
+    fn allocate_address_free_and_realloc() {
+        let mut pool = AddressAllocator::new(GuestAddress(0x1000), 0x1000, Some(0x100)).unwrap();
+
+        // First range is [0x1200:0x1a00]
+        assert_eq!(
+            pool.allocate(Some(GuestAddress(0x1200)), 0x800),
+            Some(GuestAddress(0x1200))
+        );
+
+        pool.free(GuestAddress(0x1200), 0x800);
+
+        assert_eq!(
+            pool.allocate(Some(GuestAddress(0x1200)), 0x800),
+            Some(GuestAddress(0x1200))
+        );
+    }
+
+    #[test]
+    fn allocate_address_free_fail_and_realloc() {
+        let mut pool = AddressAllocator::new(GuestAddress(0x1000), 0x1000, Some(0x100)).unwrap();
+
+        // First range is [0x1200:0x1a00]
+        assert_eq!(
+            pool.allocate(Some(GuestAddress(0x1200)), 0x800),
+            Some(GuestAddress(0x1200))
+        );
+
+        // We try to free a range smaller than the allocated one.
+        pool.free(GuestAddress(0x1200), 0x100);
+
+        assert_eq!(pool.allocate(Some(GuestAddress(0x1200)), 0x800), None);
+    }
+
+    #[test]
+    fn allocate_address_fail_free_and_realloc() {
+        let mut pool = AddressAllocator::new(GuestAddress(0x1000), 0x1000, Some(0x100)).unwrap();
+
+        // First allocation fails
+        assert_eq!(pool.allocate(Some(GuestAddress(0x1200)), 0x2000), None);
+
+        // We try to free a range that was not allocated.
+        pool.free(GuestAddress(0x1200), 0x2000);
+
+        // Now we try an allocation that should succeed.
+        assert_eq!(
+            pool.allocate(Some(GuestAddress(0x1200)), 0x800),
+            Some(GuestAddress(0x1200))
+        );
+    }
 }