k-os/crates/domain/src/memory/frame/mod.rs

//! Physical frame ownership. The kernel's single source of truth for what
//! every frame of the user pool is used for and who is charged for it.

use types::{
    address::{PhysicalAddress, VirtualAddress},
    id::{ProcessId, RegionId},
    mapping::PagePerms,
};

#[cfg(test)]
#[path = "tests.rs"]
mod tests;

#[cfg(kani)]
#[path = "proofs.rs"]
mod proofs;

pub const PAGE_SIZE: usize = 4096;

/// Frames the kernel can track. The board's user pool must fit.
pub const MAX_FRAMES: usize = 4096;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FrameUse {
    /// Mapped exactly once, at `page` in `owner`'s address space.
    User {
        owner: ProcessId,
        page: u32,
        perms: PagePerms,
    },
    /// Holds part of `owner`'s page tables.
    PageTable { owner: ProcessId },
    /// Backs a shared memory region.
    Region { region: RegionId },
}

impl FrameUse {
    /// A user mapping at `va`, which must be page aligned.
    pub fn user(owner: ProcessId, va: VirtualAddress, perms: PagePerms) -> Option<Self> {
        if !va.raw().is_multiple_of(PAGE_SIZE) {
            return None;
        }
        let page = u32::try_from(va.raw() / PAGE_SIZE).ok()?;
        Some(Self::User { owner, page, perms })
    }
}

pub struct FrameTable<const N: usize = MAX_FRAMES> {
    base: PhysicalAddress,
    total: usize,
    frames: [Option<FrameUse>; N],
    free: usize,
}

impl<const N: usize> FrameTable<N> {
    pub fn new(base: PhysicalAddress, total: usize) -> Self {
        let total = total.min(N);
        Self {
            base,
            total,
            frames: [None; N],
            free: total,
        }
    }

    pub fn base(&self) -> PhysicalAddress {
        self.base
    }

    pub fn capacity(&self) -> usize {
        self.total
    }

    pub fn free_count(&self) -> usize {
        self.free
    }

    fn index_of(&self, pa: PhysicalAddress) -> Option<usize> {
        let offset = pa.raw().checked_sub(self.base.raw())?;
        if !offset.is_multiple_of(PAGE_SIZE) {
            return None;
        }
        let idx = offset / PAGE_SIZE;
        (idx < self.total).then_some(idx)
    }

    fn address_of(&self, idx: usize) -> PhysicalAddress {
        self.base.offset(idx)
    }

    pub fn get(&self, pa: PhysicalAddress) -> Option<&FrameUse> {
        self.frames[self.index_of(pa)?].as_ref()
    }

    pub fn alloc(&mut self, usage: FrameUse) -> Option<PhysicalAddress> {
        let idx = self.frames[..self.total].iter().position(Option::is_none)?;
        self.frames[idx] = Some(usage);
        self.free -= 1;
        Some(self.address_of(idx))
    }

    /// Allocates `count` physically adjacent frames; frame `i` gets `usage(i)`.
    pub fn alloc_contiguous(
        &mut self,
        count: usize,
        usage: impl Fn(usize) -> FrameUse,
    ) -> Option<PhysicalAddress> {
        if count == 0 {
            return None;
        }
        let mut run = 0;
        let mut start = None;
        for idx in 0..self.total {
            if self.frames[idx].is_some() {
                run = 0;
                continue;
            }
            run += 1;
            if run == count {
                start = Some(idx + 1 - count);
                break;
            }
        }
        let start = start?;
        for i in 0..count {
            self.frames[start + i] = Some(usage(i));
        }
        self.free -= count;
        Some(self.address_of(start))
    }

    /// Frees an allocated frame and reports what it was used for.
    pub fn free(&mut self, pa: PhysicalAddress) -> Option<FrameUse> {
        let idx = self.index_of(pa)?;
        let usage = self.frames[idx].take()?;
        self.free += 1;
        Some(usage)
    }

    /// Changes the use of an allocated frame (e.g. moving it to another
    /// process). Fails on free frames.
    pub fn retag(&mut self, pa: PhysicalAddress, usage: FrameUse) -> bool {
        match self.index_of(pa).and_then(|i| self.frames[i].as_mut()) {
            Some(slot) => {
                *slot = usage;
                true
            }
            None => false,
        }
    }

    /// The frame mapped at `va` in `owner`, if `owner` holds it as a user page.
    pub fn find_user(
        &self,
        owner: ProcessId,
        va: VirtualAddress,
    ) -> Option<(PhysicalAddress, PagePerms)> {
        if !va.raw().is_multiple_of(PAGE_SIZE) {
            return None;
        }
        let page = u32::try_from(va.raw() / PAGE_SIZE).ok()?;
        self.frames[..self.total]
            .iter()
            .enumerate()
            .find_map(|(idx, f)| match f {
                Some(FrameUse::User {
                    owner: o,
                    page: p,
                    perms,
                }) if *o == owner && *p == page => Some((self.address_of(idx), *perms)),
                _ => None,
            })
    }

    /// Every allocated frame and its use.
    pub fn iter(&self) -> impl Iterator<Item = (PhysicalAddress, &FrameUse)> + '_ {
        self.frames[..self.total]
            .iter()
            .enumerate()
            .filter_map(|(idx, f)| f.as_ref().map(|u| (self.address_of(idx), u)))
    }

    /// User pages of `owner` as (frame, virtual address) pairs.
    pub fn user_frames_of(
        &self,
        owner: ProcessId,
    ) -> impl Iterator<Item = (PhysicalAddress, VirtualAddress)> + '_ {
        self.frames[..self.total]
            .iter()
            .enumerate()
            .filter_map(move |(idx, f)| match f {
                Some(FrameUse::User { owner: o, page, .. }) if *o == owner => Some((
                    self.address_of(idx),
                    VirtualAddress::from_raw(*page as usize * PAGE_SIZE),
                )),
                _ => None,
            })
    }

    /// Frees every user and page-table frame charged to `owner`; returns
    /// how many were freed.
    pub fn free_owned_by(&mut self, owner: ProcessId) -> usize {
        let mut freed = 0;
        for slot in &mut self.frames[..self.total] {
            let owned = matches!(
                slot,
                Some(FrameUse::User { owner: o, .. } | FrameUse::PageTable { owner: o }) if *o == owner
            );
            if owned {
                *slot = None;
                freed += 1;
            }
        }
        self.free += freed;
        freed
    }
}