src/region_pool.rs

// SPDX-FileCopyrightText: Copyright 2023 Arm Limited and/or its affiliates <open-source-office@arm.com>
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Region pool
//!
//! The region pool component handles regions allocations from memory areas in a generic way.

use core::ops::Range;

use alloc::vec::Vec;

/// Region interface
///
/// This trait provides the necessary information about a region to the `RegionPool`.
pub trait Region: Sized {
    type Resource;

    /// Get base address
    fn base(&self) -> usize;

    // Get length
    fn length(&self) -> usize;

    /// Check if the region is used
    fn used(&self) -> bool;

    /// Check if an area defined by its base address and length is inside the region
    fn contains(&self, base: usize, length: usize) -> bool;

    /// Try append the parameter region. Return true on success.
    fn try_append(&mut self, other: &Self) -> bool;

    /// Split region into multiple regions by the area passed as a parameter. It returns the region
    /// of the area as the first member of the return tuple and all the new sliced regions as the
    /// second member of the tuple.
    /// The function has to handle four cases:
    /// * The area matches the region exactly -> return one region
    /// * The area is at the beginning of the region -> return two areas
    /// * The area is at the end of the region -> return two areas
    /// * The area is in the middle of the region -> return three areas
    fn create_split(
        &self,
        base: usize,
        length: usize,
        resource: Option<Self::Resource>,
    ) -> (Self, Vec<Self>);
}

/// Region pool error type
#[derive(Debug, PartialEq)]
pub enum RegionPoolError {
    NoSpace,
    AlreadyUsed,
    NotFound,
}

/// Region pool
pub struct RegionPool<T: Region> {
    regions: Vec<T>,
}

impl<T: Region> RegionPool<T> {
    /// Create empty region pool
    pub fn new() -> Self {
        Self {
            regions: Vec::new(),
        }
    }

    /// Add a region to the pool
    pub fn add(&mut self, region: T) -> Result<(), RegionPoolError> {
        if !self
            .regions
            .iter()
            .any(|r| r.contains(region.base(), region.length()))
        {
            self.regions.push(region);
            self.regions.sort_by_key(|a| a.base());

            Ok(())
        } else {
            Err(RegionPoolError::AlreadyUsed)
        }
    }

    /// Allocate a region from the pool of a given length. It will select an area in which the
    /// region fits and has the minimal size. Then it splits this area to get the allocated region
    /// with the exact size.
    pub fn allocate(&mut self, length: usize, resource: T::Resource) -> Result<T, RegionPoolError> {
        let region_to_allocate_from = self
            .regions
            .iter()
            .enumerate()
            .filter(|(_index, region)| region.length() >= length && !region.used())
            .min_by_key(|(_index_a, region)| region.length());

        if let Some((index, region)) = region_to_allocate_from {
            let (new_region, split_regions) =
                region.create_split(region.base(), length, Some(resource));
            self.replace(index, split_regions);
            Ok(new_region)
        } else {
            Err(RegionPoolError::NoSpace)
        }
    }

    /// Acquire a region with the given base address and length.
    pub fn acquire(
        &mut self,
        base: usize,
        length: usize,
        resource: T::Resource,
    ) -> Result<T, RegionPoolError> {
        let region_to_acquire_from = self
            .regions
            .iter()
            .enumerate()
            .find(|(_, region)| region.contains(base, length) && !region.used());

        if let Some((index, region)) = region_to_acquire_from {
            let (new_region, split_regions) = region.create_split(base, length, Some(resource));
            self.replace(index, split_regions);
            Ok(new_region)
        } else {
            Err(RegionPoolError::AlreadyUsed)
        }
    }

    /// Release region
    pub fn release(&mut self, region_to_release: T) -> Result<(), RegionPoolError> {
        assert!(region_to_release.used());

        let region_to_release_from = self.regions.iter().enumerate().find(|(_, r)| {
            r.contains(region_to_release.base(), region_to_release.length()) && r.used()
        });

        if let Some((index, region)) = region_to_release_from {
            self.replace(
                index,
                region
                    .create_split(region_to_release.base(), region_to_release.length(), None)
                    .1,
            );
            self.regions.dedup_by(|a, b| b.try_append(a));

            Ok(())
        } else {
            Err(RegionPoolError::NotFound)
        }
    }

    /// Find the region which contains the given area
    pub fn find_containing_region(
        &self,
        base: usize,
        length: usize,
    ) -> Result<&T, RegionPoolError> {
        let region_index = match self.regions.binary_search_by(|r| r.base().cmp(&base)) {
            Ok(exact_index) => Ok(exact_index),
            Err(insert_index) => {
                if insert_index > 0 {
                    Ok(insert_index - 1)
                } else {
                    Err(RegionPoolError::NotFound)
                }
            }
        }?;

        if self.regions[region_index].contains(base, length) {
            Ok(&self.regions[region_index])
        } else {
            Err(RegionPoolError::NotFound)
        }
    }

    fn replace(&mut self, index: usize, replacements: Vec<T>) {
        self.regions.splice(
            Range {
                start: index,
                end: index + 1,
            },
            replacements,
        );
    }
}

#[cfg(test)]
#[derive(Debug, PartialEq, Eq)]
struct RegionExample {
    base: usize,
    length: usize,
    used: bool,
}

#[cfg(test)]
impl RegionExample {
    fn new(base: usize, length: usize, used: bool) -> Self {
        Self { base, length, used }
    }
}

#[cfg(test)]
impl Region for RegionExample {
    type Resource = ();

    fn base(&self) -> usize {
        self.base
    }

    fn length(&self) -> usize {
        self.length
    }

    fn used(&self) -> bool {
        self.used
    }

    fn contains(&self, base: usize, length: usize) -> bool {
        self.base <= base && base + length <= self.base + self.length
    }

    fn try_append(&mut self, other: &Self) -> bool {
        if self.used == other.used && self.base + self.length == other.base {
            self.length += other.length;
            true
        } else {
            false
        }
    }

    fn create_split(
        &self,
        base: usize,
        length: usize,
        resource: Option<Self::Resource>,
    ) -> (Self, Vec<Self>) {
        let mut res = Vec::new();
        if self.base != base {
            res.push(RegionExample::new(self.base, base - self.base, self.used));
        }
        res.push(RegionExample::new(base, length, resource.is_some()));
        if self.base + self.length != base + length {
            res.push(RegionExample::new(
                base + length,
                (self.base + self.length) - (base + length),
                self.used,
            ));
        }

        (RegionExample::new(base, length, resource.is_some()), res)
    }
}

#[test]
fn test_region_contains() {
    let region = RegionExample::new(0x4000_0000, 0x1000_0000, false);

    assert!(region.contains(0x4000_0000, 0x1000_0000));
    assert!(!region.contains(0x4000_0000 - 1, 0x1000_0000 + 1));
    assert!(!region.contains(0x4000_0000, 0x1000_0000 + 1));
    assert!(region.contains(0x4000_0000 + 1, 0x1000_0000 - 1));
}

#[test]
fn test_region_try_append() {
    // Normal append
    let mut region = RegionExample::new(0x4000_0000, 0x1000_0000, false);
    let appending = RegionExample::new(0x5000_0000, 0x0000_1000, false);

    assert!(region.try_append(&appending));
    assert_eq!(RegionExample::new(0x4000_0000, 0x1000_1000, false), region);

    // Different used flags
    let mut region = RegionExample::new(0x4000_0000, 0x1000_0000, false);
    let appending = RegionExample::new(0x5000_0000, 0x0000_1000, true);

    assert!(!region.try_append(&appending));
    assert_eq!(RegionExample::new(0x4000_0000, 0x1000_0000, false), region);

    // Not continious
    let mut region = RegionExample::new(0x4000_0000, 0x1000_0000, false);
    let appending = RegionExample::new(0x5000_1000, 0x0000_1000, false);

    assert!(!region.try_append(&appending));
    assert_eq!(RegionExample::new(0x4000_0000, 0x1000_0000, false), region);
}

#[test]
fn test_region_create_split() {
    // Cut first part
    let region = RegionExample::new(0x4000_0000, 0x1000_0000, false);

    let res = region.create_split(0x4000_0000, 0x0000_1000, Some(())).1;
    assert_eq!(RegionExample::new(0x4000_0000, 0x0000_1000, true), res[0]);
    assert_eq!(RegionExample::new(0x4000_1000, 0x0fff_f000, false), res[1]);

    // Cut last part
    let region = RegionExample::new(0x4000_0000, 0x1000_0000, false);

    let res = region.create_split(0x4fff_f000, 0x0000_1000, Some(())).1;
    assert_eq!(RegionExample::new(0x4000_0000, 0x0fff_f000, false), res[0]);
    assert_eq!(RegionExample::new(0x4fff_f000, 0x0000_1000, true), res[1]);

    // Cut middle part
    let region = RegionExample::new(0x4000_0000, 0x1000_0000, false);

    let res = region.create_split(0x4020_0000, 0x0000_1000, Some(())).1;
    assert_eq!(RegionExample::new(0x4000_0000, 0x0020_0000, false), res[0]);
    assert_eq!(RegionExample::new(0x4020_0000, 0x0000_1000, true), res[1]);
    assert_eq!(RegionExample::new(0x4020_1000, 0x0fdf_f000, false), res[2]);
}

#[test]
fn test_region_pool_add() {
    let mut pool = RegionPool::new();
    assert_eq!(
        Ok(()),
        pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
    );
    assert_eq!(
        Ok(()),
        pool.add(RegionExample::new(0x5000_0000, 0x1000_0000, false))
    );
    assert_eq!(
        Err(RegionPoolError::AlreadyUsed),
        pool.add(RegionExample::new(0x4000_1000, 0x1000, false))
    );
}

#[test]
fn test_pool_allocate() {
    let mut pool = RegionPool::new();
    pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
        .unwrap();

    let res = pool.allocate(0x1000, ());
    assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x1000, true)), res);
    let res = pool.allocate(0x1_0000, ());
    assert_eq!(Ok(RegionExample::new(0x4000_1000, 0x1_0000, true)), res);
    let res = pool.allocate(0x2000_0000, ());
    assert_eq!(Err(RegionPoolError::NoSpace), res);
}

#[test]
fn test_region_pool_acquire() {
    let mut pool = RegionPool::new();
    pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
        .unwrap();

    let res = pool.acquire(0x4000_1000, 0x1000, ());
    assert_eq!(Ok(RegionExample::new(0x4000_1000, 0x1000, true)), res);
    let res = pool.allocate(0x1_0000, ());
    assert_eq!(Ok(RegionExample::new(0x4000_2000, 0x10000, true)), res);
    let res = pool.acquire(0x4000_1000, 0x1000, ());
    assert_eq!(Err(RegionPoolError::AlreadyUsed), res);
}

#[test]
fn test_region_pool_release() {
    let mut pool = RegionPool::new();
    pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
        .unwrap();

    let res = pool.allocate(0x1000, ());
    assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x1000, true)), res);
    assert_eq!(Ok(()), pool.release(res.unwrap()));

    let res = pool.allocate(0x1_0000, ());
    assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x10000, true)), res);
    assert_eq!(Ok(()), pool.release(res.unwrap()));

    assert_eq!(
        Err(RegionPoolError::NotFound),
        pool.release(RegionExample::new(0x4000_0000, 0x1000, true))
    );
}

#[test]
#[should_panic]
fn test_region_pool_release_not_used() {
    let mut pool = RegionPool::new();
    pool.release(RegionExample::new(0x4000_0000, 0x1000, false))
        .unwrap();
}

#[test]
fn test_region_pool_find_containing_region() {
    let mut pool = RegionPool::<RegionExample>::new();
    pool.add(RegionExample::new(0x4000_0000, 0x1000_0000, false))
        .unwrap();

    assert_eq!(
        Err(RegionPoolError::NotFound),
        pool.find_containing_region(0x0000_0000, 0x1000)
    );

    assert_eq!(
        Err(RegionPoolError::NotFound),
        pool.find_containing_region(0x5000_0000, 0x1000)
    );

    assert_eq!(
        Err(RegionPoolError::NotFound),
        pool.find_containing_region(0x4000_0000, 0x2000_0000)
    );

    assert_eq!(
        Ok(&RegionExample::new(0x4000_0000, 0x1000_0000, false)),
        pool.find_containing_region(0x4000_0000, 0x1000_0000)
    );

    assert_eq!(
        Ok(&RegionExample::new(0x4000_0000, 0x1000_0000, false)),
        pool.find_containing_region(0x4000_1000, 0x1000)
    );
}