// 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;
    type Base: Ord + Copy;
    type Length: Ord + Copy;
    type Alignment: Copy;

    /// Get base address
    fn base(&self) -> Self::Base;

    // Get length
    fn length(&self) -> Self::Length;

    /// 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: Self::Base, length: Self::Length) -> bool;

    /// Try to allocate a base address from the region. If there's no alignment specified, the
    /// function will check if the requested length fits into the buffer and if yes, it will
    /// return the base of the region.
    /// If an alignment is specified, it will call try_alloc_aligned which has to be implemeted
    /// by the trait users.
    /// This function is part of the trait so the trait users can override it if they want to
    /// achieve a better implementation specific behavior.
    fn try_alloc(
        &self,
        length: Self::Length,
        alignment: Option<Self::Alignment>,
    ) -> Option<Self::Base> {
        if let Some(alignment) = alignment {
            self.try_alloc_aligned(length, alignment)
        } else if length <= self.length() {
            Some(self.base())
        } else {
            None
        }
    }

    /// Try to allocate aligned address from the region.
    fn try_alloc_aligned(
        &self,
        length: Self::Length,
        alignment: Self::Alignment,
    ) -> Option<Self::Base>;

    /// 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: Self::Base,
        length: Self::Length,
        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: T::Length,
        resource: T::Resource,
        alignment: Option<T::Alignment>,
    ) -> Result<T, RegionPoolError> {
        let region_to_allocate_from = self
            .regions
            .iter()
            .enumerate()
            .filter(|(_index, region)| {
                !region.used() && region.try_alloc(length, alignment).is_some()
            })
            .min_by_key(|(_index_a, region)| region.length());

        if let Some((index, region)) = region_to_allocate_from {
            // The previous call to try_alloc ensures that the following unwrap will not fail.
            let base = region.try_alloc(length, alignment).unwrap();
            let (new_region, split_regions) = region.create_split(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: T::Base,
        length: T::Length,
        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: T::Base,
        length: T::Length,
    ) -> 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)]
mod tests {
    use super::*;

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

    impl RegionExample {
        fn new(base: usize, length: usize, used: bool) -> Self {
            Self { base, length, used }
        }
    }

    impl Region for RegionExample {
        type Resource = ();
        type Base = usize;
        type Length = usize;
        type Alignment = usize;

        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_alloc_aligned(
            &self,
            length: Self::Length,
            alignment: Self::Alignment,
        ) -> Option<Self::Base> {
            let aligned_base = self.base.next_multiple_of(alignment);
            let base_offset = aligned_base.checked_sub(self.base)?;

            let required_length = base_offset.checked_add(length)?;
            if required_length <= self.length {
                Some(aligned_base)
            } else {
                None
            }
        }

        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_alloc() {
        let region = RegionExample::new(0x4000_1000, 0x1_0000, false);

        assert_eq!(Some(0x4000_1000), region.try_alloc(0x1000, None));
        assert_eq!(Some(0x4000_2000), region.try_alloc(0x1000, Some(0x2000)));
        assert_eq!(None, region.try_alloc(0x1000, Some(0x10_0000)));
    }

    #[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, (), None);
        assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x1000, true)), res);
        let res = pool.allocate(0x1_0000, (), None);
        assert_eq!(Ok(RegionExample::new(0x4000_1000, 0x1_0000, true)), res);
        let res = pool.allocate(0x2000_0000, (), None);
        assert_eq!(Err(RegionPoolError::NoSpace), res);
    }

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

        let res = pool.allocate(0x1000, (), Some(0x2000));
        assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x1000, true)), res);
        let res = pool.allocate(0x2000, (), Some(0x1_0000));
        assert_eq!(Ok(RegionExample::new(0x4001_0000, 0x2000, true)), res);
        let res = pool.allocate(0x1000, (), None);
        // This region is allocated from the empty space between the first two regions
        assert_eq!(Ok(RegionExample::new(0x4000_1000, 0x1000, true)), 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, (), None);
        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, (), None);
        assert_eq!(Ok(RegionExample::new(0x4000_0000, 0x1000, true)), res);
        assert_eq!(Ok(()), pool.release(res.unwrap()));

        let res = pool.allocate(0x1_0000, (), None);
        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)
        );
    }
}