src/region.rs

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

//! Module for handling physical and virtual memory regions
//!
//! A region is a continuous memory area in the given memory space.

use alloc::vec::Vec;
use log::debug;

use super::{
    page_pool::{PagePool, Pages},
    region_pool::Region,
};

/// Physical region
///
/// A physical memory region can be in three different state
/// * Unused
/// * Points to a page pool allocated address
/// * Points to a physical address without allocation
pub enum PhysicalRegion {
    Unused,
    Allocated(PagePool, Pages),
    PhysicalAddress(usize),
}

impl PhysicalRegion {
    /// Get physical memory address
    fn get_pa(&self) -> usize {
        match self {
            PhysicalRegion::Unused => panic!("Unused area has no PA"),
            PhysicalRegion::Allocated(_page_pool, pages) => pages.get_pa(),
            PhysicalRegion::PhysicalAddress(pa) => *pa,
        }
    }
}

/// Virtual region
///
/// A virtual memory region has a virtual address, a length and a physical region.
pub struct VirtualRegion {
    va: usize,
    length: usize,
    physical_region: PhysicalRegion,
}

impl VirtualRegion {
    /// Create new virtual memory region without a physical region
    pub fn new(va: usize, length: usize) -> Self {
        Self::new_from_fields(va, length, PhysicalRegion::Unused)
    }

    /// Create virtual region with points to a given physical address
    pub fn new_with_pa(pa: usize, va: usize, length: usize) -> Self {
        Self::new_from_fields(va, length, PhysicalRegion::PhysicalAddress(pa))
    }

    /// Create virtual region by defining all the fields of the object
    fn new_from_fields(va: usize, length: usize, physical_region: PhysicalRegion) -> Self {
        Self {
            va,
            length,
            physical_region,
        }
    }

    /// Get the base address of the linked physical region
    pub fn get_pa(&self) -> usize {
        self.physical_region.get_pa()
    }

    /// Get physical address for a virtual address
    pub fn get_pa_for_va(&self, va: usize) -> usize {
        let offset = va.checked_sub(self.va).unwrap();

        assert!(offset < self.length);
        self.get_pa().checked_add(offset).unwrap()
    }
}

impl Region for VirtualRegion {
    type Resource = PhysicalRegion;
    type Base = usize;
    type Length = usize;

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

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

    fn used(&self) -> bool {
        !matches!(self.physical_region, PhysicalRegion::Unused)
    }

    fn contains(&self, base: usize, length: usize) -> bool {
        if let (Some(end), Some(self_end)) =
            (base.checked_add(length), self.va.checked_add(self.length))
        {
            self.va <= base && end <= self_end
        } else {
            false
        }
    }

    fn try_append(&mut self, other: &Self) -> bool {
        if let (Some(self_end), Some(new_length)) = (
            self.va.checked_add(self.length),
            self.length.checked_add(other.length),
        ) {
            if self_end == other.va {
                // VA is consecutive
                match (&self.physical_region, &other.physical_region) {
                    (PhysicalRegion::Unused, PhysicalRegion::Unused) => {
                        // Unused range can be merged without further conditions
                        self.length = new_length;
                        return true;
                    }
                    (
                        PhysicalRegion::PhysicalAddress(self_pa),
                        PhysicalRegion::PhysicalAddress(other_pa),
                    ) => {
                        // Used ranges can be only merged if the PA doesn't overflow and it is
                        // consecutive
                        if let Some(self_end_pa) = self_pa.checked_add(self.length) {
                            if self_end_pa == *other_pa {
                                self.length = new_length;
                                return true;
                            }
                        }
                    }

                    // PhyisicalRegion::Allocated instances cannot be merged at the moment. Not sure
                    // if it's a valid use case. If needed the pages has to be merged which might
                    // require tricks to invalidate the pages of 'other'.
                    _ => {}
                }
            }
        }

        false
    }

    fn create_split(
        &self,
        base: usize,
        length: usize,
        resource: Option<Self::Resource>,
    ) -> (Self, Vec<Self>) {
        assert!(self.contains(base, length));
        assert!(self.used() != resource.is_some());

        if let Some(physical_region) = resource {
            // Self is unused, setting part of it to used
            let pa = physical_region.get_pa();

            let mut res = Vec::new();
            if self.va != base {
                res.push(Self::new(self.va, base.checked_sub(self.va).unwrap()));
            }

            res.push(Self::new_from_fields(base, length, physical_region));

            let end = base.checked_add(length).unwrap();
            let self_end = self.va.checked_add(self.length).unwrap();
            if end != self_end {
                res.push(Self::new(end, self_end.checked_sub(end).unwrap()));
            }

            (
                Self::new_from_fields(base, length, PhysicalRegion::PhysicalAddress(pa)),
                res,
            )
        } else {
            // Self is used, mark part of it unused
            let mut res = Vec::new();
            if self.va != base {
                let physical_region = match &self.physical_region {
                    PhysicalRegion::Allocated(_page_pool, _pages) => {
                        todo!("Implement Pages::split");
                    }
                    PhysicalRegion::PhysicalAddress(pa) => PhysicalRegion::PhysicalAddress(*pa),
                    _ => {
                        panic!("Splitting unused region by other unused")
                    }
                };

                res.push(Self::new_from_fields(
                    self.va,
                    base.checked_sub(self.va).unwrap(),
                    physical_region,
                ));
            }

            res.push(Self::new(base, length));

            let end = base.checked_add(length).unwrap();
            let self_end = self.va.checked_add(self.length).unwrap();
            if end != self_end {
                let physical_region = match &self.physical_region {
                    PhysicalRegion::Allocated(_page_pool, _pages) => {
                        todo!("Implement Pages::split");
                    }
                    PhysicalRegion::PhysicalAddress(pa) => {
                        let offset = end.checked_sub(self.va).unwrap();
                        PhysicalRegion::PhysicalAddress(pa.checked_add(offset).unwrap())
                    }
                    _ => {
                        panic!("Splitting unused region by other unused")
                    }
                };

                res.push(Self::new_from_fields(
                    end,
                    self_end.checked_sub(end).unwrap(),
                    physical_region,
                ));
            }

            (Self::new(base, length), res)
        }
    }
}

impl Drop for VirtualRegion {
    /// If the virtual region has a linked physical region which was allocated then release the
    /// allocated pages.
    fn drop(&mut self) {
        let mut physical_region = PhysicalRegion::Unused;

        core::mem::swap(&mut self.physical_region, &mut physical_region);

        if let PhysicalRegion::Allocated(page_pool, pages) = physical_region {
            debug!(
                "Dropping physical region with pages: PA={:#010x} VA={:#010x}",
                pages.get_pa(),
                self.base(),
            );

            page_pool.release_pages(pages).unwrap();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::super::page_pool::PagePoolArea;
    use super::*;

    #[test]
    #[should_panic]
    fn test_physical_region_unused() {
        let region = PhysicalRegion::Unused;
        region.get_pa();
    }

    #[test]
    fn test_physical_region() {
        const PA: usize = 0x0123_4567_89ab_cdef;
        const LENGTH: usize = 0x8000_0000_0000;

        static PAGE_POOL_AREA: PagePoolArea<16> = PagePoolArea::new();
        let region =
            PhysicalRegion::Allocated(PagePool::new(&PAGE_POOL_AREA), Pages::new(PA, LENGTH, true));
        assert_eq!(PA, region.get_pa());

        let region = PhysicalRegion::PhysicalAddress(PA);
        assert_eq!(PA, region.get_pa());
    }

    #[test]
    fn test_virtual_region() {
        const VA: usize = 0x0123_4567_89ab_cdef;
        const PA: usize = 0xfedc_ba98_7654_3210;
        const LENGTH: usize = 0x8000_0000_0000;

        let region = VirtualRegion::new(VA, LENGTH);
        assert_eq!(VA, region.va);
        assert_eq!(VA, region.base());
        assert_eq!(LENGTH, region.length);
        assert_eq!(LENGTH, region.length());
        assert!(matches!(region.physical_region, PhysicalRegion::Unused));
        assert!(!region.used());

        let region = VirtualRegion::new_with_pa(PA, VA, LENGTH);
        assert_eq!(VA, region.va);
        assert_eq!(VA, region.base());
        assert_eq!(LENGTH, region.length);
        assert_eq!(LENGTH, region.length());
        assert!(matches!(
            region.physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));
        assert_eq!(PA, region.get_pa());
        assert!(region.used());
    }

    #[test]
    fn test_virtual_region_get_pa_for_va() {
        let region =
            VirtualRegion::new_with_pa(0x8000_0000_0000_0000, 0x4000_0000_0000_0000, 0x1000);
        assert_eq!(
            0x8000_0000_0000_0000,
            region.get_pa_for_va(0x4000_0000_0000_0000)
        );
        assert_eq!(
            0x8000_0000_0000_0001,
            region.get_pa_for_va(0x4000_0000_0000_0001)
        );
        assert_eq!(
            0x8000_0000_0000_0fff,
            region.get_pa_for_va(0x4000_0000_0000_0fff)
        );
    }

    #[test]
    #[should_panic]
    fn test_virtual_region_get_pa_for_va_low_va() {
        let region =
            VirtualRegion::new_with_pa(0x8000_0000_0000_0000, 0x4000_0000_0000_0000, 0x1000);
        region.get_pa_for_va(0x3fff_ffff_ffff_ffff);
    }

    #[test]
    #[should_panic]
    fn test_virtual_region_get_pa_for_va_high_va() {
        let region =
            VirtualRegion::new_with_pa(0x8000_0000_0000_0000, 0x4000_0000_0000_0000, 0x1000);
        region.get_pa_for_va(0x4000_0000_0000_1000);
    }

    #[test]
    fn test_virtual_region_contains() {
        const VA: usize = 0x8000_0000_0000_0000;
        const LENGTH: usize = 0x8000_0000_0000;

        let region_overflow_end = VirtualRegion::new(0x8000_0000_0000_0000, 0x8000_0000_0000_0000);
        assert!(!region_overflow_end.contains(0x8000_0000_0000_0000, 1));

        let region = VirtualRegion::new(0x4000_0000_0000_0000, 0x8000_0000_0000_0000);
        assert!(!region.contains(0x8000_0000_0000_0000, 0x8000_0000_0000_0000));

        assert!(!region.contains(0x4000_0000_0000_0000, 0x8000_0000_0000_0001));
        assert!(!region.contains(0x3fff_ffff_ffff_ffff, 0x8000_0000_0000_0000));
        assert!(region.contains(0x4000_0000_0000_0000, 0x8000_0000_0000_0000));
        assert!(region.contains(0x4000_0000_0000_0000, 0x7fff_ffff_ffff_ffff));
        assert!(region.contains(0x4000_0000_0000_0001, 0x7fff_ffff_ffff_ffff));
    }

    #[test]
    fn test_virtual_region_try_append() {
        // Both unused
        let mut region_unused0 = VirtualRegion::new(0x4000_0000, 0x1000);
        let mut region_unused1 = VirtualRegion::new(0x4000_1000, 0x1000);

        assert!(!region_unused1.try_append(&region_unused0));
        assert_eq!(0x4000_0000, region_unused0.va);
        assert_eq!(0x1000, region_unused0.length);
        assert_eq!(0x4000_1000, region_unused1.va);
        assert_eq!(0x1000, region_unused1.length);

        assert!(region_unused0.try_append(&region_unused1));
        assert_eq!(0x4000_0000, region_unused0.va);
        assert_eq!(0x2000, region_unused0.length);
        assert_eq!(0x4000_1000, region_unused1.va);
        assert_eq!(0x1000, region_unused1.length);

        // Unused and PA region
        let mut region_unused = VirtualRegion::new(0x4000_0000, 0x1000);
        let region_physical = VirtualRegion::new_with_pa(0x8000_0000, 0x4000_1000, 0x1000);
        assert!(!region_unused.try_append(&region_physical));
        assert_eq!(0x4000_0000, region_unused.va);
        assert_eq!(0x1000, region_unused.length);
        assert_eq!(0x4000_1000, region_physical.va);
        assert_eq!(0x1000, region_physical.length);

        // Both PA regions but non-consecutive PA ranges
        let mut region_physical0 = VirtualRegion::new_with_pa(0x8000_0000, 0x4000_0000, 0x1000);
        let region_physical1 = VirtualRegion::new_with_pa(0x9000_0000, 0x4000_1000, 0x1000);
        assert!(!region_physical0.try_append(&region_physical1));
        assert_eq!(0x4000_0000, region_physical0.va);
        assert_eq!(0x1000, region_physical0.length);
        assert_eq!(0x4000_1000, region_physical1.va);
        assert_eq!(0x1000, region_physical1.length);

        // Both PA regions with consecutive PA ranges
        let mut region_physical0 = VirtualRegion::new_with_pa(0x8000_0000, 0x4000_0000, 0x1000);
        let region_physical1 = VirtualRegion::new_with_pa(0x8000_1000, 0x4000_1000, 0x1000);
        assert!(region_physical0.try_append(&region_physical1));
        assert_eq!(0x4000_0000, region_physical0.va);
        assert_eq!(0x2000, region_physical0.length);
        assert_eq!(0x4000_1000, region_physical1.va);
        assert_eq!(0x1000, region_physical1.length);

        // VA overflow
        let mut region_unused0 = VirtualRegion::new(0x8000_0000_0000_0000, 0x8000_0000_0000_0000);
        let mut region_unused1 = VirtualRegion::new(0x4000_1000, 0x1000);

        assert!(!region_unused0.try_append(&region_unused1));
        assert_eq!(0x8000_0000_0000_0000, region_unused0.va);
        assert_eq!(0x8000_0000_0000_0000, region_unused0.length);
        assert_eq!(0x4000_1000, region_unused1.va);
        assert_eq!(0x1000, region_unused1.length);

        assert!(!region_unused1.try_append(&region_unused0));
        assert_eq!(0x8000_0000_0000_0000, region_unused0.va);
        assert_eq!(0x8000_0000_0000_0000, region_unused0.length);
        assert_eq!(0x4000_1000, region_unused1.va);
        assert_eq!(0x1000, region_unused1.length);

        // PA overflow
        let mut region_physical0 =
            VirtualRegion::new_with_pa(0x8000_0000_0000_0000, 0x4000_0000, 0x8000_0000_0000_0000);
        let region_physical1 =
            VirtualRegion::new_with_pa(0x9000_0000, 0x8000_0000_4000_0000, 0x1000);
        assert!(!region_physical0.try_append(&region_physical1));
        assert_eq!(0x4000_0000, region_physical0.va);
        assert_eq!(0x8000_0000_0000_0000, region_physical0.length);
        assert_eq!(0x8000_0000_4000_0000, region_physical1.va);
        assert_eq!(0x1000, region_physical1.length);
    }

    #[test]
    fn test_virtual_region_create_split_by_used() {
        let region_unused = VirtualRegion::new(0x4000_0000, 0x4000);

        // New region at the start
        let (new_region, splitted_regions) = region_unused.create_split(
            0x4000_0000,
            0x1000,
            Some(PhysicalRegion::PhysicalAddress(0x8000_0000)),
        );

        assert_eq!(0x4000_0000, new_region.va);
        assert_eq!(0x1000, new_region.length);
        assert_eq!(0x8000_0000, new_region.get_pa());
        assert!(matches!(
            new_region.physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        assert_eq!(0x4000_0000, splitted_regions[0].va);
        assert_eq!(0x1000, splitted_regions[0].length);
        assert_eq!(0x8000_0000, splitted_regions[0].get_pa());
        assert!(matches!(
            splitted_regions[0].physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        assert_eq!(0x4000_1000, splitted_regions[1].va);
        assert_eq!(0x3000, splitted_regions[1].length);
        assert!(matches!(
            splitted_regions[1].physical_region,
            PhysicalRegion::Unused
        ));

        // New region in the middle
        let (new_region, splitted_regions) = region_unused.create_split(
            0x4000_1000,
            0x1000,
            Some(PhysicalRegion::PhysicalAddress(0x8000_0000)),
        );

        assert_eq!(0x4000_1000, new_region.va);
        assert_eq!(0x1000, new_region.length);
        assert_eq!(0x8000_0000, new_region.get_pa());
        assert!(matches!(
            new_region.physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        assert_eq!(0x4000_0000, splitted_regions[0].va);
        assert_eq!(0x1000, splitted_regions[0].length);
        assert!(matches!(
            splitted_regions[0].physical_region,
            PhysicalRegion::Unused
        ));

        assert_eq!(0x4000_1000, splitted_regions[1].va);
        assert_eq!(0x1000, splitted_regions[1].length);
        assert_eq!(0x8000_0000, splitted_regions[1].get_pa());
        assert!(matches!(
            splitted_regions[1].physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        assert_eq!(0x4000_2000, splitted_regions[2].va);
        assert_eq!(0x2000, splitted_regions[2].length);
        assert!(matches!(
            splitted_regions[2].physical_region,
            PhysicalRegion::Unused
        ));

        // New region at the end
        let (new_region, splitted_regions) = region_unused.create_split(
            0x4000_3000,
            0x1000,
            Some(PhysicalRegion::PhysicalAddress(0x8000_0000)),
        );

        assert_eq!(0x4000_3000, new_region.va);
        assert_eq!(0x1000, new_region.length);
        assert_eq!(0x8000_0000, new_region.get_pa());
        assert!(matches!(
            new_region.physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        assert_eq!(0x4000_0000, splitted_regions[0].va);
        assert_eq!(0x3000, splitted_regions[0].length);
        assert!(matches!(
            splitted_regions[0].physical_region,
            PhysicalRegion::Unused
        ));

        assert_eq!(0x4000_3000, splitted_regions[1].va);
        assert_eq!(0x1000, splitted_regions[1].length);
        assert_eq!(0x8000_0000, splitted_regions[1].get_pa());
        assert!(matches!(
            splitted_regions[1].physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));
    }

    #[test]
    fn test_virtual_region_create_split_by_unused() {
        let region_unused = VirtualRegion::new_with_pa(0x8000_0000, 0x4000_0000, 0x4000);

        // New region at the start
        let (new_region, splitted_regions) = region_unused.create_split(0x4000_0000, 0x1000, None);

        assert_eq!(0x4000_0000, new_region.va);
        assert_eq!(0x1000, new_region.length);
        assert!(matches!(new_region.physical_region, PhysicalRegion::Unused));

        assert_eq!(0x4000_0000, splitted_regions[0].va);
        assert_eq!(0x1000, splitted_regions[0].length);
        assert!(matches!(
            splitted_regions[0].physical_region,
            PhysicalRegion::Unused
        ));

        assert_eq!(0x4000_1000, splitted_regions[1].va);
        assert_eq!(0x3000, splitted_regions[1].length);
        assert_eq!(0x8000_1000, splitted_regions[1].get_pa());
        assert!(matches!(
            splitted_regions[1].physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        // New region in the middle
        let (new_region, splitted_regions) = region_unused.create_split(0x4000_1000, 0x1000, None);

        assert_eq!(0x4000_1000, new_region.va);
        assert_eq!(0x1000, new_region.length);
        assert!(matches!(new_region.physical_region, PhysicalRegion::Unused));

        assert_eq!(0x4000_0000, splitted_regions[0].va);
        assert_eq!(0x1000, splitted_regions[0].length);
        assert_eq!(0x8000_0000, splitted_regions[0].get_pa());
        assert!(matches!(
            splitted_regions[0].physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        assert_eq!(0x4000_1000, splitted_regions[1].va);
        assert_eq!(0x1000, splitted_regions[1].length);
        assert!(matches!(
            splitted_regions[1].physical_region,
            PhysicalRegion::Unused
        ));

        assert_eq!(0x4000_2000, splitted_regions[2].va);
        assert_eq!(0x2000, splitted_regions[2].length);
        assert_eq!(0x8000_2000, splitted_regions[2].get_pa());
        assert!(matches!(
            splitted_regions[2].physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        // New region at the end
        let (new_region, splitted_regions) = region_unused.create_split(0x4000_3000, 0x1000, None);

        assert_eq!(0x4000_3000, new_region.va);
        assert_eq!(0x1000, new_region.length);
        assert!(matches!(new_region.physical_region, PhysicalRegion::Unused));

        assert_eq!(0x4000_0000, splitted_regions[0].va);
        assert_eq!(0x3000, splitted_regions[0].length);
        assert_eq!(0x8000_0000, splitted_regions[0].get_pa());
        assert!(matches!(
            splitted_regions[0].physical_region,
            PhysicalRegion::PhysicalAddress(_)
        ));

        assert_eq!(0x4000_3000, splitted_regions[1].va);
        assert_eq!(0x1000, splitted_regions[1].length);

        assert!(matches!(
            splitted_regions[1].physical_region,
            PhysicalRegion::Unused
        ));
    }

    #[test]
    #[should_panic]
    fn test_virtual_region_does_not_contain() {
        let region = VirtualRegion::new(0x4000_0000, 0x1000);
        region.create_split(
            0x8000_0000,
            0x1000,
            Some(PhysicalRegion::PhysicalAddress(0xc000_0000)),
        );
    }

    #[test]
    #[should_panic]
    fn test_virtual_region_create_split_same_used() {
        let region = VirtualRegion::new(0x4000_0000, 0x1000);
        region.create_split(0x4000_0000, 0x1000, Some(PhysicalRegion::Unused));
    }

    #[test]
    fn test_virtual_region_drop() {
        const PA: usize = 0x0123_4567_89ab_cdef;
        const LENGTH: usize = 0x8000_0000_0000;

        static PAGE_POOL_AREA: PagePoolArea<8192> = PagePoolArea::new();
        let page_pool = PagePool::new(&PAGE_POOL_AREA);
        let page = page_pool.allocate_pages(4096).unwrap();

        let physical_region = PhysicalRegion::Allocated(page_pool, page);

        // Testing physical region drop through virtualregion
        let virtual_region = VirtualRegion::new_from_fields(0x4000_0000, 1000, physical_region);
        drop(virtual_region);
    }
}