//! Virtual memory allocator for the Yggdrasil kernel.
//!
//! The allocator uses a [DiscreteRangeMap] to track the memory regions and allows attaching
//! metadata values to each region. "Touching" region coalescing is enabled through the [Eq] trait
//! implemented for [RangeData].

#![deny(missing_docs)]
#![no_std]
#![feature(linked_list_cursors, let_chains, btree_extract_if)]

extern crate alloc;

use core::ops::Range;

use discrete_range_map::{DiscreteRangeMap, InclusiveInterval, InclusiveRange};
use yggdrasil_abi::error::Error;

#[cfg(any(target_pointer_width = "64", rust_analyzer))]
type PfnIndex = u64;
#[cfg(any(target_pointer_width = "32", rust_analyzer))]
type PfnIndex = u32;

/// Metadata associated with an allocated memory region. The [Eq] trait is used to coalesce "equal"
/// regions if they "touch".
pub trait RangeData: Eq + Clone {}

fn ie(from: PfnIndex, to: PfnIndex) -> InclusiveInterval<PfnIndex> {
    InclusiveInterval::from(from..to)
}

/// Main virtual memory allocator
#[derive(Clone)]
pub struct VirtualMemoryAllocator<D: RangeData> {
    map: DiscreteRangeMap<PfnIndex, InclusiveInterval<PfnIndex>, D>,
    outer_range: InclusiveInterval<PfnIndex>,
}

impl<D: RangeData> VirtualMemoryAllocator<D> {
    /// Creates a new virtual memory allocator, bounded by `lower_limit_pfn..upper_limit_pfn`
    pub fn new(lower_limit_pfn: usize, upper_limit_pfn: usize) -> Self {
        Self {
            map: DiscreteRangeMap::new(),
            outer_range: ie(lower_limit_pfn as _, upper_limit_pfn as _),
        }
    }

    /// Returns an iterator over the regions within the allocator
    pub fn regions(&self) -> impl Iterator<Item = (Range<usize>, &D)> {
        self.map.iter().map(|(range, data)| {
            let range = range.start() as usize..range.end() as usize + 1;
            (range, data)
        })
    }

    /// Allocates a contiguous range of virtual address space and associates metadata with it
    pub fn allocate(&mut self, page_count: usize, data: D) -> Result<usize, Error> {
        let start_pfn = self
            .map
            .gaps_trimmed(self.outer_range)
            .find_map(|range| {
                if range.size() >= page_count as _ {
                    Some(range.start() as usize)
                } else {
                    None
                }
            })
            .ok_or(Error::OutOfMemory)?;

        // Should not fail
        self.insert(start_pfn, page_count, data)?;

        Ok(start_pfn)
    }

    /// Retrieves an entry corresponding to given PF#
    pub fn get(&self, pfn: usize) -> Option<(Range<usize>, &D)> {
        match self.map.get_entry_at_point(pfn as u64) {
            Ok((range, data)) => {
                let range = range.start() as usize..range.end() as usize + 1;
                Some((range, data))
            }
            Err(_) => None,
        }
    }

    /// Tries to insert given PF# range with its associated metadata as allocated memory,
    /// returning [Error] if requested range overlaps any existing allocated ranges
    pub fn insert(&mut self, start_pfn: usize, page_count: usize, data: D) -> Result<(), Error> {
        let end_pfn = (start_pfn + page_count) as PfnIndex;
        let start_pfn = start_pfn as PfnIndex;

        self.map
            .insert_merge_touching_if_values_equal(ie(start_pfn, end_pfn), data)
            .map_err(|_| Error::AlreadyExists)?;

        Ok(())
    }

    /// Releases any pages overlapping the requested range, calling `release` on the ranges
    pub fn free<F: FnMut(usize, Range<usize>, D) -> Result<(), Error>>(
        &mut self,
        start_pfn: usize,
        page_count: usize,
        mut release: F,
    ) -> Result<(), Error> {
        let end_pfn = (start_pfn + page_count) as PfnIndex;
        let start_pfn = start_pfn as PfnIndex;

        self.map
            .cut_with_origin(ie(start_pfn, end_pfn))
            .try_for_each(|(origin, range, data)| {
                let range = range.start() as usize..range.end() as usize + 1;
                release(origin.start() as _, range, data)
            })
    }

    /// Removes all allocations, invoking a function on each of them
    pub fn clear<F: FnMut(Range<usize>, D) -> Result<(), Error>>(
        &mut self,
        mut release: F,
    ) -> Result<(), Error> {
        self.map.drain().try_for_each(|(range, data)| {
            let range = range.start() as usize..range.end() as usize + 1;
            release(range, data)
        })
    }
}
