ahrightos-memory 0.1.0

#![cfg_attr(not(test), no_std)]

pub const PAGE_SIZE: usize = 4096;
pub const PAGE_TABLE_ENTRIES: usize = 512;
pub const VM_REGION_MAX: usize = 32;
pub const VM_REGION_MAX_PAGES: usize = 128;

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct VirtAddr(usize);

impl VirtAddr {
    pub const fn new(addr: usize) -> Self {
        Self(addr)
    }

    pub const fn as_usize(self) -> usize {
        self.0
    }

    pub fn as_mut_ptr<T>(self) -> *mut T {
        self.0 as *mut T
    }

    pub const fn page_offset(self) -> usize {
        self.0 & 0xFFF
    }

    pub const fn vpn0(self) -> usize {
        (self.0 >> 12) & 0x1FF
    }

    pub const fn vpn1(self) -> usize {
        (self.0 >> 21) & 0x1FF
    }

    pub const fn vpn2(self) -> usize {
        (self.0 >> 30) & 0x1FF
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
#[repr(transparent)]
pub struct PhysAddr(usize);

impl PhysAddr {
    pub const fn new(addr: usize) -> Self {
        Self(addr)
    }

    pub const fn as_usize(self) -> usize {
        self.0
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Page {
    start_address: VirtAddr,
}

impl Page {
    pub fn containing_address(addr: VirtAddr) -> Self {
        Self {
            start_address: VirtAddr::new(addr.as_usize() & !(PAGE_SIZE - 1)),
        }
    }

    pub fn start_address(self) -> VirtAddr {
        self.start_address
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PhysFrame {
    start_address: PhysAddr,
}

impl PhysFrame {
    pub fn containing_address(addr: PhysAddr) -> Self {
        Self {
            start_address: PhysAddr::new(addr.as_usize() & !(PAGE_SIZE - 1)),
        }
    }

    pub fn start_address(self) -> PhysAddr {
        self.start_address
    }
}

#[repr(u64)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PageTableFlags {
    Valid = 1 << 0,
    Read = 1 << 1,
    Write = 1 << 2,
    Execute = 1 << 3,
    User = 1 << 4,
    Global = 1 << 5,
    Accessed = 1 << 6,
    Dirty = 1 << 7,
}

#[repr(transparent)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct PageTableEntry {
    entry: usize,
}

impl PageTableEntry {
    pub const fn new() -> Self {
        Self { entry: 0 }
    }

    pub fn is_valid(&self) -> bool {
        (self.entry & PageTableFlags::Valid as usize) != 0
    }

    pub fn is_leaf(&self) -> bool {
        (self.entry & 0xE) != 0
    }

    pub fn ppn(&self) -> usize {
        (self.entry >> 10) & 0xFFF_FFFF_FFFF
    }

    pub fn phys_addr(&self) -> PhysAddr {
        PhysAddr::new(self.ppn() << 12)
    }

    pub fn set(&mut self, ppn: usize, flags: usize) {
        self.entry = (ppn << 10) | flags;
    }

    pub fn flags(&self) -> usize {
        self.entry & 0xFF
    }
}

impl Default for PageTableEntry {
    fn default() -> Self {
        Self::new()
    }
}

#[repr(align(4096))]
pub struct PageTable {
    pub entries: [PageTableEntry; PAGE_TABLE_ENTRIES],
}

impl PageTable {
    pub const fn new() -> Self {
        Self {
            entries: [PageTableEntry::new(); PAGE_TABLE_ENTRIES],
        }
    }

    pub fn get_entry(&self, index: usize) -> &PageTableEntry {
        &self.entries[index]
    }

    pub fn get_entry_mut(&mut self, index: usize) -> &mut PageTableEntry {
        &mut self.entries[index]
    }

    pub fn zero(&mut self) {
        for entry in self.entries.iter_mut() {
            *entry = PageTableEntry::new();
        }
    }
}

impl Default for PageTable {
    fn default() -> Self {
        Self::new()
    }
}

pub struct SimpleFrameAllocator {
    start_frame: usize,
    next_frame: usize,
    end_frame: usize,
    free_stack: [usize; 256],
    free_top: usize,
}

impl SimpleFrameAllocator {
    pub fn new(kernel_end: usize, memory_end: usize) -> Self {
        let next_frame = kernel_end.div_ceil(PAGE_SIZE);
        let end_frame = memory_end / PAGE_SIZE;
        Self {
            start_frame: next_frame,
            next_frame,
            end_frame,
            free_stack: [0; 256],
            free_top: 0,
        }
    }

    pub fn allocate(&mut self) -> Option<PhysFrame> {
        if self.free_top > 0 {
            self.free_top -= 1;
            let frame_num = self.free_stack[self.free_top];
            return Some(PhysFrame::containing_address(PhysAddr::new(
                frame_num * PAGE_SIZE,
            )));
        }

        if self.next_frame >= self.end_frame {
            return None;
        }

        let frame = PhysFrame::containing_address(PhysAddr::new(self.next_frame * PAGE_SIZE));
        self.next_frame += 1;
        Some(frame)
    }

    pub fn deallocate(&mut self, frame: PhysFrame) {
        let frame_num = frame.start_address().as_usize() / PAGE_SIZE;
        if frame_num < self.start_frame
            || frame_num >= self.next_frame
            || self.free_top >= self.free_stack.len()
            || self.free_stack[..self.free_top].contains(&frame_num)
        {
            return;
        }
        self.free_stack[self.free_top] = frame_num;
        self.free_top += 1;
    }

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

    pub fn set_next_frame(&mut self, value: usize) {
        self.next_frame = value;
    }
}

pub struct MemoryManager {
    pub frame_allocator: SimpleFrameAllocator,
}

impl MemoryManager {
    pub fn new(kernel_end: usize, memory_end: usize) -> Self {
        Self {
            frame_allocator: SimpleFrameAllocator::new(kernel_end, memory_end),
        }
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct VmPermissions {
    bits: u8,
}

impl VmPermissions {
    pub const READ: Self = Self { bits: 1 << 0 };
    pub const WRITE: Self = Self { bits: 1 << 1 };
    pub const EXECUTE: Self = Self { bits: 1 << 2 };
    pub const USER: Self = Self { bits: 1 << 3 };

    pub const fn empty() -> Self {
        Self { bits: 0 }
    }

    pub const fn from_bits(bits: u8) -> Self {
        Self { bits }
    }

    pub const fn bits(self) -> u8 {
        self.bits
    }

    pub const fn union(self, other: Self) -> Self {
        Self {
            bits: self.bits | other.bits,
        }
    }

    pub const fn contains(self, other: Self) -> bool {
        (self.bits & other.bits) == other.bits
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum VmRegionKind {
    /// Anonymous mapping, suitable for mmap-like heap/scratch regions.
    Anonymous,
    /// File-backed mapping bookkeeping. Actual file page-in is a later kernel step.
    FileBacked,
    /// Grow-down user stack reservation with a bounded maximum size.
    Stack { max_pages: usize },
    /// Swap-backed anonymous reservation. Swap I/O is a later kernel step.
    SwapBacked,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct VmRegion {
    start: VirtAddr,
    end: VirtAddr,
    kind: VmRegionKind,
    permissions: VmPermissions,
    committed_pages: u128,
}

impl VmRegion {
    pub fn new(
        start: VirtAddr,
        len: usize,
        kind: VmRegionKind,
        permissions: VmPermissions,
    ) -> Result<Self, &'static str> {
        if len == 0 {
            return Err("VM region length must be non-zero");
        }
        if start.page_offset() != 0 || !len.is_multiple_of(PAGE_SIZE) {
            return Err("VM region must be page aligned");
        }
        let page_count = len / PAGE_SIZE;
        if page_count == 0 || page_count > VM_REGION_MAX_PAGES {
            return Err("VM region page count is outside teaching bounds");
        }
        if let VmRegionKind::Stack { max_pages } = kind {
            if max_pages == 0 {
                return Err("Stack maximum pages must be non-zero");
            }
            if max_pages > VM_REGION_MAX_PAGES {
                return Err("Stack maximum pages exceed teaching bounds");
            }
            if page_count > max_pages {
                return Err("Stack region length exceeds maximum pages");
            }
        }
        let end = start
            .as_usize()
            .checked_add(len)
            .ok_or("VM region end overflows")?;
        Ok(Self {
            start,
            end: VirtAddr::new(end),
            kind,
            permissions,
            committed_pages: 0,
        })
    }

    pub const fn start(&self) -> VirtAddr {
        self.start
    }

    pub const fn end(&self) -> VirtAddr {
        self.end
    }

    pub const fn kind(&self) -> VmRegionKind {
        self.kind
    }

    pub const fn permissions(&self) -> VmPermissions {
        self.permissions
    }

    pub fn len(&self) -> usize {
        self.end.as_usize() - self.start.as_usize()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn page_count(&self) -> usize {
        self.len() / PAGE_SIZE
    }

    pub fn contains(&self, addr: VirtAddr) -> bool {
        addr.as_usize() >= self.start.as_usize() && addr.as_usize() < self.end.as_usize()
    }

    pub fn overlaps(&self, other: &VmRegion) -> bool {
        self.start.as_usize() < other.end.as_usize() && other.start.as_usize() < self.end.as_usize()
    }

    pub fn committed_page_count(&self) -> usize {
        self.committed_pages.count_ones() as usize
    }

    pub fn is_page_committed(&self, addr: VirtAddr) -> bool {
        let Some(index) = self.page_index(addr) else {
            return false;
        };
        (self.committed_pages & (1u128 << index)) != 0
    }

    pub fn commit_page(&mut self, addr: VirtAddr) -> Result<bool, &'static str> {
        let index = self
            .page_index(addr)
            .ok_or("VM page is outside the region")?;
        let mask = 1u128 << index;
        let was_new = (self.committed_pages & mask) == 0;
        self.committed_pages |= mask;
        Ok(was_new)
    }

    pub fn grow_down_to(&mut self, fault_addr: VirtAddr) -> Result<bool, &'static str> {
        let VmRegionKind::Stack { max_pages } = self.kind else {
            return Err("Only stack VM regions can grow down");
        };
        let fault_page = align_down(fault_addr.as_usize());
        if fault_page >= self.start.as_usize() {
            return Ok(false);
        }
        let added_pages = (self.start.as_usize() - fault_page) / PAGE_SIZE;
        let max_len = max_pages
            .checked_mul(PAGE_SIZE)
            .ok_or("Stack maximum size overflows")?;
        let new_len = self.end.as_usize() - fault_page;
        if new_len > max_len || new_len / PAGE_SIZE > VM_REGION_MAX_PAGES {
            return Err("Stack growth exceeds reserved teaching bound");
        }
        self.start = VirtAddr::new(fault_page);
        self.committed_pages <<= added_pages;
        Ok(true)
    }

    fn page_index(&self, addr: VirtAddr) -> Option<usize> {
        if !self.contains(addr) {
            return None;
        }
        Some((align_down(addr.as_usize()) - self.start.as_usize()) / PAGE_SIZE)
    }
}

#[derive(Debug, Clone)]
pub struct VmRegionTable {
    regions: [Option<VmRegion>; VM_REGION_MAX],
}

impl VmRegionTable {
    pub const fn new() -> Self {
        Self {
            regions: [None; VM_REGION_MAX],
        }
    }

    pub fn insert(&mut self, region: VmRegion) -> Result<(), &'static str> {
        if self
            .regions
            .iter()
            .flatten()
            .any(|existing| existing.overlaps(&region))
        {
            return Err("VM region overlaps an existing mapping");
        }
        let slot = self
            .regions
            .iter_mut()
            .find(|slot| slot.is_none())
            .ok_or("VM region table is full")?;
        *slot = Some(region);
        Ok(())
    }

    pub fn find(&self, addr: VirtAddr) -> Option<&VmRegion> {
        self.regions
            .iter()
            .flatten()
            .find(|region| region.contains(addr))
    }

    pub fn find_mut(&mut self, addr: VirtAddr) -> Option<&mut VmRegion> {
        self.regions
            .iter_mut()
            .flatten()
            .find(|region| region.contains(addr))
    }

    pub fn remove(&mut self, start: VirtAddr) -> Option<VmRegion> {
        let slot = self
            .regions
            .iter_mut()
            .find(|region| region.map(|r| r.start == start).unwrap_or(false))?;
        slot.take()
    }

    pub fn len(&self) -> usize {
        self.regions.iter().flatten().count()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    pub fn commit_page(&mut self, addr: VirtAddr) -> Result<bool, &'static str> {
        let region = self.find_mut(addr).ok_or("No VM region contains page")?;
        region.commit_page(addr)
    }

    pub fn grow_stack_for_fault(&mut self, fault_addr: VirtAddr) -> Result<bool, &'static str> {
        let fault_page = align_down(fault_addr.as_usize());
        let mut stack_slot = None;
        let mut proposed_start = 0usize;
        let mut proposed_end = 0usize;

        for (index, region) in self.regions.iter().enumerate() {
            let Some(region) = region else {
                continue;
            };
            let VmRegionKind::Stack { max_pages } = region.kind else {
                continue;
            };
            let lower_bound = region
                .end
                .as_usize()
                .saturating_sub(max_pages.saturating_mul(PAGE_SIZE));
            if fault_page >= lower_bound && fault_page < region.start.as_usize() {
                stack_slot = Some(index);
                proposed_start = fault_page;
                proposed_end = region.end.as_usize();
                break;
            }
        }

        let stack_slot =
            stack_slot.ok_or("Fault address is not inside a growable stack reservation")?;
        if self.regions.iter().enumerate().any(|(index, region)| {
            index != stack_slot
                && region
                    .as_ref()
                    .is_some_and(|existing| ranges_overlap(proposed_start, proposed_end, existing))
        }) {
            return Err("Stack growth would overlap an existing VM region");
        }

        self.regions[stack_slot]
            .as_mut()
            .expect("stack slot was selected")
            .grow_down_to(fault_addr)
    }
}

impl Default for VmRegionTable {
    fn default() -> Self {
        Self::new()
    }
}

const fn align_down(addr: usize) -> usize {
    addr & !(PAGE_SIZE - 1)
}

fn ranges_overlap(start: usize, end: usize, region: &VmRegion) -> bool {
    start < region.end.as_usize() && region.start.as_usize() < end
}

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

    #[test]
    fn virt_addr_decomposes_sv39_parts() {
        let addr = VirtAddr::new(0x1234_5678);
        assert_eq!(addr.page_offset(), 0x678);
        assert_eq!(addr.vpn0(), 0x145);
    }

    #[test]
    fn frame_allocator_reuses_deallocated_frames() {
        let mut allocator = SimpleFrameAllocator::new(0x8020_0000, 0x8030_0000);
        let f1 = allocator.allocate().expect("first frame");
        let f2 = allocator.allocate().expect("second frame");
        assert_ne!(f1.start_address(), f2.start_address());

        let f1_addr = f1.start_address();
        allocator.deallocate(f1);
        let f3 = allocator.allocate().expect("reused frame");
        assert_eq!(f3.start_address(), f1_addr);
    }

    #[test]
    fn page_table_entries_encode_ppn_and_flags() {
        let mut entry = PageTableEntry::new();
        let flags = PageTableFlags::Valid as usize | PageTableFlags::Read as usize;
        entry.set(0x80200, flags);
        assert!(entry.is_valid());
        assert!(entry.is_leaf());
        assert_eq!(entry.ppn(), 0x80200);
        assert_eq!(entry.phys_addr(), PhysAddr::new(0x8020_0000));
        assert_eq!(entry.flags(), flags);
    }

    #[test]
    fn vm_region_table_tracks_lazy_mmap_commitment() {
        let permissions = VmPermissions::READ
            .union(VmPermissions::WRITE)
            .union(VmPermissions::USER);
        let region = VmRegion::new(
            VirtAddr::new(0x4000_0000),
            PAGE_SIZE * 4,
            VmRegionKind::Anonymous,
            permissions,
        )
        .expect("anonymous region");

        let mut table = VmRegionTable::new();
        table.insert(region).expect("insert region");
        assert_eq!(table.len(), 1);
        assert!(table.find(VirtAddr::new(0x4000_0123)).is_some());

        assert!(table
            .commit_page(VirtAddr::new(0x4000_0123))
            .expect("first commit"));
        assert!(!table
            .commit_page(VirtAddr::new(0x4000_0abc))
            .expect("duplicate commit"));
        let region = table.find(VirtAddr::new(0x4000_0123)).unwrap();
        assert_eq!(region.committed_page_count(), 1);
        assert!(region.is_page_committed(VirtAddr::new(0x4000_0000)));
        assert!(!region.is_page_committed(VirtAddr::new(0x4000_1000)));
    }

    #[test]
    fn vm_region_table_rejects_overlap_and_bounds_stack_growth() {
        let permissions = VmPermissions::READ
            .union(VmPermissions::WRITE)
            .union(VmPermissions::USER);
        let mut table = VmRegionTable::new();
        table
            .insert(
                VmRegion::new(
                    VirtAddr::new(0x7000_0000),
                    PAGE_SIZE,
                    VmRegionKind::Stack { max_pages: 4 },
                    permissions,
                )
                .unwrap(),
            )
            .unwrap();

        let overlapping = VmRegion::new(
            VirtAddr::new(0x7000_0000),
            PAGE_SIZE,
            VmRegionKind::Anonymous,
            permissions,
        )
        .unwrap();
        assert!(table.insert(overlapping).is_err());

        assert!(table
            .grow_stack_for_fault(VirtAddr::new(0x6fff_f000))
            .expect("grow one page"));
        assert!(table.find(VirtAddr::new(0x6fff_f000)).is_some());
        assert!(table
            .grow_stack_for_fault(VirtAddr::new(0x6fff_b000))
            .is_err());
    }

    #[test]
    fn vm_region_rejects_invalid_stack_bounds_at_construction() {
        let permissions = VmPermissions::READ
            .union(VmPermissions::WRITE)
            .union(VmPermissions::USER);

        assert_eq!(
            VmRegion::new(
                VirtAddr::new(0x7000_0000),
                PAGE_SIZE,
                VmRegionKind::Stack { max_pages: 0 },
                permissions,
            ),
            Err("Stack maximum pages must be non-zero")
        );
        assert_eq!(
            VmRegion::new(
                VirtAddr::new(0x7000_0000),
                PAGE_SIZE * 2,
                VmRegionKind::Stack { max_pages: 1 },
                permissions,
            ),
            Err("Stack region length exceeds maximum pages")
        );
        assert_eq!(
            VmRegion::new(
                VirtAddr::new(0x7000_0000),
                PAGE_SIZE,
                VmRegionKind::Stack {
                    max_pages: VM_REGION_MAX_PAGES + 1,
                },
                permissions,
            ),
            Err("Stack maximum pages exceed teaching bounds")
        );
    }

    #[test]
    fn vm_region_stack_growth_preserves_overlap_invariant() {
        let permissions = VmPermissions::READ
            .union(VmPermissions::WRITE)
            .union(VmPermissions::USER);
        let mut table = VmRegionTable::new();
        table
            .insert(
                VmRegion::new(
                    VirtAddr::new(0x7fff_e000),
                    PAGE_SIZE,
                    VmRegionKind::Anonymous,
                    permissions,
                )
                .unwrap(),
            )
            .unwrap();
        table
            .insert(
                VmRegion::new(
                    VirtAddr::new(0x7fff_f000),
                    PAGE_SIZE,
                    VmRegionKind::Stack { max_pages: 4 },
                    permissions,
                )
                .unwrap(),
            )
            .unwrap();

        assert_eq!(
            table.grow_stack_for_fault(VirtAddr::new(0x7fff_e800)),
            Err("Stack growth would overlap an existing VM region")
        );
        let stack = table.find(VirtAddr::new(0x7fff_f000)).unwrap();
        assert_eq!(stack.start(), VirtAddr::new(0x7fff_f000));
        assert!(matches!(stack.kind(), VmRegionKind::Stack { max_pages: 4 }));
    }

    #[test]
    fn vm_region_stack_growth_keeps_committed_pages_at_original_addresses() {
        let permissions = VmPermissions::READ
            .union(VmPermissions::WRITE)
            .union(VmPermissions::USER);
        let mut table = VmRegionTable::new();
        let stack_start = VirtAddr::new(0x7000_0000);
        table
            .insert(
                VmRegion::new(
                    stack_start,
                    PAGE_SIZE,
                    VmRegionKind::Stack { max_pages: 4 },
                    permissions,
                )
                .unwrap(),
            )
            .unwrap();
        assert!(table.commit_page(stack_start).unwrap());
        assert!(table
            .grow_stack_for_fault(VirtAddr::new(0x6fff_f800))
            .unwrap());

        let grown = table.find(stack_start).unwrap();
        assert_eq!(grown.start(), VirtAddr::new(0x6fff_f000));
        assert!(grown.is_page_committed(stack_start));
        assert!(!grown.is_page_committed(VirtAddr::new(0x6fff_f000)));
    }
}