memory_range/
lib.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! The [`MemoryRange`] type, which represents a 4KB-page-aligned byte range of
5//! memory, plus algorithms operating on the type.
6
7#![forbid(unsafe_code)]
8#![no_std]
9
10use core::iter::Iterator;
11use core::iter::Peekable;
12use core::ops::Range;
13
14const PAGE_SIZE: u64 = 4096;
15const TWO_MB: u64 = 0x20_0000;
16const ONE_GB: u64 = 0x4000_0000;
17
18/// Represents a page-aligned byte range of memory.
19///
20/// This type has a stable `Protobuf` representation, and can be directly used
21/// in saved state.
22// TODO: enforce invariants during de/serialization
23#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
24#[cfg_attr(
25    feature = "mesh",
26    derive(mesh_protobuf::Protobuf),
27    mesh(package = "topology")
28)]
29#[cfg_attr(feature = "inspect", derive(inspect::Inspect), inspect(display))]
30pub struct MemoryRange {
31    #[cfg_attr(feature = "mesh", mesh(1))]
32    start: u64,
33    #[cfg_attr(feature = "mesh", mesh(2))]
34    end: u64,
35}
36
37impl core::fmt::Display for MemoryRange {
38    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
39        write!(f, "{:#x}-{:#x}", self.start(), self.end())
40    }
41}
42
43impl TryFrom<Range<u64>> for MemoryRange {
44    type Error = InvalidMemoryRange;
45
46    fn try_from(range: Range<u64>) -> Result<Self, Self::Error> {
47        Self::try_new(range)
48    }
49}
50
51impl TryFrom<Range<usize>> for MemoryRange {
52    type Error = InvalidMemoryRange;
53
54    fn try_from(range: Range<usize>) -> Result<Self, Self::Error> {
55        Self::try_new(range.start as u64..range.end as u64)
56    }
57}
58
59/// Error returned by [`MemoryRange::try_new`].
60#[derive(Debug, thiserror::Error)]
61#[error("unaligned or invalid memory range: {start:#x}-{end:#x}")]
62pub struct InvalidMemoryRange {
63    start: u64,
64    end: u64,
65}
66
67impl MemoryRange {
68    /// The maximum address that can be represented by a `MemoryRange`.
69    pub const MAX_ADDRESS: u64 = !(PAGE_SIZE - 1);
70
71    /// Returns a new range for the given guest address range.
72    ///
73    /// Panics if the start or end are not 4KB aligned or if the start is after
74    /// the end.
75    #[track_caller]
76    pub const fn new(range: Range<u64>) -> Self {
77        assert!(range.start & (PAGE_SIZE - 1) == 0);
78        assert!(range.end & (PAGE_SIZE - 1) == 0);
79        assert!(range.start <= range.end);
80        Self {
81            start: range.start,
82            end: range.end,
83        }
84    }
85
86    /// Returns a new range for the given guest address range.
87    ///
88    /// Returns `None` if the start or end are not 4KB aligned or if the start
89    /// is after the end.
90    pub const fn try_new(range: Range<u64>) -> Result<Self, InvalidMemoryRange> {
91        if range.start & (PAGE_SIZE - 1) != 0
92            || range.end & (PAGE_SIZE - 1) != 0
93            || range.start > range.end
94        {
95            return Err(InvalidMemoryRange {
96                start: range.start,
97                end: range.end,
98            });
99        }
100        Ok(Self {
101            start: range.start,
102            end: range.end,
103        })
104    }
105
106    /// Returns the smallest 4K-aligned range that contains the given address
107    /// range.
108    ///
109    /// Panics if the start is after the end or if the end address is in the
110    /// last page of the 64-bit space.
111    pub fn bounding(range: Range<u64>) -> Self {
112        Self::bounding_aligned(range, PAGE_SIZE)
113    }
114
115    /// Returns the smallest range with the specified alignment that contains
116    /// the given address range.
117    ///
118    /// Panics if the start is after the end or if the aligned end address
119    /// would overflow.
120    pub fn bounding_aligned(range: Range<u64>, alignment: u64) -> Self {
121        assert!(range.start <= range.end);
122        assert!(alignment.is_power_of_two());
123        assert!(range.end <= u64::MAX - (alignment - 1));
124        let start = range.start & !(alignment - 1);
125        let end = (range.end + (alignment - 1)) & !(alignment - 1);
126        Self::new(start..end)
127    }
128
129    /// Returns a new range for the given guest 4KB page range.
130    ///
131    /// Panics if the start is after the end or if the start address or end
132    /// address overflow.
133    pub fn from_4k_gpn_range(range: Range<u64>) -> Self {
134        const MAX: u64 = u64::MAX / PAGE_SIZE;
135        assert!(range.start <= MAX);
136        assert!(range.end <= MAX);
137        Self::new(range.start * PAGE_SIZE..range.end * PAGE_SIZE)
138    }
139
140    /// The empty range, with start and end addresses of zero.
141    pub const EMPTY: Self = Self::new(0..0);
142
143    /// The start address.
144    pub fn start(&self) -> u64 {
145        self.start
146    }
147
148    /// The start address as a 4KB page number.
149    pub fn start_4k_gpn(&self) -> u64 {
150        self.start / PAGE_SIZE
151    }
152
153    /// The end address as a 4KB page number.
154    pub fn end_4k_gpn(&self) -> u64 {
155        self.end / PAGE_SIZE
156    }
157
158    /// The number of 4KB pages in the range.
159    pub fn page_count_4k(&self) -> u64 {
160        (self.end - self.start) / PAGE_SIZE
161    }
162
163    /// The number of 2MB pages in the range.
164    pub fn page_count_2m(&self) -> u64 {
165        (self.end - self.start).div_ceil(TWO_MB)
166    }
167
168    /// The end address.
169    pub fn end(&self) -> u64 {
170        self.end
171    }
172
173    /// The length of the range in bytes.
174    pub fn len(&self) -> u64 {
175        self.end() - self.start()
176    }
177
178    /// Check if the range is empty.
179    pub fn is_empty(&self) -> bool {
180        self.start == self.end
181    }
182
183    /// Gets the biggest page size possible for the range.
184    pub fn alignment(&self, base: u64) -> u64 {
185        let order = ((base + self.start()) | (base + self.end())).trailing_zeros();
186        1 << order
187    }
188
189    /// Returns the largest range contained in this range whose start and end
190    /// are aligned to `alignment` bytes. This may be the empty range.
191    ///
192    /// Panics if `alignment` is not a power of two.
193    pub fn aligned_subrange(&self, alignment: u64) -> Self {
194        assert!(alignment.is_power_of_two());
195        let start = (self.start + alignment - 1) & !(alignment - 1);
196        let end = self.end & !(alignment - 1);
197        if start <= end {
198            Self::new(start..end)
199        } else {
200            Self::EMPTY
201        }
202    }
203
204    /// Returns whether `self` and `other` overlap.
205    pub fn overlaps(&self, other: &Self) -> bool {
206        self.end > other.start && self.start < other.end
207    }
208
209    /// Returns whether `self` contains `other`.
210    pub fn contains(&self, other: &Self) -> bool {
211        self.start <= other.start && self.end >= other.end
212    }
213
214    /// Returns whether `self` contains the byte at `addr`.
215    pub fn contains_addr(&self, addr: u64) -> bool {
216        (self.start..self.end).contains(&addr)
217    }
218
219    /// Returns the byte offset of `addr` within the range, if it is contained.
220    pub fn offset_of(&self, addr: u64) -> Option<u64> {
221        if self.contains_addr(addr) {
222            Some(addr - self.start)
223        } else {
224            None
225        }
226    }
227
228    /// Returns the intersection of `self` and `other`.
229    pub fn intersection(&self, other: &Self) -> Self {
230        let start = self.start.max(other.start);
231        let end = self.end.min(other.end);
232        if start <= end {
233            Self::new(start..end)
234        } else {
235            Self::EMPTY
236        }
237    }
238
239    /// Split the range at the given byte offset within the range.
240    ///
241    /// Panics if `offset` is not within the range or is not page-aligned.
242    #[track_caller]
243    pub fn split_at_offset(&self, offset: u64) -> (Self, Self) {
244        assert!(offset <= self.len());
245        assert!(offset.is_multiple_of(PAGE_SIZE));
246        (
247            Self {
248                start: self.start,
249                end: self.start + offset,
250            },
251            Self {
252                start: self.start + offset,
253                end: self.end,
254            },
255        )
256    }
257}
258
259impl From<MemoryRange> for Range<u64> {
260    fn from(range: MemoryRange) -> Self {
261        Range {
262            start: range.start(),
263            end: range.end(),
264        }
265    }
266}
267
268/// Iterator over aligned subranges of a memory range.
269///
270/// Each subrange will be made up of aligned pages of size 4KB, 2MB, or 1GB.
271/// Subrange boundaries are chosen to output the maximum number of the largest
272/// pages.
273#[derive(Debug, Clone)]
274pub struct AlignedSubranges {
275    range: MemoryRange,
276    offset: u64,
277    max_len: u64,
278}
279
280impl AlignedSubranges {
281    /// Returns a new iterator of subranges in `range`.
282    pub fn new(range: MemoryRange) -> Self {
283        Self {
284            range,
285            offset: 0,
286            max_len: u64::MAX,
287        }
288    }
289
290    /// Returns an iterator that considers subrange alignment offset by `offset`
291    /// bytes.
292    pub fn with_offset(self, offset: u64) -> Self {
293        Self { offset, ..self }
294    }
295
296    /// Returns an iterator that outputs subranges only up to `max_len` bytes.
297    pub fn with_max_range_len(self, max_len: u64) -> Self {
298        Self { max_len, ..self }
299    }
300}
301
302impl Iterator for AlignedSubranges {
303    type Item = MemoryRange;
304
305    fn next(&mut self) -> Option<Self::Item> {
306        if self.range.is_empty() {
307            return None;
308        }
309
310        let start = self.range.start() + self.offset;
311        let mut end = self.range.end() + self.offset;
312        if end - start > self.max_len {
313            end = start + self.max_len;
314        }
315
316        let mut align = |page_size| {
317            let page_mask: u64 = page_size - 1;
318            if (start + page_mask) & !page_mask >= end & !page_mask {
319                // No sense in aligning this, since we won't get any aligned
320                // pages.
321                return;
322            }
323            if start & page_mask != 0 {
324                // Align the next range's start.
325                end = end.min((start + page_mask) & !page_mask);
326            } else {
327                // Align this range's end.
328                end &= !page_mask;
329            }
330        };
331
332        align(TWO_MB);
333        align(ONE_GB);
334        let start = start - self.offset;
335        let end = end - self.offset;
336        self.range = MemoryRange::new(end..self.range.end());
337        Some(MemoryRange::new(start..end))
338    }
339}
340
341/// Returns an iterator over memory ranges that are in both `left` and `right`.
342///
343/// For example, if `left` is `[0..4MB, 8MB..12MB]` and `right` is `[2MB..6MB, 10MB..11MB]`,
344/// the resulting iterator will yield `[2MB..4MB, 10MB..11MB]`.
345///
346/// Panics if `left` or `right` are not sorted or are overlapping.
347pub fn overlapping_ranges(
348    left: impl IntoIterator<Item = MemoryRange>,
349    right: impl IntoIterator<Item = MemoryRange>,
350) -> impl Iterator<Item = MemoryRange> {
351    walk_ranges(
352        left.into_iter().map(|r| (r, ())),
353        right.into_iter().map(|r| (r, ())),
354    )
355    .filter_map(|(r, c)| match c {
356        RangeWalkResult::Both((), ()) => Some(r),
357        _ => None,
358    })
359}
360
361/// Returns an iterator over the ranges in `left` that are not in `right`.
362///
363/// For example, if `left` is `[0..4MB, 8MB..12MB]` and `right` is `[2MB..6MB,
364/// 10MB..11MB]`, the resulting iterator will yield `[0..2MB, 8MB..10MB,
365/// 11MB..12MB]`.
366///
367/// Panics if `left` or `right` are not sorted or are overlapping.
368pub fn subtract_ranges(
369    left: impl IntoIterator<Item = MemoryRange>,
370    right: impl IntoIterator<Item = MemoryRange>,
371) -> impl Iterator<Item = MemoryRange> {
372    walk_ranges(
373        left.into_iter().map(|r| (r, ())),
374        right.into_iter().map(|r| (r, ())),
375    )
376    .filter_map(|(r, c)| match c {
377        RangeWalkResult::Left(()) => Some(r),
378        RangeWalkResult::Neither | RangeWalkResult::Right(()) | RangeWalkResult::Both((), ()) => {
379            None
380        }
381    })
382}
383
384/// Returns an iterator that computes the overlapping state of the ranges in
385/// `left` and `right`.
386///
387/// The iterator yields a tuple of a [`MemoryRange`] and a [`RangeWalkResult`]
388/// enum that indicates whether each subrange is only in `left`, only in
389/// `right`, in both, or in neither.
390///
391/// Panics if `left` or `right` are not sorted.
392///
393/// # Examples
394///
395/// ```
396/// # use memory_range::{MemoryRange, RangeWalkResult, walk_ranges};
397/// let left = [(MemoryRange::new(0x100000..0x400000), "first"), (MemoryRange::new(0x800000..0xc00000), "second")];
398/// let right = [(MemoryRange::new(0x200000..0x900000), 1000), (MemoryRange::new(0x900000..0xa00000), 2000)];
399/// let v: Vec<_> = walk_ranges(left, right).collect();
400/// let expected = [
401///     (MemoryRange::new(0..0x100000), RangeWalkResult::Neither),
402///     (MemoryRange::new(0x100000..0x200000), RangeWalkResult::Left("first")),
403///     (MemoryRange::new(0x200000..0x400000), RangeWalkResult::Both("first", 1000)),
404///     (MemoryRange::new(0x400000..0x800000), RangeWalkResult::Right(1000)),
405///     (MemoryRange::new(0x800000..0x900000), RangeWalkResult::Both("second", 1000)),
406///     (MemoryRange::new(0x900000..0xa00000), RangeWalkResult::Both("second", 2000)),
407///     (MemoryRange::new(0xa00000..0xc00000), RangeWalkResult::Left("second")),
408///     (MemoryRange::new(0xc00000..MemoryRange::MAX_ADDRESS), RangeWalkResult::Neither),
409/// ];
410/// assert_eq!(v.as_slice(), expected.as_slice());
411/// ```
412pub fn walk_ranges<T: Clone, U: Clone>(
413    left: impl IntoIterator<Item = (MemoryRange, T)>,
414    right: impl IntoIterator<Item = (MemoryRange, U)>,
415) -> impl Iterator<Item = (MemoryRange, RangeWalkResult<T, U>)> {
416    RangeWalkIter {
417        pos: 0,
418        left: PeekableSorted::new(left),
419        right: PeekableSorted::new(right),
420    }
421}
422
423/// The result of an iteration of [`walk_ranges`].
424#[derive(Copy, Clone, Debug, PartialEq, Eq)]
425pub enum RangeWalkResult<T, U> {
426    /// Neither iterator contains this range.
427    Neither,
428    /// Only the left iterator contains this range, in the element with the
429    /// given value.
430    Left(T),
431    /// Only the right iterator contains this range, in the element with the
432    /// given value.
433    Right(U),
434    /// Both iterators contain this range, in the elements with the given
435    /// values.
436    Both(T, U),
437}
438
439struct RangeWalkIter<I: Iterator, J: Iterator> {
440    pos: u64,
441    left: PeekableSorted<I>,
442    right: PeekableSorted<J>,
443}
444
445struct PeekableSorted<I: Iterator> {
446    iter: I,
447    #[expect(clippy::option_option)] // `Some(None)` is used to remember that `iter` is empty.
448    item: Option<Option<I::Item>>,
449}
450
451impl<I: Iterator<Item = (MemoryRange, T)>, T> PeekableSorted<I> {
452    fn new(iter: impl IntoIterator<IntoIter = I>) -> Self {
453        Self {
454            iter: iter.into_iter(),
455            item: None,
456        }
457    }
458
459    fn peek_in_range_ensure_sorted(&mut self, pos: u64, msg: &str) -> Option<&(MemoryRange, T)> {
460        loop {
461            let r = self
462                .item
463                .get_or_insert_with(|| {
464                    let r = self.iter.next()?;
465                    assert!(r.0.start() >= pos, "{msg} not sorted");
466                    Some(r)
467                })
468                .as_ref()?;
469            if !r.0.is_empty() && r.0.end() > pos {
470                return Some(self.item.as_ref().unwrap().as_ref().unwrap());
471            }
472            self.item = None;
473        }
474    }
475}
476
477impl<I: Iterator<Item = (MemoryRange, T)>, J: Iterator<Item = (MemoryRange, U)>, T: Clone, U: Clone>
478    Iterator for RangeWalkIter<I, J>
479{
480    type Item = (MemoryRange, RangeWalkResult<T, U>);
481
482    fn next(&mut self) -> Option<Self::Item> {
483        if self.pos == MemoryRange::MAX_ADDRESS {
484            return None;
485        }
486        let left = self.left.peek_in_range_ensure_sorted(self.pos, "left");
487        let right = self.right.peek_in_range_ensure_sorted(self.pos, "right");
488        let (end, c) = match (left, right) {
489            (Some(&(left, ref t)), Some(&(right, ref u))) => {
490                if self.pos < left.start() {
491                    if self.pos < right.start() {
492                        (left.start().min(right.start()), RangeWalkResult::Neither)
493                    } else {
494                        (
495                            left.start().min(right.end()),
496                            RangeWalkResult::Right(u.clone()),
497                        )
498                    }
499                } else if self.pos < right.start() {
500                    (
501                        right.start().min(left.end()),
502                        RangeWalkResult::Left(t.clone()),
503                    )
504                } else {
505                    (
506                        left.end().min(right.end()),
507                        RangeWalkResult::Both(t.clone(), u.clone()),
508                    )
509                }
510            }
511            (Some(&(left, ref t)), None) => {
512                if self.pos < left.start() {
513                    (left.start, RangeWalkResult::Neither)
514                } else {
515                    (left.end(), RangeWalkResult::Left(t.clone()))
516                }
517            }
518            (None, Some(&(right, ref u))) => {
519                if self.pos < right.start() {
520                    (right.start, RangeWalkResult::Neither)
521                } else {
522                    (right.end(), RangeWalkResult::Right(u.clone()))
523                }
524            }
525            (None, None) => (MemoryRange::MAX_ADDRESS, RangeWalkResult::Neither),
526        };
527        let r = MemoryRange::new(self.pos..end);
528        self.pos = end;
529        Some((r, c))
530    }
531}
532
533/// Takes a sequence of memory ranges, sorted by their start address, and
534/// returns an iterator over the flattened ranges, where overlapping and
535/// adjacent ranges are merged and deduplicated.
536///
537/// Panics if the input ranges are not sorted by their start address.
538///
539/// # Example
540/// ```rust
541/// # use memory_range::{flatten_ranges, MemoryRange};
542/// let ranges = [
543///     MemoryRange::new(0x1000..0x2000),
544///     MemoryRange::new(0x2000..0x5000),
545///     MemoryRange::new(0x4000..0x6000),
546///     MemoryRange::new(0x5000..0x6000),
547///     MemoryRange::new(0x8000..0x9000),
548/// ];
549/// let flattened = [
550///     MemoryRange::new(0x1000..0x6000),
551///     MemoryRange::new(0x8000..0x9000),
552/// ];
553/// assert!(flatten_ranges(ranges).eq(flattened));
554/// ```
555pub fn flatten_ranges(
556    ranges: impl IntoIterator<Item = MemoryRange>,
557) -> impl Iterator<Item = MemoryRange> {
558    FlattenIter {
559        iter: ranges.into_iter().peekable(),
560    }
561}
562
563struct FlattenIter<I: Iterator> {
564    iter: Peekable<I>,
565}
566
567impl<I: Iterator<Item = MemoryRange>> Iterator for FlattenIter<I> {
568    type Item = MemoryRange;
569
570    fn next(&mut self) -> Option<Self::Item> {
571        let first = self.iter.next()?;
572        let mut start = first.start();
573        let mut end = first.end();
574        while let Some(r) = self.iter.next_if(|r| {
575            assert!(r.start() >= start, "ranges are not sorted");
576            r.start() <= end
577        }) {
578            start = r.start();
579            end = end.max(r.end());
580        }
581        Some(MemoryRange::new(first.start()..end))
582    }
583}
584
585/// Similar to [`flatten_ranges`], but considers ranges non-equivalent if their
586/// associated tags differ.
587///
588/// Panics if the input ranges are not sorted by their start address, or if
589/// ranges overlap.
590///
591/// # Example
592/// ```rust
593/// # use memory_range::{merge_adjacent_ranges, MemoryRange};
594///
595/// #[derive(Clone, Copy, Debug, PartialEq, Eq)]
596/// enum Color {
597///    Red,
598///    Blue,
599/// }
600///
601/// let ranges = [
602///     (MemoryRange::new(0x1000..0x2000), Color::Red),
603///     (MemoryRange::new(0x2000..0x5000), Color::Red),
604///     (MemoryRange::new(0x5000..0x6000), Color::Blue),
605///     (MemoryRange::new(0x8000..0x9000), Color::Red),
606/// ];
607/// let flattened = [
608///     (MemoryRange::new(0x1000..0x5000), Color::Red),
609///     (MemoryRange::new(0x5000..0x6000), Color::Blue),
610///     (MemoryRange::new(0x8000..0x9000), Color::Red),
611/// ];
612/// assert!(merge_adjacent_ranges(ranges).eq(flattened));
613/// ```
614pub fn merge_adjacent_ranges<T: PartialEq>(
615    ranges: impl IntoIterator<Item = (MemoryRange, T)>,
616) -> impl Iterator<Item = (MemoryRange, T)> {
617    MergeAdjacentIter {
618        iter: ranges.into_iter().peekable(),
619    }
620}
621
622struct MergeAdjacentIter<I: Iterator> {
623    iter: Peekable<I>,
624}
625
626impl<I: Iterator<Item = (MemoryRange, T)>, T: PartialEq> Iterator for MergeAdjacentIter<I> {
627    type Item = (MemoryRange, T);
628
629    fn next(&mut self) -> Option<Self::Item> {
630        let (first, typ) = self.iter.next()?;
631        let mut start = first.start();
632        let mut end = first.end();
633        while let Some((r, _t)) = self.iter.next_if(|(r, t)| {
634            assert!(r.start() >= start, "ranges are not sorted");
635            assert!(r.start() >= end, "ranges overlap");
636            r.start() == end && &typ == t
637        }) {
638            start = r.start();
639            end = end.max(r.end());
640        }
641        Some((MemoryRange::new(first.start()..end), typ))
642    }
643}
644
645#[cfg(test)]
646mod tests {
647    extern crate alloc;
648    use super::MemoryRange;
649    use super::TWO_MB;
650    use crate::AlignedSubranges;
651    use crate::flatten_ranges;
652    use crate::merge_adjacent_ranges;
653    use crate::overlapping_ranges;
654    use crate::subtract_ranges;
655    use alloc::vec;
656    use alloc::vec::Vec;
657
658    const KB: u64 = 1024;
659    const MB: u64 = 1024 * KB;
660    const GB: u64 = 1024 * MB;
661
662    #[test]
663    fn test_align() {
664        #[derive(Clone, Debug, PartialEq, Copy)]
665        struct AlignedRangeResult {
666            range: MemoryRange,
667            page_size: u64,
668        }
669
670        let compare_with_base = |r1: Vec<MemoryRange>, base: u64, er: Vec<AlignedRangeResult>| {
671            let result: Vec<_> = r1
672                .iter()
673                .flat_map(|range| AlignedSubranges::new(*range).with_offset(base))
674                .collect();
675            assert_eq!(result.len(), er.len());
676            for (pos, range) in result.iter().enumerate() {
677                assert_eq!(*range, er[pos].range);
678                assert_eq!(range.alignment(base), er[pos].page_size);
679            }
680        };
681
682        let compare = |r1: Vec<MemoryRange>, er: Vec<AlignedRangeResult>| {
683            compare_with_base(r1, 0, er);
684        };
685
686        /// Builds a memory range with short hand.
687        fn b_mr(start: u64, end: u64) -> MemoryRange {
688            MemoryRange::new(start..end)
689        }
690
691        /// Builds a aligned range result with short hand.
692        fn b_arr(range: MemoryRange, page_size: u64) -> AlignedRangeResult {
693            AlignedRangeResult { range, page_size }
694        }
695
696        // [0, 4KB]
697        let ram: Vec<_> = vec![b_mr(0x0, 4 * KB)];
698        let expected_res: Vec<_> = vec![b_arr(ram[0], 1 << 12)];
699
700        compare(ram, expected_res);
701
702        // [0, 2MB]
703        let ram: Vec<_> = vec![b_mr(0x0, 2 * MB)];
704        let expected_res: Vec<_> = vec![b_arr(ram[0], 1 << 21)];
705        compare(ram, expected_res);
706
707        // [0, 1MB]
708        let ram: Vec<_> = vec![b_mr(0x0, GB)];
709        let expected_res: Vec<_> = vec![b_arr(ram[0], 1 << 30)];
710        compare(ram, expected_res);
711
712        // [6MB, 12.004MB]
713        let ram: Vec<_> = vec![b_mr(6 * MB, 12 * MB + 4 * KB)];
714        let expected_res: Vec<_> = vec![
715            b_arr(b_mr(6 * MB, 12 * MB), TWO_MB),
716            b_arr(b_mr(12 * MB, 12 * MB + 4 * KB), 1 << 12),
717        ];
718        compare(ram, expected_res);
719
720        // [5.4MB, 12.2MB]
721        let ram: Vec<_> = vec![b_mr(5 * MB + 400 * KB, 12 * MB + 400 * KB)];
722        let expected_res: Vec<_> = vec![
723            b_arr(b_mr(5 * MB + 400 * KB, 6 * MB), 1 << 14),
724            b_arr(b_mr(6 * MB, 12 * MB), TWO_MB),
725            b_arr(b_mr(12 * MB, 12 * MB + 400 * KB), 1 << 14),
726        ];
727        compare(ram, expected_res);
728
729        // [1.501GB, 3.503GB]
730        let ram: Vec<_> = vec![b_mr(GB + 501 * MB, 3 * GB + 503 * MB)];
731        let expected_res: Vec<_> = vec![
732            b_arr(b_mr(GB + 501 * MB, GB + 502 * MB), 1 << 20),
733            b_arr(b_mr(GB + 502 * MB, 2 * GB), 1 << 21),
734            b_arr(b_mr(2 * GB, 3 * GB), 1 << 30),
735            b_arr(b_mr(3 * GB, 3 * GB + 502 * MB), 1 << 21),
736            b_arr(b_mr(3 * GB + 502 * MB, 3 * GB + 503 * MB), 1 << 20),
737        ];
738        compare(ram, expected_res);
739
740        // [4.008MB, 6.008MB] with necessary base to align up to 2MB
741        let ram: Vec<_> = vec![b_mr(4 * MB + 8 * KB, 6 * MB + 8 * KB)];
742        let base = 2 * MB - 8 * KB;
743        let expected_res: Vec<_> = vec![b_arr(ram[0], TWO_MB)];
744        compare_with_base(ram, base, expected_res);
745
746        // [4.008MB, 6.008MB] without any base
747        let ram: Vec<_> = vec![b_mr(4 * MB + 8 * KB, 6 * MB + 8 * KB)];
748        let expected_res: Vec<_> = vec![b_arr(ram[0], 1 << 13)];
749        compare_with_base(ram, 0, expected_res);
750    }
751
752    #[test]
753    fn test_overlapping_ranges() {
754        let left = [
755            MemoryRange::new(0..4 * MB),
756            MemoryRange::new(8 * MB..12 * MB),
757            MemoryRange::new(12 * MB..12 * MB),
758            MemoryRange::new(16 * MB..20 * MB),
759            MemoryRange::new(24 * MB..32 * MB),
760            MemoryRange::new(40 * MB..48 * MB),
761        ];
762        let right = [
763            MemoryRange::new(2 * MB..6 * MB),
764            MemoryRange::new(10 * MB..11 * MB),
765            MemoryRange::new(11 * MB..11 * MB),
766            MemoryRange::new(11 * MB..13 * MB),
767            MemoryRange::new(15 * MB..22 * MB),
768            MemoryRange::new(26 * MB..30 * MB),
769        ];
770
771        let result: Vec<_> = overlapping_ranges(left, right).collect();
772        assert_eq!(
773            result.as_slice(),
774            &[
775                MemoryRange::new(2 * MB..4 * MB),
776                MemoryRange::new(10 * MB..11 * MB),
777                MemoryRange::new(11 * MB..12 * MB),
778                MemoryRange::new(16 * MB..20 * MB),
779                MemoryRange::new(26 * MB..30 * MB),
780            ]
781        );
782    }
783
784    #[test]
785    fn test_subtract_ranges() {
786        let left = [
787            MemoryRange::new(0..4 * MB),
788            MemoryRange::new(8 * MB..12 * MB),
789            MemoryRange::new(12 * MB..12 * MB),
790            MemoryRange::new(16 * MB..20 * MB),
791            MemoryRange::new(24 * MB..32 * MB),
792            MemoryRange::new(40 * MB..48 * MB),
793        ];
794        let right = [
795            MemoryRange::new(2 * MB..6 * MB),
796            MemoryRange::new(10 * MB..11 * MB),
797            MemoryRange::new(11 * MB..11 * MB),
798            MemoryRange::new(11 * MB..13 * MB),
799            MemoryRange::new(15 * MB..22 * MB),
800            MemoryRange::new(26 * MB..30 * MB),
801        ];
802
803        let result: Vec<_> = subtract_ranges(left, right).collect();
804        assert_eq!(
805            result.as_slice(),
806            &[
807                MemoryRange::new(0..2 * MB),
808                MemoryRange::new(8 * MB..10 * MB),
809                MemoryRange::new(24 * MB..26 * MB),
810                MemoryRange::new(30 * MB..32 * MB),
811                MemoryRange::new(40 * MB..48 * MB),
812            ]
813        );
814    }
815
816    #[test]
817    #[should_panic(expected = "left not sorted")]
818    fn test_panic_unsorted_overlapping_left() {
819        overlapping_ranges(
820            [MemoryRange::new(MB..2 * MB), MemoryRange::new(0..MB)],
821            [MemoryRange::new(3 * MB..4 * MB)],
822        )
823        .for_each(|_| ());
824    }
825
826    #[test]
827    #[should_panic(expected = "right not sorted")]
828    fn test_panic_unsorted_overlapping_right() {
829        overlapping_ranges(
830            [
831                MemoryRange::new(MB..2 * MB),
832                MemoryRange::new(3 * MB..4 * MB),
833            ],
834            [MemoryRange::new(0..MB), MemoryRange::new(0..MB)],
835        )
836        .for_each(|_| ());
837    }
838
839    #[test]
840    #[should_panic(expected = "left not sorted")]
841    fn test_panic_unsorted_subtract_left() {
842        subtract_ranges(
843            [MemoryRange::new(MB..2 * MB), MemoryRange::new(0..MB)],
844            [MemoryRange::new(MB..2 * MB)],
845        )
846        .for_each(|_| ());
847    }
848
849    #[test]
850    #[should_panic(expected = "right not sorted")]
851    fn test_panic_unsorted_subtract_right() {
852        subtract_ranges(
853            [
854                MemoryRange::new(MB..2 * MB),
855                MemoryRange::new(3 * MB..4 * MB),
856            ],
857            [MemoryRange::new(MB..2 * MB), MemoryRange::new(MB..2 * MB)],
858        )
859        .for_each(|_| ());
860    }
861
862    #[test]
863    fn test_aligned_subrange() {
864        let test_cases = &[
865            (0..0, MB, 0..0),
866            (0..MB, MB, 0..MB),
867            (4 * KB..MB + 4 * KB, MB, MB..MB),
868            (MB..5 * MB, 2 * MB, 2 * MB..4 * MB),
869        ];
870        for (range, alignment, expected_aligned_range) in test_cases.iter().cloned() {
871            assert_eq!(
872                MemoryRange::new(range).aligned_subrange(alignment),
873                MemoryRange::new(expected_aligned_range)
874            );
875        }
876    }
877
878    #[test]
879    fn test_flatten_ranges() {
880        let ranges =
881            [0..4, 5..7, 6..11, 13..20, 20..25, 22..24, 35..36].map(MemoryRange::from_4k_gpn_range);
882        let result = [0..4, 5..11, 13..25, 35..36].map(MemoryRange::from_4k_gpn_range);
883        assert!(flatten_ranges(ranges).eq(result));
884    }
885
886    #[test]
887    #[should_panic(expected = "ranges are not sorted")]
888    fn test_flatten_ranges_not_sorted() {
889        flatten_ranges([0..4, 5..7, 3..8].map(MemoryRange::from_4k_gpn_range)).for_each(|_| ());
890    }
891
892    #[test]
893    fn test_merge_adjacent_ranges() {
894        #[derive(Clone, Copy, PartialEq, Eq)]
895        enum Color {
896            Red,
897            Blue,
898        }
899
900        let ranges = [0..4, 5..7, 7..11, 11..12, 13..20, 20..25, 35..36]
901            .map(MemoryRange::from_4k_gpn_range)
902            .into_iter()
903            .zip([
904                Color::Red,
905                Color::Red,
906                Color::Red,
907                Color::Blue,
908                Color::Red,
909                Color::Red,
910                Color::Blue,
911            ]);
912        let result = [0..4, 5..11, 11..12, 13..25, 35..36]
913            .map(MemoryRange::from_4k_gpn_range)
914            .into_iter()
915            .zip([Color::Red, Color::Red, Color::Blue, Color::Red, Color::Blue]);
916        assert!(merge_adjacent_ranges(ranges).eq(result));
917    }
918
919    #[test]
920    #[should_panic(expected = "ranges are not sorted")]
921    fn test_merge_adjacent_ranges_not_sorted() {
922        merge_adjacent_ranges([0..4, 5..7, 3..8].map(|r| (MemoryRange::from_4k_gpn_range(r), ())))
923            .for_each(|_| ());
924    }
925
926    #[test]
927    #[should_panic(expected = "ranges overlap")]
928    fn test_merge_adjacent_ranges_overlap() {
929        merge_adjacent_ranges([0..6, 5..7, 9..12].map(|r| (MemoryRange::from_4k_gpn_range(r), ())))
930            .for_each(|_| ());
931    }
932}
memory_range/lib.rs

memory_range/
lib.rs
