#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Block {
offset: u64,
size: u64,
}
fn align_up(value: u64, align: u64) -> u64 {
value.div_ceil(align) * align
}
#[derive(Clone, Copy, Debug)]
struct Pending {
offset: u64,
size: u64,
retire_frame: u64,
}
pub(crate) struct RangeAllocator {
free: Vec<Block>,
pending: Vec<Pending>,
}
impl RangeAllocator {
pub(crate) fn new() -> Self {
Self {
free: Vec::new(),
pending: Vec::new(),
}
}
pub(crate) fn alloc(&mut self, size: u64) -> Option<u64> {
self.alloc_aligned(size, 1)
}
pub(crate) fn alloc_aligned(&mut self, size: u64, align: u64) -> Option<u64> {
if size == 0 {
return Some(0);
}
let align = align.max(1);
let mut best: Option<(usize, u64, u64)> = None;
for (i, b) in self.free.iter().enumerate() {
let offset = align_up(b.offset, align);
let pad = offset - b.offset;
let Some(usable) = b.size.checked_sub(pad) else {
continue;
};
let Some(waste) = usable.checked_sub(size) else {
continue;
};
if best.is_none_or(|(_, _, best_waste)| waste < best_waste) {
best = Some((i, offset, waste));
}
}
let (i, offset, _) = best?;
let block = self.free[i];
let pad = offset - block.offset;
let tail_offset = offset + size;
let tail_size = block.offset + block.size - tail_offset;
match (pad, tail_size) {
(0, 0) => {
self.free.remove(i);
}
(0, _) => {
self.free[i] = Block {
offset: tail_offset,
size: tail_size,
};
}
(_, 0) => {
self.free[i] = Block {
offset: block.offset,
size: pad,
};
}
_ => {
self.free[i] = Block {
offset: block.offset,
size: pad,
};
self.free.insert(
i + 1,
Block {
offset: tail_offset,
size: tail_size,
},
);
}
}
Some(offset)
}
pub(crate) fn free(&mut self, offset: u64, size: u64, retire_frame: u64) {
if size == 0 {
return;
}
self.pending.push(Pending {
offset,
size,
retire_frame,
});
}
pub(crate) fn reclaim(&mut self, current_frame: u64) {
let mut i = 0;
while i < self.pending.len() {
if self.pending[i].retire_frame <= current_frame {
let p = self.pending.swap_remove(i);
self.insert_free(Block {
offset: p.offset,
size: p.size,
});
} else {
i += 1;
}
}
}
pub(crate) fn free_bytes(&self) -> u64 {
self.free.iter().map(|b| b.size).sum()
}
#[cfg(test)]
pub(crate) fn free_block_count(&self) -> usize {
self.free.len()
}
fn insert_free(&mut self, block: Block) {
let pos = self.free.partition_point(|b| b.offset < block.offset);
debug_assert!(
pos == self.free.len() || block.offset + block.size <= self.free[pos].offset,
"RangeAllocator: freed block overlaps an existing free block"
);
debug_assert!(
pos == 0 || self.free[pos - 1].offset + self.free[pos - 1].size <= block.offset,
"RangeAllocator: freed block overlaps an existing free block"
);
self.free.insert(pos, block);
while pos + 1 < self.free.len()
&& self.free[pos].offset + self.free[pos].size == self.free[pos + 1].offset
{
let next_size = self.free[pos + 1].size;
self.free[pos].size += next_size;
self.free.remove(pos + 1);
}
if pos > 0 && self.free[pos - 1].offset + self.free[pos - 1].size == self.free[pos].offset {
let this_size = self.free[pos].size;
self.free[pos - 1].size += this_size;
self.free.remove(pos);
}
}
}
impl Default for RangeAllocator {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn seed(a: &mut RangeAllocator, offset: u64, size: u64) {
a.free(offset, size, 0);
a.reclaim(0);
}
#[test]
fn alloc_from_a_single_region_shaves_the_front() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 300);
assert_eq!(a.alloc(100), Some(0));
assert_eq!(a.alloc(100), Some(100));
assert_eq!(a.alloc(100), Some(200));
assert_eq!(a.alloc(1), None);
assert_eq!(a.free_bytes(), 0);
}
#[test]
fn alloc_returns_none_when_no_block_is_large_enough() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 40);
seed(&mut a, 100, 40);
assert_eq!(a.alloc(50), None);
assert_eq!(a.alloc(40), Some(0));
}
#[test]
fn best_fit_consumes_an_exact_size_block_whole() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 200);
seed(&mut a, 500, 100);
seed(&mut a, 900, 50);
assert_eq!(a.alloc(100), Some(500));
assert_eq!(a.free_block_count(), 2);
assert_eq!(a.free_bytes(), 250);
}
#[test]
fn adjacent_regions_coalesce_into_one_block() {
let mut a = RangeAllocator::new();
seed(&mut a, 100, 50);
seed(&mut a, 0, 100); seed(&mut a, 150, 50); assert_eq!(a.free_block_count(), 1);
assert_eq!(a.alloc(200), Some(0));
}
#[test]
fn freed_region_is_withheld_until_its_retire_frame() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 100);
assert_eq!(a.alloc(100), Some(0));
a.free(0, 100, 5);
a.reclaim(4);
assert_eq!(a.alloc(100), None); a.reclaim(5);
assert_eq!(a.alloc(100), Some(0)); }
#[test]
fn reclaimed_region_coalesces_with_neighbours() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 300);
let first = a.alloc(100).unwrap();
let second = a.alloc(100).unwrap();
assert_eq!((first, second), (0, 100));
a.free(first, 100, 1);
a.free(second, 100, 1);
a.reclaim(1);
assert_eq!(a.free_block_count(), 1);
assert_eq!(a.alloc(300), Some(0));
}
#[test]
fn aligned_alloc_rounds_the_offset_up() {
let mut a = RangeAllocator::new();
seed(&mut a, 4, 500);
assert_eq!(a.alloc_aligned(100, 256), Some(256));
}
#[test]
fn alignment_padding_stays_allocatable() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 1024);
assert_eq!(a.alloc_aligned(64, 1), Some(0));
assert_eq!(a.alloc_aligned(64, 256), Some(256));
assert_eq!(a.alloc_aligned(192, 1), Some(64));
assert_eq!(a.free_bytes(), 1024 - 64 - 64 - 192);
}
#[test]
fn freeing_an_aligned_alloc_returns_exactly_what_was_handed_out() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 1024);
let off = a.alloc_aligned(100, 256).unwrap();
assert_eq!(off, 0);
let second = a.alloc_aligned(100, 256).unwrap();
assert_eq!(second, 256);
a.free(off, 100, 0);
a.free(second, 100, 0);
a.reclaim(0);
assert_eq!(a.free_block_count(), 1);
assert_eq!(a.free_bytes(), 1024);
}
#[test]
fn best_fit_counts_alignment_padding_as_waste() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 140);
seed(&mut a, 512, 32);
assert_eq!(a.alloc_aligned(24, 256), Some(512));
}
#[test]
fn aligned_alloc_fails_when_padding_pushes_it_past_the_block() {
let mut a = RangeAllocator::new();
seed(&mut a, 100, 200);
assert_eq!(a.alloc_aligned(100, 256), None);
assert_eq!(a.alloc_aligned(100, 1), Some(100));
}
#[test]
fn evict_and_reuse_places_a_different_mesh_in_freed_space() {
let mut a = RangeAllocator::new();
seed(&mut a, 0, 128);
seed(&mut a, 128, 64);
let a_off = a.alloc(128).unwrap();
let b_off = a.alloc(64).unwrap();
assert_eq!((a_off, b_off), (0, 128));
a.free(a_off, 128, 12);
a.reclaim(12);
assert_eq!(a.alloc(128), Some(0));
}
}