use crate::allocation_engine::IntervalTree;
use crate::{AllocPolicy, Constraint, Error, RangeInclusive, Result};
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug)]
pub struct AddressAllocator {
address_space: RangeInclusive,
interval_tree: IntervalTree,
}
impl AddressAllocator {
pub fn new(base: u64, size: u64) -> std::result::Result<Self, Error> {
let end = base
.checked_add(size.checked_sub(1).ok_or(Error::Underflow)?)
.ok_or(Error::Overflow)?;
let aux_range = RangeInclusive::new(base, end)?;
Ok(AddressAllocator {
address_space: aux_range,
interval_tree: IntervalTree::new(aux_range),
})
}
pub fn allocate(
&mut self,
size: u64,
alignment: u64,
policy: AllocPolicy,
) -> Result<RangeInclusive> {
let constraint = Constraint::new(size, alignment, policy)?;
self.interval_tree.allocate(constraint)
}
pub fn free(&mut self, key: &RangeInclusive) -> Result<()> {
self.interval_tree.free(key)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_regression_exact_match_length_check() {
let mut pool = AddressAllocator::new(0x0, 0x2000).unwrap();
let res = pool
.allocate(0x1000, 0x1000, AllocPolicy::ExactMatch(0x1000))
.unwrap();
assert_eq!(
pool.allocate(0x0, 0x1000, AllocPolicy::FirstMatch)
.unwrap_err(),
Error::InvalidSize(0x0)
);
assert_eq!(
pool.allocate(0x1000, 0x1000, AllocPolicy::ExactMatch(0x3))
.unwrap_err(),
Error::UnalignedAddress
);
assert_eq!(res, RangeInclusive::new(0x1000, 0x1FFF).unwrap());
let res = pool
.allocate(0x1000, 0x1000, AllocPolicy::ExactMatch(0x0))
.unwrap();
assert_eq!(res, RangeInclusive::new(0x0, 0x0FFF).unwrap());
}
#[test]
fn test_new_fails_overflow() {
assert_eq!(
AddressAllocator::new(u64::max_value(), 0x100).unwrap_err(),
Error::Overflow
);
}
#[test]
fn test_new_fails_size_zero() {
assert_eq!(
AddressAllocator::new(0x1000, 0x0).unwrap_err(),
Error::Underflow
);
}
#[test]
fn test_allocate_fails_alignment_zero() {
let mut pool = AddressAllocator::new(0x1000, 0x10000).unwrap();
assert_eq!(
pool.allocate(0x100, 0, AllocPolicy::FirstMatch)
.unwrap_err(),
Error::InvalidAlignment
);
}
#[test]
fn test_allocate_fails_alignment_non_power_of_two() {
let mut pool = AddressAllocator::new(0x1000, 0x10000).unwrap();
assert_eq!(
pool.allocate(0x100, 200, AllocPolicy::FirstMatch)
.unwrap_err(),
Error::InvalidAlignment
);
}
#[test]
fn test_allocate_fails_not_enough_space() {
let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap();
assert_eq!(
pool.allocate(0x800, 0x100, AllocPolicy::LastMatch).unwrap(),
RangeInclusive::new(0x1800, 0x1FFF).unwrap()
);
assert_eq!(
pool.allocate(0x900, 0x100, AllocPolicy::FirstMatch)
.unwrap_err(),
Error::ResourceNotAvailable
);
assert_eq!(
pool.allocate(0x400, 0x100, AllocPolicy::FirstMatch)
.unwrap(),
RangeInclusive::new(0x1000, 0x13FF).unwrap()
);
}
#[test]
fn test_allocate_with_alignment_first_ok() {
let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap();
assert_eq!(
pool.allocate(0x110, 0x100, AllocPolicy::FirstMatch)
.unwrap(),
RangeInclusive::new(0x1000, 0x110F).unwrap()
);
assert_eq!(
pool.allocate(0x100, 0x100, AllocPolicy::FirstMatch)
.unwrap(),
RangeInclusive::new(0x1200, 0x12FF).unwrap()
);
assert_eq!(
pool.allocate(0x10, 0x100, AllocPolicy::FirstMatch).unwrap(),
RangeInclusive::new(0x1300, 0x130F).unwrap()
);
}
#[test]
fn test_allocate_with_alignment_last_ok() {
let mut pool_reverse = AddressAllocator::new(0x1000, 0x10000).unwrap();
assert_eq!(
pool_reverse
.allocate(0x110, 0x100, AllocPolicy::LastMatch)
.unwrap(),
RangeInclusive::new(0x10E00, 0x10F0F).unwrap()
);
assert_eq!(
pool_reverse
.allocate(0x100, 0x100, AllocPolicy::LastMatch)
.unwrap(),
RangeInclusive::new(0x10D00, 0x10DFF).unwrap()
);
assert_eq!(
pool_reverse
.allocate(0x10, 0x100, AllocPolicy::LastMatch)
.unwrap(),
RangeInclusive::new(0x10C00, 0x10C0F).unwrap()
);
}
#[test]
fn test_allocate_address_not_enough_space() {
let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap();
assert_eq!(
pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch)
.unwrap(),
RangeInclusive::new(0x1000, 0x17FF).unwrap()
);
assert_eq!(
pool.allocate(0x200, 0x100, AllocPolicy::ExactMatch(0x1A00))
.unwrap(),
RangeInclusive::new(0x1A00, 0x1BFF).unwrap()
);
assert_eq!(
pool.allocate(0x800, 0x100, AllocPolicy::ExactMatch(0x1800))
.unwrap_err(),
Error::ResourceNotAvailable
);
assert_eq!(
pool.allocate(0x100, 0x100, AllocPolicy::ExactMatch(0x1800))
.unwrap(),
RangeInclusive::new(0x1800, 0x18FF).unwrap()
);
}
#[test]
fn test_tree_allocate_address_free_and_realloc() {
let mut pool = AddressAllocator::new(0x1000, 0x1000).unwrap();
assert_eq!(
pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch)
.unwrap(),
RangeInclusive::new(0x1000, 0x17FF).unwrap()
);
let _ = pool.free(&RangeInclusive::new(0x1000, 0x17FF).unwrap());
assert_eq!(
pool.allocate(0x800, 0x100, AllocPolicy::FirstMatch)
.unwrap(),
RangeInclusive::new(0x1000, 0x17FF).unwrap()
);
}
#[test]
fn test_allocate_address_fail_free_and_realloc() {
let mut pool = AddressAllocator::new(0x0, 0x1000).unwrap();
assert_eq!(
pool.allocate(0x2000, 0x100, AllocPolicy::FirstMatch)
.unwrap_err(),
Error::ResourceNotAvailable
);
assert_eq!(
pool.free(&RangeInclusive::new(0x1200, 0x3200).unwrap())
.unwrap_err(),
Error::ResourceNotAvailable
);
assert_eq!(
pool.allocate(0x4FE, 0x100, AllocPolicy::ExactMatch(0x500))
.unwrap(),
RangeInclusive::new(0x500, 0x9FD).unwrap()
);
assert!(pool
.free(&RangeInclusive::new(0x500, 0x9FD).unwrap())
.is_ok());
}
}