resources/
address_allocator.rs

1// Copyright 2018 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use 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/// Manages allocating address ranges.
16/// Use `AddressAllocator` whenever an address range needs to be allocated to different users.
17/// Allocations must be uniquely tagged with an Alloc enum, which can be used for lookup.
18/// An human-readable tag String must also be provided for debugging / reference.
19#[derive(Debug, Eq, PartialEq)]
20pub struct AddressAllocator {
21    /// The list of pools from which address are allocated. The union
22    /// of all regions from |allocs| and |regions| equals the pools.
23    pools: Vec<AddressRange>,
24    min_align: u64,
25    preferred_align: u64,
26    /// The region that is allocated.
27    allocs: HashMap<Alloc, (AddressRange, String)>,
28    /// The region that is not allocated yet.
29    regions: BTreeSet<AddressRange>,
30}
31
32impl AddressAllocator {
33    /// Creates a new `AddressAllocator` for managing a range of addresses.
34    /// Can return an error if `pool` is empty or if alignment isn't a power of two.
35    ///
36    /// * `pool` - The address range to allocate from.
37    /// * `min_align` - The minimum size of an address region to align to, defaults to four.
38    /// * `preferred_align` - The preferred alignment of an address region, used if possible.
39    ///
40    /// If an allocation cannot be satisfied with the preferred alignment, the minimum alignment
41    /// will be used instead.
42    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    /// Creates a new `AddressAllocator` for managing a range of addresses.
51    /// Can return `None` if all pools are empty alignment isn't a power of two.
52    ///
53    /// * `pools` - The list of pools to initialize the allocator with.
54    /// * `min_align` - The minimum size of an address region to align to, defaults to four.
55    /// * `preferred_align` - The preferred alignment of an address region, used if possible.
56    ///
57    /// If an allocation cannot be satisfied with the preferred alignment, the minimum alignment
58    /// will be used instead.
59    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    /// Gets the regions managed by the allocator.
93    ///
94    /// This returns the original `pools` value provided to `AddressAllocator::new()`.
95    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            // finds first region matching alignment and size.
144            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            // finds last region matching alignment and size.
159            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    /// Allocates a range of addresses from the reverse managed region with an optional tag
191    /// and minimal alignment. Returns allocated_address. (allocated_address, size, tag)
192    /// can be retrieved through the `get` method.
193    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    /// Allocates a range of addresses, preferring to allocate from high rather than low addresses.
204    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    /// Allocates a range of addresses from the managed region with an optional tag
214    /// and minimal alignment. Returns allocated_address. (allocated_address, size, tag)
215    /// can be retrieved through the `get` method.
216    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    /// Allocates a range of addresses from the managed region with an optional tag
236    /// and required location. Allocation alignment is not enforced.
237    /// Returns OutOfSpace if requested range is not available or ExistingAlloc if the requested
238    /// range overlaps an existing allocation.
239    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    /// Releases exising allocation back to free pool and returns the range that was released.
268    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    /// Release a allocation contains the value.
278    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    // Find an existing allocation that overlaps the region defined by `range`. If more
290    // than one allocation overlaps the given region, any of them may be returned, since the HashMap
291    // iterator is not ordered in any particular way.
292    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    // Return the max address of the allocated address ranges.
304    pub fn get_max_addr(&self) -> u64 {
305        self.regions.iter().fold(0, |x, range| x.max(range.end))
306    }
307
308    /// Returns allocation associated with `alloc`, or None if no such allocation exists.
309    pub fn get(&self, alloc: &Alloc) -> Option<&(AddressRange, String)> {
310        self.allocs.get(alloc)
311    }
312
313    /// Insert range of addresses into the pool, coalescing neighboring regions.
314    fn insert_at(&mut self, mut slot: AddressRange) -> Result<()> {
315        if slot.is_empty() {
316            return Err(Error::AllocSizeZero);
317        }
318
319        // Find the region with the highest starting address that is at most
320        // |slot.start|. Check if it overlaps with |slot|, or if it is adjacent to
321        // (and thus can be coalesced with) |slot|.
322        let mut smaller_merge = None;
323        if let Some(smaller) = self
324            .regions
325            .range((Bound::Unbounded, Bound::Included(slot)))
326            .max()
327        {
328            // If there is overflow, then |smaller| covers up through u64::MAX
329            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        // Find the region with the smallest starting address that is greater than
341        // |slot.start|. Check if it overlaps with |slot|, or if it is adjacent to
342        // (and thus can be coalesced with) |slot|.
343        let mut larger_merge = None;
344        if let Some(larger) = self
345            .regions
346            .range((Bound::Excluded(slot), Bound::Unbounded))
347            .min()
348        {
349            // If there is underflow, then |larger| covers down through 0
350            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    /// Returns an address from associated PCI `alloc` given an allocation offset and size.
375    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
400/// Contains a set of `AddressAllocator`s for allocating address ranges.
401/// When attempting an allocation, each allocator will be tried in order until
402/// the allocation is successful.
403/// See `AddressAllocator` for function documentation.
404pub 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        // Anon is used for brevity. Don't manually instantiate Anon allocs!
492        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        // 0x1200..0x1a00
651        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        // 0x600..0x2000
667        assert_eq!(
668            pool.allocate(0x600, Alloc::Anon(2), String::from("bar2")),
669            Ok(0x1a00)
670        );
671        // 0x1000..0x1200
672        assert_eq!(
673            pool.allocate(0x200, Alloc::Anon(3), String::from("bar3")),
674            Ok(0x1000)
675        );
676        // 0x1b00..0x1c00 (overlaps with 0x600..0x2000)
677        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        // 0x1fff..0x2000 (overlaps with 0x600..0x2000)
689        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        // 0x1200..0x1201 (overlaps with 0x1200..0x1a00)
701        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        // 0x11ff..0x1200 (overlaps with 0x1000..0x1200)
713        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        // 0x1100..0x1300 (overlaps with 0x1000..0x1200 and 0x1200..0x1a00)
725        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}