add available() API to AddressAllocator

Signed-off-by: ylzh10 <46687083+ylzh10@users.noreply.github.com>

Author: ylzh10

CHANGELOG.md

@@ -3,6 +3,10 @@
 ## Upcoming version
 
 ### Added
+
+- Added `available()` API in `AddressAllocator` to allow
+  getting the available memories after `allocate/free()`s
+
 ### Changed
 ### Fixed
 ### Removed

src/address_allocator.rs

@@ -29,6 +29,8 @@ pub struct AddressAllocator {
     // tree will represent a memory location and can have two states either
     // `NodeState::Free` or `NodeState::Allocated`.
     interval_tree: IntervalTree,
+    // Available memory space
+    available: usize,
 }
 
 impl AddressAllocator {
@@ -43,6 +45,7 @@ impl AddressAllocator {
         Ok(AddressAllocator {
             address_space: aux_range,
             interval_tree: IntervalTree::new(aux_range),
+            available: aux_range.len() as usize,
         })
     }
 
@@ -63,13 +66,22 @@ impl AddressAllocator {
         policy: AllocPolicy,
     ) -> Result<RangeInclusive> {
         let constraint = Constraint::new(size, alignment, policy)?;
-        self.interval_tree.allocate(constraint)
+        let allocated = self.interval_tree.allocate(constraint)?;
+        self.available -= allocated.len() as usize;
+        Ok(allocated)
     }
 
     /// Deletes the specified memory slot or returns `ResourceNotAvailable` if
     /// the node was not allocated before.
     pub fn free(&mut self, key: &RangeInclusive) -> Result<()> {
-        self.interval_tree.free(key)
+        self.interval_tree.free(key)?;
+        self.available += key.len() as usize;
+        Ok(())
+    }
+
+    /// Returns the available memory size in this allocator.
+    pub fn available(&self) -> usize {
+        self.available
     }
 }
 
@@ -158,20 +170,27 @@ mod tests {
     #[test]
     fn test_allocate_with_alignment_first_ok() {
         let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap();
+        assert_eq!(pool.available(), 0x1000);
+        // Allocate 0x110
         assert_eq!(
             pool.allocate(0x110, 0x100, AllocPolicy::FirstMatch)
                 .unwrap(),
             RangeInclusive::new(0x1000, 0x110F).unwrap()
         );
+        assert_eq!(pool.available(), 0x1000 - 0x110);
+        // Allocate 0x100
         assert_eq!(
             pool.allocate(0x100, 0x100, AllocPolicy::FirstMatch)
                 .unwrap(),
             RangeInclusive::new(0x1200, 0x12FF).unwrap()
         );
+        assert_eq!(pool.available(), 0x1000 - 0x110 - 0x100);
+        // Allocate 0x10
         assert_eq!(
             pool.allocate(0x10, 0x100, AllocPolicy::FirstMatch).unwrap(),
             RangeInclusive::new(0x1300, 0x130F).unwrap()
         );
+        assert_eq!(pool.available(), 0x1000 - 0x110 - 0x100 - 0x10);
     }
 
     #[test]
@@ -230,18 +249,24 @@ mod tests {
     #[test]
     fn test_tree_allocate_address_free_and_realloc() {
         let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap();
+        assert_eq!(pool.available(), 0x1000);
+        // Allocate 0x800
         assert_eq!(
             pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch)
                 .unwrap(),
             RangeInclusive::new(0x1000, 0x17FF).unwrap()
         );
-
+        assert_eq!(pool.available(), 0x1000 - 0x800);
+        // Free 0x800
         let _ = pool.free(&RangeInclusive::new(0x1000, 0x17FF).unwrap());
+        assert_eq!(pool.available(), 0x1000);
+        // Allocate 0x800 again
         assert_eq!(
             pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch)
                 .unwrap(),
             RangeInclusive::new(0x1000, 0x17FF).unwrap()
         );
+        assert_eq!(pool.available(), 0x1000 - 0x800);
     }
 
     #[test]