diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c87434..81e3287 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/address_allocator.rs b/src/address_allocator.rs index 5580154..7b79d14 100644 --- a/src/address_allocator.rs +++ b/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 { 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]