region.rs

use core::{
    fmt,
    ops::{Range, RangeInclusive},
};

/// A non-empty range of `usize` values
///
/// Like [`RangeInclusive`], `Region` can repesent any `usize` range, even those
/// that contain `usize::MAX`.
#[derive(Clone, Copy, PartialEq, Eq)]
pub struct Region {
    /// Start of the region, inclusive
    start: usize,
    /// End of the region, inclusive
    end: usize,
}

impl fmt::Debug for Region {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        (self.start..=self.end).fmt(f)
    }
}

fn mask_of_order(order: u32) -> usize {
    if order == usize::BITS {
        usize::MAX
    } else {
        let val = 1usize
            .checked_shl(order)
            .expect("Order should be in range 0..=usize::BITS");
        val - 1
    }
}

impl Region {
    /// Create a `Region` with range `start..end`
    ///
    /// Returns `None` if `start..end` is empty.
    ///
    /// [`Region::new`] cannot create regions that contain `usize::MAX`.
    pub fn new(range: Range<usize>) -> Option<Self> {
        Self::inclusive(range.start..=range.end.checked_sub(1)?)
    }

    /// Create a `Region` from starting value and length
    ///
    /// Returns `None` if `len` is `0` or if the maximum would be larger than
    /// `usize::MAX`.
    ///
    /// [`Region::new_len`] cannot create the region containing every `usize`
    pub fn new_len(start: usize, len: usize) -> Option<Self> {
        Self::inclusive(start..=len.checked_sub(1)?.checked_add(start)?)
    }

    /// Create a `Region` with range `start..=end`
    ///
    /// Returns `None` if `start..=end` is empty.
    ///
    /// [`Region::inclusive`] can create all possible regions.
    pub fn inclusive(range: RangeInclusive<usize>) -> Option<Self> {
        (range.start() <= range.end()).then_some(Self {
            start: *range.start(),
            end: *range.end(),
        })
    }

    fn inclusive_assert(start: usize, end_incl: usize) -> Self {
        debug_assert!(start <= end_incl);
        Self {
            start,
            end: end_incl,
        }
    }

    /// Create a `Region` ending at `usize::MAX` from a length
    ///
    /// Returns `None` if `len` is `0`.
    ///
    /// [`Region::new_len_end`] cannot create the region containing every
    /// `usize`
    pub fn new_len_end(len: usize) -> Option<Self> {
        len.checked_sub(1)
            .map(|len_1| Self::inclusive_assert(usize::MAX - len_1, usize::MAX))
    }

    /// The region containing every `usize`
    pub fn full() -> Self {
        Self {
            start: usize::MIN,
            end: usize::MAX,
        }
    }

    /// Start of the region
    pub fn start(&self) -> usize {
        self.start
    }

    /// End of the region (inclusive)
    pub fn end_inclusive(&self) -> usize {
        self.end
    }

    /// End of the region (exclusive)
    ///
    /// Returns `None` if the end would overflow, i.e. the region contains
    /// `usize::MAX`
    pub fn end(&self) -> Option<usize> {
        self.end.checked_add(1)
    }

    /// Length of the region
    ///
    /// Since `Regions` are all non-empty, the size is never zero.
    ///
    /// Returns `None` if the length would overflow, i.e. the region contains
    /// every `usize`
    #[allow(clippy::len_without_is_empty)] // Always non-empty
    pub fn len(&self) -> Option<usize> {
        (self.end - self.start).checked_add(1)
    }

    /// Largest `order`-power-of-two-aligned region contained within this region
    ///
    /// Returns `None` if no such (non-empty) region would fit
    pub fn align_fit(&self, order: u32) -> Option<Self> {
        let mask = mask_of_order(order);
        let start = self.start.checked_add(mask)? & !mask;
        let end = self.end.checked_sub(mask)? | mask;

        // Check if range is still non-empty
        Self::inclusive(start..=end)
    }

    /// Smallest `order`-power-of-two-aligned region that covers this region
    pub fn align_cover(&self, order: u32) -> Self {
        let mask = mask_of_order(order);
        let start = self.start & !mask;
        let end = self.end | mask;

        Self::inclusive_assert(start, end)
    }

    /// Allocate `len` from the start of region
    ///
    /// Returns `None` if the allocation does not fit, `Some((start, remain))`
    /// if the allocation fits, starting at `start`, and `remain` is the
    /// remaining unallocated part of the region or `None` if the entire region
    /// is allocated.
    pub fn allocate(&self, len: usize) -> Option<(usize, Option<Self>)> {
        self.len().map_or(true, |l| len <= l).then(|| {
            (
                self.start,
                (self.len() != Some(len))
                    .then(|| Self::inclusive_assert(self.start + len, self.end)),
            )
        })
    }

    fn split_at_right(&self, pos: usize) -> (Option<Self>, Option<Self>) {
        (
            pos.checked_sub(1)
                .and_then(|pos_1| Self::inclusive(self.start..=self.end.min(pos_1))),
            Self::inclusive(self.start.max(pos)..=self.end),
        )
    }

    fn split_at_left(&self, pos: usize) -> (Option<Self>, Option<Self>) {
        (
            Self::inclusive(self.start..=self.end.min(pos)),
            pos.checked_add(1)
                .and_then(|pos_1| Self::inclusive(self.start.max(pos_1)..=self.end)),
        )
    }

    /// Remove values in `other` from region, return remaining parts
    ///
    /// Carving out `other` would leave at most two parts:
    ///
    /// ```text
    /// other:         4 5 6
    /// self:  0 1 2 3 4 5 6 7 8 9
    /// left:  0 1 2 3
    /// right:               7 8 9
    /// ```
    ///
    /// Returns the two parts as `(left, right)`. `None` means that part is
    /// empty.
    ///
    /// If any parts of `other` are not contained in this region they are
    /// ignored.
    pub fn carve_out(&self, other: &Self) -> (Option<Self>, Option<Self>) {
        let (left, mid_right) = self.split_at_right(other.start);
        let right = mid_right.and_then(|mr| mr.split_at_left(other.end).1);
        (left, right)
    }

    /// Offset of `value` into this region
    ///
    /// If this region does not contain `value`, returns `None`.
    pub fn offset(&self, value: usize) -> Option<usize> {
        // https://github.com/rust-lang/rust-clippy/issues/9422
        #[allow(clippy::unnecessary_lazy_evaluations)]
        self.contains(value).then(|| value - self.start)
    }

    /// Returns `true` if this region contains `value`
    pub fn contains(&self, value: usize) -> bool {
        (self.start..=self.end).contains(&value)
    }

    /// Convert into [`RangeInclusive`]
    pub fn range(&self) -> RangeInclusive<usize> {
        self.start..=self.end
    }
}

#[cfg(test)]
mod test {
    use super::Region;

    #[test]
    fn test_new() {
        assert_eq!(Region::new(0..1).unwrap().range(), 0..=0);
        assert_eq!(
            Region::new(0..usize::MAX).unwrap().range(),
            0..=usize::MAX - 1
        );

        assert_eq!(Region::new(0..0), None); // Empty
        assert_eq!(Region::new(usize::MAX..0), None); // Empty
    }

    #[test]
    fn test_new_len() {
        assert_eq!(
            Region::new_len(0, usize::MAX).unwrap().range(),
            0..=usize::MAX - 1
        );

        // Should allow going all the way to the end of usize range
        assert_eq!(
            Region::new_len(1, usize::MAX).unwrap().range(),
            1..=usize::MAX
        );

        assert_eq!(Region::new_len(0, 0), None); // Empty
        assert_eq!(Region::new_len(2, usize::MAX), None); // Overflow
    }

    #[test]
    fn test_new_len_end() {
        assert_eq!(
            Region::new_len_end(1).unwrap().range(),
            usize::MAX..=usize::MAX
        );
        assert_eq!(
            Region::new_len_end(usize::MAX).unwrap().range(),
            1..=usize::MAX
        );

        assert_eq!(Region::new_len_end(0), None); // Empty
    }

    #[test]
    fn test_new_inclusive() {
        assert_eq!(Region::inclusive(0..=0).unwrap().range(), 0..=0);
        assert_eq!(
            Region::inclusive(0..=usize::MAX).unwrap().range(),
            0..=usize::MAX
        );

        assert_eq!(Region::inclusive(usize::MAX..=0), None); // Empty
    }

    #[test]
    fn test_end() {
        assert_eq!(Region::new(10..20).unwrap().end(), Some(20));
        assert_eq!(Region::new(0..usize::MAX).unwrap().end(), Some(usize::MAX));
        assert_eq!(Region::inclusive(0..=usize::MAX).unwrap().end(), None);
        assert_eq!(Region::new_len_end(42).unwrap().end(), None);
    }

    #[test]
    fn test_len() {
        assert_eq!(Region::new(10..20).unwrap().len(), Some(10));
        assert_eq!(Region::new_len(10, 20).unwrap().len(), Some(20));
        assert_eq!(Region::inclusive(10..=20).unwrap().len(), Some(11));

        assert_eq!(Region::full().len(), None); // Overflow

        assert_eq!(
            Region::new(0..usize::MAX - 1).unwrap().len(),
            Some(usize::MAX - 1)
        );

        assert_eq!(
            Region::inclusive(0..=usize::MAX - 1).unwrap().len(),
            Some(usize::MAX)
        );

        assert_eq!(
            Region::new_len(0, usize::MAX).unwrap().len(),
            Some(usize::MAX)
        );
    }

    #[test]
    fn test_align_fit() {
        let region = Region::new(1000..2500).unwrap();
        assert_eq!(region.align_fit(0), Some(region));
        assert_eq!(region.align_fit(10), Region::new(1024..2048));
        assert_eq!(region.align_fit(11), None);
        assert_eq!(region.align_fit(usize::BITS), None);
    }

    #[test]
    fn test_align_fit_end() {
        let region = Region::new_len_end(2000).unwrap();
        assert_eq!(region.align_fit(0), Some(region));
        assert_eq!(region.align_fit(10), Region::new_len_end(1024));
        assert_eq!(region.align_fit(11), None);
        assert_eq!(region.align_fit(usize::BITS), None);
    }

    #[test]
    fn test_align_fit_full() {
        assert_eq!(Region::full().align_fit(10), Some(Region::full()));
        assert_eq!(Region::full().align_fit(usize::BITS), Some(Region::full()));
    }

    #[test]
    fn test_align_cover() {
        let region = Region::new(1000..2600).unwrap();
        assert_eq!(region.align_cover(0), region);
        assert_eq!(region.align_cover(9), Region::new(512..3072).unwrap());
        assert_eq!(region.align_cover(10), Region::new(0..3072).unwrap());
        assert_eq!(region.align_cover(usize::BITS), Region::full());
    }

    #[test]
    fn test_align_cover_end() {
        let region = Region::new_len_end(2000).unwrap();
        assert_eq!(region.align_cover(0), region);
        assert_eq!(region.align_cover(10), Region::new_len_end(2048).unwrap());
        assert_eq!(region.align_cover(usize::BITS), Region::full());
    }

    #[test]
    fn test_align_cover_full() {
        assert_eq!(Region::full().align_cover(10), Region::full());
        assert_eq!(Region::full().align_cover(usize::BITS), Region::full());
    }

    #[test]
    fn test_allocate() {
        let region = Region::new(1000..2000).unwrap();
        assert_eq!(region.allocate(0), Some((1000, Some(region))));
        assert_eq!(region.allocate(500), Some((1000, Region::new(1500..2000))));
        assert_eq!(region.allocate(1000), Some((1000, None)));
        assert_eq!(region.allocate(1001), None);
    }

    #[test]
    fn test_allocate_end() {
        let region = Region::new_len_end(1000).unwrap();
        assert_eq!(
            region.allocate(500),
            Some((region.start(), Region::new_len_end(500)))
        );
        assert_eq!(region.allocate(1000), Some((region.start(), None)));
        assert_eq!(region.allocate(1001), None);
    }

    #[test]
    fn test_allocate_full() {
        assert_eq!(
            Region::full().allocate(1024),
            Some((0, Region::inclusive(1024..=usize::MAX)))
        );

        assert_eq!(
            Region::full().allocate(usize::MAX),
            Some((0, Region::inclusive(usize::MAX..=usize::MAX)))
        );
    }

    #[test]
    fn test_carve_out() {
        let region = Region::new(1000..2000).unwrap();
        assert_eq!(
            region.carve_out(&Region::new(0..1000).unwrap()),
            (None, Some(region))
        );
        assert_eq!(
            region.carve_out(&Region::new(2000..3000).unwrap()),
            (Some(region), None)
        );
        assert_eq!(
            region.carve_out(&Region::new(0..1500).unwrap()),
            (None, Region::new(1500..2000))
        );
        assert_eq!(
            region.carve_out(&Region::new(1000..1500).unwrap()),
            (None, Region::new(1500..2000))
        );
        assert_eq!(
            region.carve_out(&Region::inclusive(1500..=1500).unwrap()),
            (Region::new(1000..1500), Region::new(1501..2000))
        );
        assert_eq!(
            region.carve_out(&Region::new(1200..1500).unwrap()),
            (Region::new(1000..1200), Region::new(1500..2000))
        );
        assert_eq!(
            region.carve_out(&Region::inclusive(1500..=usize::MAX).unwrap()),
            (Region::new(1000..1500), None)
        );
        assert_eq!(region.carve_out(&region), (None, None));
    }

    #[test]
    fn test_carve_out_end() {
        let region = Region::new_len_end(1000).unwrap();
        assert_eq!(
            region.carve_out(&Region::new(0..1000).unwrap()),
            (None, Some(region))
        );
        assert_eq!(
            region.carve_out(&Region::new(0..usize::MAX - 499).unwrap()),
            (None, Region::new_len_end(500))
        );
        assert_eq!(
            region.carve_out(&Region::new_len_end(500).unwrap()),
            (Region::new(usize::MAX - 999..usize::MAX - 499), None)
        );
        assert_eq!(region.carve_out(&region), (None, None));
    }

    #[test]
    fn test_carve_out_full() {
        assert_eq!(
            Region::full().carve_out(&Region::new(0..1500).unwrap()),
            (None, Region::inclusive(1500..=usize::MAX))
        );
        assert_eq!(
            Region::full().carve_out(&Region::new(1000..1500).unwrap()),
            (Region::new(0..1000), Region::inclusive(1500..=usize::MAX))
        );
        assert_eq!(
            Region::full().carve_out(&Region::inclusive(1500..=usize::MAX).unwrap()),
            (Region::new(0..1500), None)
        );
        assert_eq!(
            Region::full().carve_out(&Region::inclusive(usize::MAX..=usize::MAX).unwrap()),
            (Region::new(0..usize::MAX), None)
        );
        assert_eq!(Region::full().carve_out(&Region::full()), (None, None));
    }
}