1use std::{
28 fmt::Debug,
29 iter::Sum,
30 ops::{Add, AddAssign, Range, Rem, Sub},
31};
32
33#[derive(Debug)]
50pub struct RangeAllocator<T> {
51 initial_range: Range<T>,
53 free_ranges: Vec<Range<T>>,
57}
58
59#[derive(Clone, Debug, PartialEq)]
64pub struct RangeAllocationError<T> {
65 pub fragmented_free_length: T,
69}
70
71impl<T> RangeAllocator<T>
72where
73 T: Clone + Copy + Add<Output = T> + AddAssign + Sub<Output = T> + Eq + PartialOrd + Debug,
74{
75 pub fn new(range: Range<T>) -> Self {
88 RangeAllocator {
89 initial_range: range.clone(),
90 free_ranges: vec![range],
91 }
92 }
93
94 pub fn initial_range(&self) -> &Range<T> {
97 &self.initial_range
98 }
99
100 pub fn grow_to(&mut self, new_end: T) {
119 let initial_range_end = self.initial_range.end;
120 if let Some(last_range) = self
121 .free_ranges
122 .last_mut()
123 .filter(|last_range| last_range.end == initial_range_end)
124 {
125 last_range.end = new_end;
126 } else {
127 self.free_ranges.push(self.initial_range.end..new_end);
128 }
129
130 self.initial_range.end = new_end;
131 }
132
133 fn allocate_range_impl(
134 &mut self,
135 length: T,
136 align_start: impl Fn(T) -> T,
137 ) -> Result<Range<T>, RangeAllocationError<T>> {
138 assert_ne!(length + length, length);
139
140 #[allow(clippy::eq_op)]
145 let mut fragmented_free_length = length - length;
146 let mut best_fit: Option<(usize, T)> = None;
147
148 for (index, range) in self.free_ranges.iter().cloned().enumerate() {
149 let range_length = range.end - range.start;
150 fragmented_free_length += range_length;
151
152 let aligned_start = align_start(range.start);
153
154 if aligned_start >= range.end {
155 continue;
156 }
157 let usable_length = range.end - aligned_start;
158 if usable_length < length {
159 continue;
160 } else if usable_length == length {
161 best_fit = Some((index, aligned_start));
163 break;
164 }
165 best_fit = Some(match best_fit {
166 Some((best_index, best_aligned_start)) => {
167 let best_usable = self.free_ranges[best_index].end - best_aligned_start;
169 if usable_length < best_usable {
170 (index, aligned_start)
171 } else {
172 (best_index, best_aligned_start)
173 }
174 }
175 None => (index, aligned_start),
176 });
177 }
178
179 match best_fit {
180 Some((index, aligned_start)) => {
181 let range = self.free_ranges[index].clone();
182 let alloc_end = aligned_start + length;
183
184 let has_prefix = aligned_start > range.start;
185 let has_suffix = alloc_end < range.end;
186
187 match (has_prefix, has_suffix) {
188 (false, false) => {
189 self.free_ranges.remove(index);
190 }
191 (false, true) => {
192 self.free_ranges[index].start = alloc_end;
193 }
194 (true, false) => {
195 self.free_ranges[index].end = aligned_start;
196 }
197 (true, true) => {
198 self.free_ranges[index].end = aligned_start;
199 self.free_ranges.insert(index + 1, alloc_end..range.end);
200 }
201 }
202
203 Ok(aligned_start..alloc_end)
204 }
205 None => Err(RangeAllocationError {
206 fragmented_free_length,
207 }),
208 }
209 }
210
211 pub fn allocate_range(&mut self, length: T) -> Result<Range<T>, RangeAllocationError<T>> {
232 self.allocate_range_impl(length, |start| start)
233 }
234
235 pub fn allocate_range_aligned(
266 &mut self,
267 length: T,
268 alignment: T,
269 ) -> Result<Range<T>, RangeAllocationError<T>>
270 where
271 T: Rem<Output = T>,
272 {
273 assert_ne!(alignment + alignment, alignment);
274 self.allocate_range_impl(length, |start| {
275 let padding = (alignment - start % alignment) % alignment;
276 start + padding
277 })
278 }
279
280 pub fn free_range(&mut self, range: Range<T>) {
301 assert!(self.initial_range.start <= range.start && range.end <= self.initial_range.end);
302 assert!(range.start < range.end);
303
304 let i = self
306 .free_ranges
307 .iter()
308 .position(|r| r.start > range.start)
309 .unwrap_or(self.free_ranges.len());
310
311 if i > 0 && range.start == self.free_ranges[i - 1].end {
314 self.free_ranges[i - 1].end =
316 if i < self.free_ranges.len() && range.end == self.free_ranges[i].start {
317 let right = self.free_ranges.remove(i);
319 right.end
320 } else {
321 range.end
322 };
323
324 return;
325 } else if i < self.free_ranges.len() && range.end == self.free_ranges[i].start {
326 self.free_ranges[i].start = if i > 0 && range.start == self.free_ranges[i - 1].end {
328 let left = self.free_ranges.remove(i - 1);
330 left.start
331 } else {
332 range.start
333 };
334
335 return;
336 }
337
338 assert!(
340 (i == 0 || self.free_ranges[i - 1].end < range.start)
341 && (i >= self.free_ranges.len() || range.end < self.free_ranges[i].start)
342 );
343
344 self.free_ranges.insert(i, range);
345 }
346
347 pub fn allocated_ranges(&self) -> impl Iterator<Item = Range<T>> + '_ {
365 let first = match self.free_ranges.first() {
366 Some(Range { ref start, .. }) if *start > self.initial_range.start => {
367 Some(self.initial_range.start..*start)
368 }
369 None => Some(self.initial_range.clone()),
370 _ => None,
371 };
372
373 let last = match self.free_ranges.last() {
374 Some(Range { end, .. }) if *end < self.initial_range.end => {
375 Some(*end..self.initial_range.end)
376 }
377 _ => None,
378 };
379
380 let mid = self
381 .free_ranges
382 .iter()
383 .zip(self.free_ranges.iter().skip(1))
384 .map(|(ra, rb)| ra.end..rb.start);
385
386 first.into_iter().chain(mid).chain(last)
387 }
388
389 pub fn reset(&mut self) {
402 self.free_ranges.clear();
403 self.free_ranges.push(self.initial_range.clone());
404 }
405
406 pub fn is_empty(&self) -> bool {
408 self.free_ranges.len() == 1 && self.free_ranges[0] == self.initial_range
409 }
410}
411
412impl<T: Copy + Sub<Output = T> + Sum> RangeAllocator<T> {
413 pub fn total_available(&self) -> T {
428 self.free_ranges
429 .iter()
430 .map(|range| range.end - range.start)
431 .sum()
432 }
433}
434
435#[cfg(test)]
436mod tests {
437 use super::*;
438
439 #[test]
440 fn test_basic_allocation() {
441 let mut alloc = RangeAllocator::new(0..10);
442 assert_eq!(alloc.allocate_range(4), Ok(0..4));
444 assert!(alloc.allocated_ranges().eq(std::iter::once(0..4)));
445 alloc.free_range(0..4);
447 assert_eq!(alloc.free_ranges, vec![0..10]);
449 assert!(alloc.allocated_ranges().eq(std::iter::empty()));
450 }
451
452 #[test]
453 fn test_out_of_space() {
454 let mut alloc = RangeAllocator::new(0..10);
455 assert_eq!(alloc.allocate_range(10), Ok(0..10));
457 assert!(alloc.allocated_ranges().eq(std::iter::once(0..10)));
458 assert!(alloc.allocate_range(4).is_err());
459 alloc.free_range(0..10);
460 }
461
462 #[test]
463 fn test_grow() {
464 let mut alloc = RangeAllocator::new(0..11);
465 assert_eq!(alloc.allocate_range(10), Ok(0..10));
467 assert!(alloc.allocated_ranges().eq(std::iter::once(0..10)));
468 assert!(alloc.allocate_range(4).is_err());
469 alloc.grow_to(20);
470 assert_eq!(alloc.allocate_range(4), Ok(10..14));
471 alloc.free_range(0..14);
472 }
473
474 #[test]
475 fn test_grow_with_hole_at_start() {
476 let mut alloc = RangeAllocator::new(0..6);
477
478 assert_eq!(alloc.allocate_range(3), Ok(0..3));
479 assert_eq!(alloc.allocate_range(3), Ok(3..6));
480 alloc.free_range(0..3);
481
482 alloc.grow_to(9);
483 assert_eq!(alloc.allocated_ranges().collect::<Vec<_>>(), [3..6]);
484 }
485 #[test]
486 fn test_grow_with_hole_in_middle() {
487 let mut alloc = RangeAllocator::new(0..6);
488
489 assert_eq!(alloc.allocate_range(2), Ok(0..2));
490 assert_eq!(alloc.allocate_range(2), Ok(2..4));
491 assert_eq!(alloc.allocate_range(2), Ok(4..6));
492 alloc.free_range(2..4);
493
494 alloc.grow_to(9);
495 assert_eq!(alloc.allocated_ranges().collect::<Vec<_>>(), [0..2, 4..6]);
496 }
497
498 #[test]
499 fn test_dont_use_block_that_is_too_small() {
500 let mut alloc = RangeAllocator::new(0..10);
501 assert_eq!(alloc.allocate_range(3), Ok(0..3));
503 assert_eq!(alloc.allocate_range(3), Ok(3..6));
504 assert_eq!(alloc.allocate_range(3), Ok(6..9));
505 alloc.free_range(3..6);
506 assert_eq!(alloc.free_ranges, vec![3..6, 9..10]);
507 assert_eq!(
508 alloc.allocated_ranges().collect::<Vec<Range<i32>>>(),
509 vec![0..3, 6..9]
510 );
511 assert_eq!(alloc.allocate_range(3), Ok(3..6));
513 }
514
515 #[test]
516 fn test_free_blocks_in_middle() {
517 let mut alloc = RangeAllocator::new(0..100);
518 assert_eq!(alloc.allocate_range(10), Ok(0..10));
520 assert_eq!(alloc.allocate_range(10), Ok(10..20));
521 assert_eq!(alloc.allocate_range(10), Ok(20..30));
522 assert_eq!(alloc.allocate_range(10), Ok(30..40));
523 assert_eq!(alloc.allocate_range(10), Ok(40..50));
524 assert_eq!(alloc.allocate_range(10), Ok(50..60));
525 assert_eq!(alloc.allocate_range(10), Ok(60..70));
526 assert_eq!(alloc.allocate_range(10), Ok(70..80));
527 assert_eq!(alloc.allocate_range(10), Ok(80..90));
528 assert_eq!(alloc.allocate_range(10), Ok(90..100));
529 assert_eq!(alloc.free_ranges, vec![]);
530 assert!(alloc.allocated_ranges().eq(std::iter::once(0..100)));
531 alloc.free_range(10..20);
532 alloc.free_range(30..40);
533 alloc.free_range(50..60);
534 alloc.free_range(70..80);
535 alloc.free_range(90..100);
536 assert_eq!(
538 alloc.free_ranges,
539 vec![10..20, 30..40, 50..60, 70..80, 90..100]
540 );
541 assert_eq!(
542 alloc.allocated_ranges().collect::<Vec<Range<i32>>>(),
543 vec![0..10, 20..30, 40..50, 60..70, 80..90]
544 );
545 assert_eq!(alloc.allocate_range(6), Ok(10..16));
547 assert_eq!(alloc.allocate_range(6), Ok(30..36));
548 assert_eq!(alloc.allocate_range(6), Ok(50..56));
549 assert_eq!(alloc.allocate_range(6), Ok(70..76));
550 assert_eq!(alloc.allocate_range(6), Ok(90..96));
551 assert_eq!(
553 alloc.free_ranges,
554 vec![16..20, 36..40, 56..60, 76..80, 96..100]
555 );
556 assert_eq!(
557 alloc.allocated_ranges().collect::<Vec<Range<i32>>>(),
558 vec![0..16, 20..36, 40..56, 60..76, 80..96]
559 );
560 assert_eq!(alloc.allocate_range(4), Ok(16..20));
562 assert_eq!(alloc.allocate_range(4), Ok(36..40));
563 assert_eq!(alloc.allocate_range(4), Ok(56..60));
564 assert_eq!(alloc.allocate_range(4), Ok(76..80));
565 assert_eq!(alloc.allocate_range(4), Ok(96..100));
566 assert_eq!(alloc.free_ranges, vec![]);
568 assert!(alloc.allocated_ranges().eq(std::iter::once(0..100)));
569 }
570
571 #[test]
572 fn test_ignore_block_if_another_fits_better() {
573 let mut alloc = RangeAllocator::new(0..10);
574 assert_eq!(alloc.allocate_range(3), Ok(0..3));
577 assert_eq!(alloc.allocate_range(3), Ok(3..6));
578 assert_eq!(alloc.allocate_range(3), Ok(6..9));
579 alloc.free_range(3..6);
580 assert_eq!(alloc.free_ranges, vec![3..6, 9..10]);
581 assert_eq!(
582 alloc.allocated_ranges().collect::<Vec<Range<i32>>>(),
583 vec![0..3, 6..9]
584 );
585 assert_eq!(alloc.allocate_range(1), Ok(9..10));
588 }
589
590 #[test]
591 fn test_merge_neighbors() {
592 let mut alloc = RangeAllocator::new(0..9);
593 assert_eq!(alloc.allocate_range(3), Ok(0..3));
594 assert_eq!(alloc.allocate_range(3), Ok(3..6));
595 assert_eq!(alloc.allocate_range(3), Ok(6..9));
596 alloc.free_range(0..3);
597 alloc.free_range(6..9);
598 alloc.free_range(3..6);
599 assert_eq!(alloc.free_ranges, vec![0..9]);
600 assert!(alloc.allocated_ranges().eq(std::iter::empty()));
601 }
602
603 #[test]
604 fn test_aligned_already_aligned() {
605 let mut alloc = RangeAllocator::new(0..20);
606 assert_eq!(alloc.allocate_range_aligned(4, 4), Ok(0..4));
608 assert_eq!(alloc.free_ranges, vec![4..20]);
609 }
610
611 #[test]
612 fn test_aligned_with_padding() {
613 let mut alloc = RangeAllocator::new(0..20);
614 assert_eq!(alloc.allocate_range(1), Ok(0..1));
616 assert_eq!(alloc.allocate_range_aligned(4, 4), Ok(4..8));
618 assert_eq!(alloc.free_ranges, vec![1..4, 8..20]);
620 }
621
622 #[test]
623 fn test_aligned_prefix_is_reusable() {
624 let mut alloc = RangeAllocator::new(0..20);
625 assert_eq!(alloc.allocate_range(1), Ok(0..1));
626 assert_eq!(alloc.allocate_range_aligned(4, 4), Ok(4..8));
627 assert_eq!(alloc.allocate_range(3), Ok(1..4));
629 assert_eq!(alloc.free_ranges, vec![8..20]);
630 }
631
632 #[test]
633 fn test_aligned_no_fit() {
634 let mut alloc = RangeAllocator::new(0..5);
635 assert_eq!(alloc.allocate_range(1), Ok(0..1));
636 assert!(alloc.allocate_range_aligned(4, 4).is_err());
638 }
639
640 #[test]
641 fn test_aligned_exact_fit_after_padding() {
642 let mut alloc = RangeAllocator::new(0..8);
643 assert_eq!(alloc.allocate_range(1), Ok(0..1));
644 assert_eq!(alloc.allocate_range_aligned(4, 4), Ok(4..8));
646 assert_eq!(alloc.free_ranges, vec![1..4]);
648 }
649
650 #[test]
651 fn test_aligned_best_fit() {
652 let mut alloc = RangeAllocator::new(0..32);
653 assert_eq!(alloc.allocate_range(4), Ok(0..4));
655 assert_eq!(alloc.allocate_range(4), Ok(4..8));
656 assert_eq!(alloc.allocate_range(4), Ok(8..12));
657 assert_eq!(alloc.allocate_range(4), Ok(12..16));
658 assert_eq!(alloc.allocate_range(4), Ok(16..20));
659 assert_eq!(alloc.allocate_range(12), Ok(20..32));
660
661 alloc.free_range(4..8);
664 alloc.free_range(12..20);
665 assert_eq!(alloc.free_ranges, vec![4..8, 12..20]);
666
667 assert_eq!(alloc.allocate_range_aligned(4, 4), Ok(4..8));
669 }
670
671 #[test]
672 fn test_aligned_non_power_of_two() {
673 let mut alloc = RangeAllocator::new(0..20);
674 assert_eq!(alloc.allocate_range(1), Ok(0..1));
675 assert_eq!(alloc.allocate_range_aligned(2, 3), Ok(3..5));
677 assert_eq!(alloc.free_ranges, vec![1..3, 5..20]);
678 }
679
680 #[test]
681 fn test_aligned_multiple_allocations() {
682 let mut alloc = RangeAllocator::new(0..32);
683 assert_eq!(alloc.allocate_range_aligned(4, 8), Ok(0..4));
684 assert_eq!(alloc.allocate_range_aligned(4, 8), Ok(8..12));
686 assert_eq!(alloc.allocate_range_aligned(4, 8), Ok(16..20));
688 assert_eq!(alloc.free_ranges, vec![4..8, 12..16, 20..32]);
690 }
691
692 #[test]
693 fn test_aligned_allocation_then_free_merges() {
694 let mut alloc = RangeAllocator::new(0..16);
695 assert_eq!(alloc.allocate_range(1), Ok(0..1));
696 assert_eq!(alloc.allocate_range_aligned(4, 4), Ok(4..8));
697 alloc.free_range(4..8);
700 assert_eq!(alloc.free_ranges, vec![1..16]);
702 }
703
704 #[test]
705 fn test_allocate_range_delegates_correctly() {
706 let mut alloc = RangeAllocator::new(0..10);
708 assert_eq!(alloc.allocate_range(4), Ok(0..4));
709 assert_eq!(alloc.allocate_range(3), Ok(4..7));
710 assert_eq!(alloc.free_ranges, vec![7..10]);
711 alloc.free_range(0..4);
712 assert_eq!(alloc.free_ranges, vec![0..4, 7..10]);
713 }
714}