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;

    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)]
use super::page_pool::PagePoolArea;

#[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);
}