diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c87434..2fca66b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,10 @@ ## Upcoming version ### Added + +- Added `used()` API in `AddressAllocator` to allow + getting the used memories after `allocate/free()`s + ### Changed ### Fixed ### Removed diff --git a/src/address_allocator.rs b/src/address_allocator.rs index 5580154..a93f779 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, + // Used memory space in the address space. + used: usize, } impl AddressAllocator { @@ -43,6 +45,7 @@ impl AddressAllocator { Ok(AddressAllocator { address_space: aux_range, interval_tree: IntervalTree::new(aux_range), + used: 0, }) } @@ -63,13 +66,30 @@ 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.used = self + .used + .checked_add(allocated.len() as usize) + .expect("Failed to calculate used memory"); + 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.used = self + .used + .checked_sub(key.len() as usize) + .expect("Failed to calculate used memory"); + Ok(()) + } + + /// Returns the used memory size in this allocator. + /// NOTE that due to fragmentations, not all unused memory may be available + /// for next `allocate()` call! + pub fn used(&self) -> usize { + self.used } } @@ -158,20 +178,27 @@ mod tests { #[test] fn test_allocate_with_alignment_first_ok() { let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap(); + assert_eq!(pool.used(), 0); + // Allocate 0x110 assert_eq!( pool.allocate(0x110, 0x100, AllocPolicy::FirstMatch) .unwrap(), RangeInclusive::new(0x1000, 0x110F).unwrap() ); + assert_eq!(pool.used(), 0x110); + // Allocate 0x100 assert_eq!( pool.allocate(0x100, 0x100, AllocPolicy::FirstMatch) .unwrap(), RangeInclusive::new(0x1200, 0x12FF).unwrap() ); + assert_eq!(pool.used(), 0x110 + 0x100); + // Allocate 0x10 assert_eq!( pool.allocate(0x10, 0x100, AllocPolicy::FirstMatch).unwrap(), RangeInclusive::new(0x1300, 0x130F).unwrap() ); + assert_eq!(pool.used(), 0x110 + 0x100 + 0x10); } #[test] @@ -230,18 +257,24 @@ mod tests { #[test] fn test_tree_allocate_address_free_and_realloc() { let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap(); + assert_eq!(pool.used(), 0); + // Allocate 0x800 assert_eq!( pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch) .unwrap(), RangeInclusive::new(0x1000, 0x17FF).unwrap() ); - + assert_eq!(pool.used(), 0x800); + // Free 0x800 let _ = pool.free(&RangeInclusive::new(0x1000, 0x17FF).unwrap()); + assert_eq!(pool.used(), 0); + // Allocate 0x800 again assert_eq!( pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch) .unwrap(), RangeInclusive::new(0x1000, 0x17FF).unwrap() ); + assert_eq!(pool.used(), 0x800); } #[test]