1use std::cmp;
6use std::collections::BTreeSet;
7use std::collections::HashMap;
8use std::ops::Bound;
9
10use crate::AddressRange;
11use crate::Alloc;
12use crate::Error;
13use crate::Result;
14
15#[derive(Debug, Eq, PartialEq)]
20pub struct AddressAllocator {
21 pools: Vec<AddressRange>,
24 min_align: u64,
25 preferred_align: u64,
26 allocs: HashMap<Alloc, (AddressRange, String)>,
28 regions: BTreeSet<AddressRange>,
30}
31
32impl AddressAllocator {
33 pub fn new(
43 pool: AddressRange,
44 min_align: Option<u64>,
45 preferred_align: Option<u64>,
46 ) -> Result<Self> {
47 Self::new_from_list(vec![pool], min_align, preferred_align)
48 }
49
50 pub fn new_from_list<T>(
60 pools: T,
61 min_align: Option<u64>,
62 preferred_align: Option<u64>,
63 ) -> Result<Self>
64 where
65 T: IntoIterator<Item = AddressRange>,
66 {
67 let pools: Vec<AddressRange> = pools.into_iter().filter(|p| !p.is_empty()).collect();
68
69 let min_align = min_align.unwrap_or(4);
70 if !min_align.is_power_of_two() || min_align == 0 {
71 return Err(Error::BadAlignment);
72 }
73
74 let preferred_align = preferred_align.unwrap_or(min_align);
75 if !preferred_align.is_power_of_two() || preferred_align < min_align {
76 return Err(Error::BadAlignment);
77 }
78
79 let mut regions = BTreeSet::new();
80 for r in pools.iter() {
81 regions.insert(*r);
82 }
83 Ok(AddressAllocator {
84 pools,
85 min_align,
86 preferred_align,
87 allocs: HashMap::new(),
88 regions,
89 })
90 }
91
92 pub fn pools(&self) -> &[AddressRange] {
96 &self.pools
97 }
98
99 fn internal_allocate_from_slot(
100 &mut self,
101 slot: AddressRange,
102 range: AddressRange,
103 alloc: Alloc,
104 tag: String,
105 ) -> Result<u64> {
106 let slot_was_present = self.regions.remove(&slot);
107 assert!(slot_was_present);
108
109 let (before, after) = slot.non_overlapping_ranges(range);
110
111 if !before.is_empty() {
112 self.regions.insert(before);
113 }
114 if !after.is_empty() {
115 self.regions.insert(after);
116 }
117
118 self.allocs.insert(alloc, (range, tag));
119 Ok(range.start)
120 }
121
122 fn internal_allocate_with_align(
123 &mut self,
124 size: u64,
125 alloc: Alloc,
126 tag: String,
127 alignment: u64,
128 reverse: bool,
129 ) -> Result<u64> {
130 let alignment = cmp::max(self.min_align, alignment);
131
132 if self.allocs.contains_key(&alloc) {
133 return Err(Error::ExistingAlloc(alloc));
134 }
135 if size == 0 {
136 return Err(Error::AllocSizeZero);
137 }
138 if !alignment.is_power_of_two() {
139 return Err(Error::BadAlignment);
140 }
141
142 let region = if !reverse {
143 self.regions
145 .iter()
146 .find(|range| {
147 match range.start % alignment {
148 0 => range.start.checked_add(size - 1),
149 r => range
150 .start
151 .checked_add(size - 1)
152 .and_then(|start| start.checked_add(alignment - r)),
153 }
154 .is_some_and(|end| end <= range.end)
155 })
156 .cloned()
157 } else {
158 self.regions
160 .iter()
161 .rev()
162 .find(|range| {
163 range
164 .end
165 .checked_sub(size - 1)
166 .is_some_and(|start| start & !(alignment - 1) >= range.start)
167 })
168 .cloned()
169 };
170
171 match region {
172 Some(slot) => {
173 let start = if !reverse {
174 match slot.start % alignment {
175 0 => slot.start,
176 r => slot.start + (alignment - r),
177 }
178 } else {
179 (slot.end - (size - 1)) & !(alignment - 1)
180 };
181 let end = start + (size - 1);
182 let range = AddressRange { start, end };
183
184 self.internal_allocate_from_slot(slot, range, alloc, tag)
185 }
186 None => Err(Error::OutOfSpace),
187 }
188 }
189
190 pub fn reverse_allocate_with_align(
194 &mut self,
195 size: u64,
196 alloc: Alloc,
197 tag: String,
198 alignment: u64,
199 ) -> Result<u64> {
200 self.internal_allocate_with_align(size, alloc, tag, alignment, true)
201 }
202
203 pub fn reverse_allocate(&mut self, size: u64, alloc: Alloc, tag: String) -> Result<u64> {
205 if let Ok(pref_alloc) =
206 self.reverse_allocate_with_align(size, alloc, tag.clone(), self.preferred_align)
207 {
208 return Ok(pref_alloc);
209 }
210 self.reverse_allocate_with_align(size, alloc, tag, self.min_align)
211 }
212
213 pub fn allocate_with_align(
217 &mut self,
218 size: u64,
219 alloc: Alloc,
220 tag: String,
221 alignment: u64,
222 ) -> Result<u64> {
223 self.internal_allocate_with_align(size, alloc, tag, alignment, false)
224 }
225
226 pub fn allocate(&mut self, size: u64, alloc: Alloc, tag: String) -> Result<u64> {
227 if let Ok(pref_alloc) =
228 self.allocate_with_align(size, alloc, tag.clone(), self.preferred_align)
229 {
230 return Ok(pref_alloc);
231 }
232 self.allocate_with_align(size, alloc, tag, self.min_align)
233 }
234
235 pub fn allocate_at(&mut self, range: AddressRange, alloc: Alloc, tag: String) -> Result<()> {
240 if self.allocs.contains_key(&alloc) {
241 return Err(Error::ExistingAlloc(alloc));
242 }
243
244 if range.is_empty() {
245 return Err(Error::AllocSizeZero);
246 }
247
248 match self
249 .regions
250 .iter()
251 .find(|avail_range| avail_range.contains_range(range))
252 {
253 Some(&slot) => {
254 let _address = self.internal_allocate_from_slot(slot, range, alloc, tag)?;
255 Ok(())
256 }
257 None => {
258 if let Some(existing_alloc) = self.find_overlapping(range) {
259 Err(Error::ExistingAlloc(existing_alloc))
260 } else {
261 Err(Error::OutOfSpace)
262 }
263 }
264 }
265 }
266
267 pub fn release(&mut self, alloc: Alloc) -> Result<AddressRange> {
269 if let Some((range, _tag)) = self.allocs.remove(&alloc) {
270 self.insert_at(range)?;
271 Ok(range)
272 } else {
273 Err(Error::BadAlloc(alloc))
274 }
275 }
276
277 pub fn release_containing(&mut self, value: u64) -> Result<AddressRange> {
279 if let Some(alloc) = self.find_overlapping(AddressRange {
280 start: value,
281 end: value,
282 }) {
283 self.release(alloc)
284 } else {
285 Err(Error::OutOfSpace)
286 }
287 }
288
289 fn find_overlapping(&self, range: AddressRange) -> Option<Alloc> {
293 if range.is_empty() {
294 return None;
295 }
296
297 self.allocs
298 .iter()
299 .find(|(_, &(alloc_range, _))| alloc_range.overlaps(range))
300 .map(|(&alloc, _)| alloc)
301 }
302
303 pub fn get_max_addr(&self) -> u64 {
305 self.regions.iter().fold(0, |x, range| x.max(range.end))
306 }
307
308 pub fn get(&self, alloc: &Alloc) -> Option<&(AddressRange, String)> {
310 self.allocs.get(alloc)
311 }
312
313 fn insert_at(&mut self, mut slot: AddressRange) -> Result<()> {
315 if slot.is_empty() {
316 return Err(Error::AllocSizeZero);
317 }
318
319 let mut smaller_merge = None;
323 if let Some(smaller) = self
324 .regions
325 .range((Bound::Unbounded, Bound::Included(slot)))
326 .max()
327 {
328 let next_addr = smaller
330 .end
331 .checked_add(1)
332 .ok_or(Error::RegionOverlap(slot))?;
333 match next_addr.cmp(&slot.start) {
334 cmp::Ordering::Less => (),
335 cmp::Ordering::Equal => smaller_merge = Some(*smaller),
336 cmp::Ordering::Greater => return Err(Error::RegionOverlap(slot)),
337 }
338 }
339
340 let mut larger_merge = None;
344 if let Some(larger) = self
345 .regions
346 .range((Bound::Excluded(slot), Bound::Unbounded))
347 .min()
348 {
349 let prev_addr = larger
351 .start
352 .checked_sub(1)
353 .ok_or(Error::RegionOverlap(slot))?;
354 match slot.end.cmp(&prev_addr) {
355 cmp::Ordering::Less => (),
356 cmp::Ordering::Equal => larger_merge = Some(*larger),
357 cmp::Ordering::Greater => return Err(Error::RegionOverlap(slot)),
358 }
359 }
360
361 if let Some(smaller) = smaller_merge {
362 self.regions.remove(&smaller);
363 slot.start = smaller.start;
364 }
365 if let Some(larger) = larger_merge {
366 self.regions.remove(&larger);
367 slot.end = larger.end;
368 }
369 self.regions.insert(slot);
370
371 Ok(())
372 }
373
374 pub fn address_from_pci_offset(&self, alloc: Alloc, offset: u64, size: u64) -> Result<u64> {
376 match alloc {
377 Alloc::PciBar { .. } => (),
378 _ => return Err(Error::InvalidAlloc(alloc)),
379 };
380
381 match self.allocs.get(&alloc) {
382 Some((pci_bar_range, _)) => {
383 let address = pci_bar_range
384 .start
385 .checked_add(offset)
386 .ok_or(Error::OutOfBounds)?;
387 let offset_range =
388 AddressRange::from_start_and_size(address, size).ok_or(Error::OutOfBounds)?;
389 if pci_bar_range.contains_range(offset_range) {
390 Ok(address)
391 } else {
392 Err(Error::OutOfBounds)
393 }
394 }
395 None => Err(Error::InvalidAlloc(alloc)),
396 }
397 }
398}
399
400pub struct AddressAllocatorSet<'a> {
405 allocators: &'a mut [AddressAllocator],
406}
407
408impl<'a> AddressAllocatorSet<'a> {
409 pub fn new(allocators: &'a mut [AddressAllocator]) -> Self {
410 AddressAllocatorSet { allocators }
411 }
412
413 pub fn allocate_with_align(
414 &mut self,
415 size: u64,
416 alloc: Alloc,
417 tag: String,
418 alignment: u64,
419 ) -> Result<u64> {
420 let mut last_res = Err(Error::OutOfSpace);
421 for allocator in self.allocators.iter_mut() {
422 last_res = allocator.allocate_with_align(size, alloc, tag.clone(), alignment);
423 if last_res.is_ok() {
424 return last_res;
425 }
426 }
427 last_res
428 }
429
430 pub fn allocate(&mut self, size: u64, alloc: Alloc, tag: String) -> Result<u64> {
431 let mut last_res = Err(Error::OutOfSpace);
432 for allocator in self.allocators.iter_mut() {
433 last_res = allocator.allocate(size, alloc, tag.clone());
434 if last_res.is_ok() {
435 return last_res;
436 }
437 }
438 last_res
439 }
440
441 pub fn allocate_at(&mut self, range: AddressRange, alloc: Alloc, tag: String) -> Result<()> {
442 let mut last_res = Err(Error::OutOfSpace);
443 for allocator in self.allocators.iter_mut() {
444 last_res = allocator.allocate_at(range, alloc, tag.clone());
445 if last_res.is_ok() {
446 return last_res;
447 }
448 }
449 last_res
450 }
451
452 pub fn release(&mut self, alloc: Alloc) -> Result<AddressRange> {
453 let mut last_res = Err(Error::OutOfSpace);
454 for allocator in self.allocators.iter_mut() {
455 last_res = allocator.release(alloc);
456 if last_res.is_ok() {
457 return last_res;
458 }
459 }
460 last_res
461 }
462
463 pub fn get(&self, alloc: &Alloc) -> Option<&(AddressRange, String)> {
464 for allocator in self.allocators.iter() {
465 let opt = allocator.get(alloc);
466 if opt.is_some() {
467 return opt;
468 }
469 }
470 None
471 }
472
473 pub fn address_from_pci_offset(&self, alloc: Alloc, offset: u64, size: u64) -> Result<u64> {
474 let mut last_res = Err(Error::OutOfSpace);
475 for allocator in self.allocators.iter() {
476 last_res = allocator.address_from_pci_offset(alloc, offset, size);
477 if last_res.is_ok() {
478 return last_res;
479 }
480 }
481 last_res
482 }
483}
484
485#[cfg(test)]
486mod tests {
487 use super::*;
488
489 #[test]
490 fn example() {
491 let mut pool = AddressAllocator::new(
493 AddressRange {
494 start: 0x1000,
495 end: 0xFFFF,
496 },
497 Some(0x100),
498 None,
499 )
500 .unwrap();
501 assert_eq!(
502 pool.allocate(0x110, Alloc::Anon(0), "caps".to_string()),
503 Ok(0x1000)
504 );
505 assert_eq!(
506 pool.allocate(0x100, Alloc::Anon(1), "cache".to_string()),
507 Ok(0x1200)
508 );
509 assert_eq!(
510 pool.allocate(0x100, Alloc::Anon(2), "etc".to_string()),
511 Ok(0x1300)
512 );
513 assert_eq!(
514 pool.get(&Alloc::Anon(1)),
515 Some(&(
516 AddressRange {
517 start: 0x1200,
518 end: 0x12FF
519 },
520 "cache".to_string()
521 ))
522 );
523 }
524
525 #[test]
526 fn empty_allocator() {
527 let mut pool = AddressAllocator::new_from_list(Vec::new(), None, None).unwrap();
528 assert_eq!(pool.pools(), &[]);
529 assert_eq!(
530 pool.allocate(1, Alloc::Anon(0), "test".to_string()),
531 Err(Error::OutOfSpace)
532 );
533 }
534
535 #[test]
536 fn new_fails_alignment_zero() {
537 assert!(AddressAllocator::new(
538 AddressRange {
539 start: 0x1000,
540 end: 0xFFFF
541 },
542 Some(0),
543 None
544 )
545 .is_err());
546 }
547
548 #[test]
549 fn new_fails_alignment_non_power_of_two() {
550 assert!(AddressAllocator::new(
551 AddressRange {
552 start: 0x1000,
553 end: 0xFFFF
554 },
555 Some(200),
556 None
557 )
558 .is_err());
559 }
560
561 #[test]
562 fn allocate_fails_exising_alloc() {
563 let mut pool = AddressAllocator::new(
564 AddressRange {
565 start: 0x1000,
566 end: 0x1FFF,
567 },
568 Some(0x100),
569 None,
570 )
571 .unwrap();
572 assert_eq!(
573 pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
574 Ok(0x1000)
575 );
576 assert_eq!(
577 pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
578 Err(Error::ExistingAlloc(Alloc::Anon(0)))
579 );
580 }
581
582 #[test]
583 fn allocate_fails_not_enough_space() {
584 let mut pool = AddressAllocator::new(
585 AddressRange {
586 start: 0x1000,
587 end: 0x1FFF,
588 },
589 Some(0x100),
590 None,
591 )
592 .unwrap();
593 assert_eq!(
594 pool.allocate(0x800, Alloc::Anon(0), String::from("bar0")),
595 Ok(0x1000)
596 );
597 assert_eq!(
598 pool.allocate(0x900, Alloc::Anon(1), String::from("bar1")),
599 Err(Error::OutOfSpace)
600 );
601 assert_eq!(
602 pool.allocate(0x800, Alloc::Anon(2), String::from("bar2")),
603 Ok(0x1800)
604 );
605 }
606
607 #[test]
608 fn allocate_with_special_alignment() {
609 let mut pool = AddressAllocator::new(
610 AddressRange {
611 start: 0x1000,
612 end: 0x1FFF,
613 },
614 Some(0x100),
615 None,
616 )
617 .unwrap();
618 assert_eq!(
619 pool.allocate(0x10, Alloc::Anon(0), String::from("bar0")),
620 Ok(0x1000)
621 );
622 assert_eq!(
623 pool.allocate_at(
624 AddressRange {
625 start: 0x1200,
626 end: 0x13ff,
627 },
628 Alloc::Anon(1),
629 String::from("bar1")
630 ),
631 Ok(())
632 );
633 assert_eq!(
634 pool.allocate_with_align(0x800, Alloc::Anon(2), String::from("bar2"), 0x800),
635 Ok(0x1800)
636 );
637 }
638
639 #[test]
640 fn allocate_and_split_allocate_at() {
641 let mut pool = AddressAllocator::new(
642 AddressRange {
643 start: 0x1000,
644 end: 0x1fff,
645 },
646 Some(1),
647 None,
648 )
649 .unwrap();
650 assert_eq!(
652 pool.allocate_at(
653 AddressRange {
654 start: 0x1200,
655 end: 0x19ff,
656 },
657 Alloc::Anon(0),
658 String::from("bar0")
659 ),
660 Ok(())
661 );
662 assert_eq!(
663 pool.allocate(0x800, Alloc::Anon(1), String::from("bar1")),
664 Err(Error::OutOfSpace)
665 );
666 assert_eq!(
668 pool.allocate(0x600, Alloc::Anon(2), String::from("bar2")),
669 Ok(0x1a00)
670 );
671 assert_eq!(
673 pool.allocate(0x200, Alloc::Anon(3), String::from("bar3")),
674 Ok(0x1000)
675 );
676 assert_eq!(
678 pool.allocate_at(
679 AddressRange {
680 start: 0x1b00,
681 end: 0x1bff,
682 },
683 Alloc::Anon(4),
684 String::from("bar4")
685 ),
686 Err(Error::ExistingAlloc(Alloc::Anon(2)))
687 );
688 assert_eq!(
690 pool.allocate_at(
691 AddressRange {
692 start: 0x1fff,
693 end: 0x1fff,
694 },
695 Alloc::Anon(5),
696 String::from("bar5")
697 ),
698 Err(Error::ExistingAlloc(Alloc::Anon(2)))
699 );
700 assert_eq!(
702 pool.allocate_at(
703 AddressRange {
704 start: 0x1200,
705 end: 0x1200,
706 },
707 Alloc::Anon(6),
708 String::from("bar6")
709 ),
710 Err(Error::ExistingAlloc(Alloc::Anon(0)))
711 );
712 assert_eq!(
714 pool.allocate_at(
715 AddressRange {
716 start: 0x11ff,
717 end: 0x11ff,
718 },
719 Alloc::Anon(7),
720 String::from("bar7")
721 ),
722 Err(Error::ExistingAlloc(Alloc::Anon(3)))
723 );
724 match pool.allocate_at(
726 AddressRange {
727 start: 0x1100,
728 end: 0x12ff,
729 },
730 Alloc::Anon(8),
731 String::from("bar8"),
732 ) {
733 Err(Error::ExistingAlloc(Alloc::Anon(0) | Alloc::Anon(3))) => {}
734 x => panic!("unexpected result {x:?}"),
735 }
736 }
737
738 #[test]
739 fn allocate_alignment() {
740 let mut pool = AddressAllocator::new(
741 AddressRange {
742 start: 0x1000,
743 end: 0xFFFF,
744 },
745 Some(0x100),
746 None,
747 )
748 .unwrap();
749 assert_eq!(
750 pool.allocate(0x110, Alloc::Anon(0), String::from("bar0")),
751 Ok(0x1000)
752 );
753 assert_eq!(
754 pool.allocate(0x100, Alloc::Anon(1), String::from("bar1")),
755 Ok(0x1200)
756 );
757 }
758
759 #[test]
760 fn allocate_retrieve_alloc() {
761 let mut pool = AddressAllocator::new(
762 AddressRange {
763 start: 0x1000,
764 end: 0xFFFF,
765 },
766 Some(0x100),
767 None,
768 )
769 .unwrap();
770 assert_eq!(
771 pool.allocate(0x110, Alloc::Anon(0), String::from("bar0")),
772 Ok(0x1000)
773 );
774 assert_eq!(
775 pool.get(&Alloc::Anon(0)),
776 Some(&(
777 AddressRange {
778 start: 0x1000,
779 end: 0x110f,
780 },
781 String::from("bar0")
782 ))
783 );
784 }
785
786 #[test]
787 fn allocate_with_alignment_allocator_alignment() {
788 let mut pool = AddressAllocator::new(
789 AddressRange {
790 start: 0x1000,
791 end: 0xFFFF,
792 },
793 Some(0x100),
794 None,
795 )
796 .unwrap();
797 assert_eq!(
798 pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x1),
799 Ok(0x1000)
800 );
801 assert_eq!(
802 pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x1),
803 Ok(0x1200)
804 );
805 }
806
807 #[test]
808 fn allocate_with_alignment_custom_alignment() {
809 let mut pool = AddressAllocator::new(
810 AddressRange {
811 start: 0x1000,
812 end: 0xFFFF,
813 },
814 Some(0x4),
815 None,
816 )
817 .unwrap();
818 assert_eq!(
819 pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x100),
820 Ok(0x1000)
821 );
822 assert_eq!(
823 pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x100),
824 Ok(0x1200)
825 );
826 }
827
828 #[test]
829 fn allocate_with_alignment_no_allocator_alignment() {
830 let mut pool = AddressAllocator::new(
831 AddressRange {
832 start: 0x1000,
833 end: 0xFFFF,
834 },
835 None,
836 None,
837 )
838 .unwrap();
839 assert_eq!(
840 pool.allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 0x100),
841 Ok(0x1000)
842 );
843 assert_eq!(
844 pool.allocate_with_align(0x100, Alloc::Anon(1), String::from("bar1"), 0x100),
845 Ok(0x1200)
846 );
847 }
848
849 #[test]
850 fn allocate_with_alignment_alignment_non_power_of_two() {
851 let mut pool = AddressAllocator::new(
852 AddressRange {
853 start: 0x1000,
854 end: 0xFFFF,
855 },
856 None,
857 None,
858 )
859 .unwrap();
860 assert!(pool
861 .allocate_with_align(0x110, Alloc::Anon(0), String::from("bar0"), 200)
862 .is_err());
863 }
864
865 #[test]
866 fn allocate_with_release() {
867 let mut pool = AddressAllocator::new(
868 AddressRange {
869 start: 0x1000,
870 end: 0x1FFF,
871 },
872 None,
873 None,
874 )
875 .unwrap();
876 assert_eq!(
877 pool.allocate_with_align(0x100, Alloc::Anon(0), String::from("bar0"), 0x100),
878 Ok(0x1000)
879 );
880 assert!(pool.release(Alloc::Anon(0)).is_ok());
881 assert_eq!(
882 pool.allocate_with_align(0x1000, Alloc::Anon(0), String::from("bar0"), 0x100),
883 Ok(0x1000)
884 );
885 }
886
887 #[test]
888 fn coalescing_and_overlap() {
889 let mut pool = AddressAllocator::new(
890 AddressRange {
891 start: 0x1000,
892 end: 0x1FFF,
893 },
894 None,
895 None,
896 )
897 .unwrap();
898 assert!(pool
899 .insert_at(AddressRange {
900 start: 0x3000,
901 end: 0x3fff,
902 })
903 .is_ok());
904 assert!(pool
905 .insert_at(AddressRange {
906 start: 0x1fff,
907 end: 0x201e,
908 })
909 .is_err());
910 assert!(pool
911 .insert_at(AddressRange {
912 start: 0x2ff1,
913 end: 0x3000,
914 })
915 .is_err());
916 assert!(pool
917 .insert_at(AddressRange {
918 start: 0x1800,
919 end: 0x27ff,
920 })
921 .is_err());
922 assert!(pool
923 .insert_at(AddressRange {
924 start: 0x2000,
925 end: 0x2fff,
926 })
927 .is_ok());
928 assert_eq!(
929 pool.allocate(0x3000, Alloc::Anon(0), String::from("bar0")),
930 Ok(0x1000)
931 );
932 }
933
934 #[test]
935 fn coalescing_single_addresses() {
936 let mut pool = AddressAllocator::new(
937 AddressRange {
938 start: 0x1000,
939 end: 0x1FFF,
940 },
941 None,
942 None,
943 )
944 .unwrap();
945 assert!(pool
946 .insert_at(AddressRange {
947 start: 0x2001,
948 end: 0x2001,
949 })
950 .is_ok());
951 assert!(pool
952 .insert_at(AddressRange {
953 start: 0x2003,
954 end: 0x2003,
955 })
956 .is_ok());
957 assert!(pool
958 .insert_at(AddressRange {
959 start: 0x2000,
960 end: 0x2000,
961 })
962 .is_ok());
963 assert!(pool
964 .insert_at(AddressRange {
965 start: 0x2002,
966 end: 0x2002,
967 })
968 .is_ok());
969 assert_eq!(
970 pool.allocate(0x1004, Alloc::Anon(0), String::from("bar0")),
971 Ok(0x1000)
972 );
973 }
974
975 #[test]
976 fn coalescing_u64_limits() {
977 let mut pool = AddressAllocator::new(
978 AddressRange {
979 start: 0,
980 end: u64::MAX - 1,
981 },
982 None,
983 None,
984 )
985 .unwrap();
986 assert!(pool
987 .insert_at(AddressRange {
988 start: u64::MAX,
989 end: u64::MAX,
990 })
991 .is_ok());
992 assert!(pool
993 .insert_at(AddressRange {
994 start: u64::MAX,
995 end: u64::MAX,
996 })
997 .is_err());
998 assert_eq!(
999 pool.allocate(u64::MAX, Alloc::Anon(0), String::from("bar0")),
1000 Ok(0)
1001 );
1002
1003 let mut pool = AddressAllocator::new(
1004 AddressRange {
1005 start: 1,
1006 end: u64::MAX,
1007 },
1008 None,
1009 None,
1010 )
1011 .unwrap();
1012 assert!(pool.insert_at(AddressRange { start: 0, end: 0 }).is_ok());
1013 assert!(pool.insert_at(AddressRange { start: 0, end: 0 }).is_err());
1014 assert_eq!(
1015 pool.allocate(u64::MAX, Alloc::Anon(0), String::from("bar0")),
1016 Ok(0)
1017 );
1018 }
1019
1020 #[test]
1021 fn allocate_and_verify_pci_offset() {
1022 let mut pool = AddressAllocator::new(
1023 AddressRange {
1024 start: 0x1000,
1025 end: 0xFFFF,
1026 },
1027 None,
1028 None,
1029 )
1030 .unwrap();
1031 let pci_bar0 = Alloc::PciBar {
1032 bus: 1,
1033 dev: 2,
1034 func: 0,
1035 bar: 0,
1036 };
1037 let pci_bar1 = Alloc::PciBar {
1038 bus: 1,
1039 dev: 2,
1040 func: 0,
1041 bar: 1,
1042 };
1043 let pci_bar2 = Alloc::PciBar {
1044 bus: 1,
1045 dev: 2,
1046 func: 0,
1047 bar: 2,
1048 };
1049 let anon = Alloc::Anon(1);
1050
1051 assert_eq!(
1052 pool.allocate(0x800, pci_bar0, String::from("bar0")),
1053 Ok(0x1000)
1054 );
1055 assert_eq!(
1056 pool.allocate(0x800, pci_bar1, String::from("bar1")),
1057 Ok(0x1800)
1058 );
1059 assert_eq!(pool.allocate(0x800, anon, String::from("anon")), Ok(0x2000));
1060
1061 assert_eq!(
1062 pool.address_from_pci_offset(pci_bar0, 0x600, 0x100),
1063 Ok(0x1600)
1064 );
1065 assert_eq!(
1066 pool.address_from_pci_offset(pci_bar1, 0x600, 0x100),
1067 Ok(0x1E00)
1068 );
1069 assert_eq!(
1070 pool.address_from_pci_offset(pci_bar0, 0x7FE, 0x001),
1071 Ok(0x17FE)
1072 );
1073 assert_eq!(
1074 pool.address_from_pci_offset(pci_bar0, 0x7FF, 0x001),
1075 Ok(0x17FF)
1076 );
1077 assert_eq!(
1078 pool.address_from_pci_offset(pci_bar0, 0x800, 0x001),
1079 Err(Error::OutOfBounds)
1080 );
1081
1082 assert_eq!(
1083 pool.address_from_pci_offset(pci_bar2, 0x7FF, 0x001),
1084 Err(Error::InvalidAlloc(pci_bar2))
1085 );
1086
1087 assert_eq!(
1088 pool.address_from_pci_offset(anon, 0x600, 0x100),
1089 Err(Error::InvalidAlloc(anon))
1090 );
1091 }
1092
1093 #[test]
1094 fn get_max_address_of_ranges() {
1095 let ranges = vec![
1096 AddressRange {
1097 start: 0x1000,
1098 end: 0xFFFF,
1099 },
1100 AddressRange {
1101 start: 0x20000,
1102 end: 0xFFFFF,
1103 },
1104 ];
1105 let pool = AddressAllocator::new_from_list(ranges, None, None).unwrap();
1106
1107 assert_eq!(pool.get_max_addr(), 0xFFFFF);
1108 }
1109
1110 #[test]
1111 fn allocate_u64_max_overflow() {
1112 let mut pool = AddressAllocator::new(
1113 AddressRange {
1114 start: u64::MAX - 3,
1115 end: u64::MAX,
1116 },
1117 None,
1118 None,
1119 )
1120 .unwrap();
1121 assert_eq!(
1122 pool.allocate(4, Alloc::Anon(0), String::from("bar0")),
1123 Ok(u64::MAX - 3)
1124 );
1125 assert_eq!(
1126 pool.get(&Alloc::Anon(0)),
1127 Some(&(
1128 AddressRange {
1129 start: u64::MAX - 3,
1130 end: u64::MAX,
1131 },
1132 String::from("bar0"),
1133 ))
1134 );
1135
1136 let mut pool = AddressAllocator::new(
1137 AddressRange {
1138 start: 1,
1139 end: u64::MAX,
1140 },
1141 None,
1142 None,
1143 )
1144 .unwrap();
1145 assert_eq!(
1146 pool.allocate(u64::MAX, Alloc::Anon(0), String::from("bar0")),
1147 Err(Error::OutOfSpace)
1148 );
1149 }
1150}