//! A process's memory map: which address ranges exist and what they are.
//!
//! Stage 13 mapped every page of a program up front. Now an address space
//! is described by *regions* (Linux: VMAs, "virtual memory areas"), and
//! pages appear only when touched:
//!
//!   program   the ELF segments (code, rodata, data+bss), loaded eagerly
//!   heap      from the end of the program up to `brk` (grows with brk())
//!   mmap      what mmap() created: anonymous memory, or a file's bytes
//!   stack     1 MiB under USER_STACK_TOP, with a 64 KiB guard zone below
//!
//! **Demand paging.** A region is just a promise. The first access to one
//! of its pages finds no page-table entry, the CPU raises a page fault, and
//! `handle_fault` sees the address is inside a region: it allocates a
//! zeroed frame (or reads the file's bytes into it), maps it, and the
//! instruction runs again -- the program never notices. An access outside
//! every region (or one the region forbids, like a write to read-only
//! memory) is a real error: the process gets SIGSEGV. So mmap(64 MiB)
//! costs nothing until the program uses it.
//!
//! **Copy-on-write.** fork() doesn't copy memory: the child gets new page
//! tables pointing at the *same* frames, and every private writable page is
//! made read-only in both. Reads just work. The first write faults;
//! `handle_fault` sees "write to a writable region, page is read-only":
//! if somebody else still maps the frame (reference count > 1) it copies
//! the page into a fresh frame for the writer, otherwise the writer is the
//! last owner and simply gets write access back. Pages that are never
//! written (all the code, usually) stay shared forever.
//!
//! **Shared mappings** (MAP_SHARED | MAP_ANONYMOUS) stay shared across fork:
//! no copy-on-write, both processes see each other's writes. We allocate
//! their frames right away at mmap time (a page first touched *after* a
//! fork would otherwise get two different frames, one per process).
//!
//! **Kernel access to user memory.** Syscalls touch user buffers directly
//! (e.g. `read` copies into one). Before that, `user_access_ok` runs the
//! same fault handling for every page of the buffer, so the pages exist
//! and are writable (copies made) before the kernel uses them: the kernel
//! itself never page-faults on user memory.
//!
//! Locking: `Process::mm` is a plain spinlock, taken only by the process's
//! own thread (syscalls, its page faults, fork, exec) -- never by interrupt
//! handlers, and with interrupts on.

use crate::errno::Errno;
use crate::fs::Inode;
use crate::memory::{self, USER_END, USER_START};
use alloc::{sync::Arc, vec::Vec};
use x86_64::structures::paging::{PageTableFlags as Flags, PhysFrame};

pub const PAGE_SIZE: u64 = 4096;

/// The user stack: 1 MiB (pages appear as it grows), ending one unmapped
/// page below USER_END.
pub const USER_STACK_TOP: u64 = USER_END - 0x1000;
pub const USER_STACK_SIZE: u64 = 1024 * 1024;
/// Below the stack: a 64 KiB guard zone that is never mapped, so a runaway
/// recursion hits it and gets SIGSEGV ("stack overflow") instead of
/// silently growing into other memory. Big, so that one large stack frame
/// can't jump over it.
pub const USER_STACK_GUARD: u64 = 64 * 1024;
/// Lowest address of stack + guard: program segments must end below this.
pub const USER_STACK_ZONE_START: u64 = USER_STACK_TOP - USER_STACK_SIZE - USER_STACK_GUARD;
/// mmap hands out addresses downwards from here (the heap grows up towards it).
pub const MMAP_TOP: u64 = USER_STACK_ZONE_START;

/// Is user address `addr` in the guard zone under the user stack?
pub fn in_user_stack_guard(addr: u64) -> bool {
    (USER_STACK_ZONE_START..USER_STACK_TOP - USER_STACK_SIZE).contains(&addr)
}

// mmap/mprotect protection bits (Linux's values).
pub const PROT_NONE: u64 = 0;
pub const PROT_READ: u64 = 1;
pub const PROT_WRITE: u64 = 2;
pub const PROT_EXEC: u64 = 4;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Kind {
    Program,
    Heap,
    Mmap,
    Stack,
}

/// One region: `start..end` (page aligned) with the same permissions and
/// the same kind of contents.
#[derive(Clone)]
pub struct Vma {
    pub start: u64,
    pub end: u64,
    pub prot: u64,
    /// MAP_SHARED: writes are seen by every process mapping it.
    pub shared: bool,
    /// File-backed: (inode, file offset of `start`). Pages are read from
    /// the file on first touch; bytes past the end of the file are zero.
    pub file: Option<(Arc<dyn Inode>, u64)>,
    pub kind: Kind,
}

impl Vma {
    pub fn new(start: u64, end: u64, prot: u64, kind: Kind) -> Vma {
        Vma { start, end, prot, shared: false, file: None, kind }
    }

    /// May a program access this region this way? (x86 can't express
    /// write-only, so any access needs at least one of the bits.)
    fn allows(&self, write: bool, exec: bool) -> bool {
        self.prot != PROT_NONE && (!write || self.prot & PROT_WRITE != 0) && (!exec || self.prot & PROT_EXEC != 0)
    }

    /// Page-table flags for one of its pages. `writable`: may this page be
    /// writable right now? (Not if it's a private page still shared with
    /// another process: copy-on-write.) PROT_NONE pages stay present but
    /// lose USER_ACCESSIBLE, so the frame isn't lost and any user access
    /// faults.
    fn pte_flags(&self, writable: bool) -> Flags {
        let mut flags = Flags::PRESENT;
        if self.prot != PROT_NONE {
            flags |= Flags::USER_ACCESSIBLE;
        }
        if writable && self.prot & PROT_WRITE != 0 {
            flags |= Flags::WRITABLE;
        }
        if self.prot & PROT_EXEC == 0 && memory::nx_enabled() {
            flags |= Flags::NO_EXECUTE;
        }
        flags
    }

    /// Cut off and return the part from `at` on (`at` inside the region).
    fn split_off(&mut self, at: u64) -> Vma {
        let mut tail = self.clone();
        tail.start = at;
        if let Some((_, offset)) = &mut tail.file {
            *offset += at - self.start;
        }
        self.end = at;
        tail
    }
}

/// Why a page fault couldn't be fixed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Fault {
    /// No region there: SIGSEGV (SEGV_MAPERR).
    Unmapped,
    /// The region doesn't allow this access: SIGSEGV (SEGV_ACCERR).
    Protection,
    /// No free frame left (Linux would wake the OOM killer).
    OutOfMemory,
    /// Reading the file failed: SIGBUS.
    Io,
}

/// A process's address space: its page tables and its regions.
pub struct Mm {
    /// Level-4 page table. Never changes (execve empties and reuses it).
    pub l4: PhysFrame,
    /// Regions, sorted by address, never overlapping.
    vmas: Vec<Vma>,
    /// The heap: brk_start..brk (brk_start is page aligned).
    pub brk_start: u64,
    pub brk: u64,
}

fn page_up(x: u64) -> u64 {
    x.next_multiple_of(PAGE_SIZE)
}

impl Mm {
    /// A new, empty address space.
    pub fn new() -> Option<Mm> {
        Some(Mm { l4: memory::new_address_space()?, vmas: Vec::new(), brk_start: 0, brk: 0 })
    }

    pub fn vmas(&self) -> &[Vma] {
        &self.vmas
    }

    /// The region containing `addr`.
    fn find(&self, addr: u64) -> Option<&Vma> {
        self.vmas.iter().find(|v| (v.start..v.end).contains(&addr))
    }

    /// Does any region overlap `start..end`?
    fn overlaps(&self, start: u64, end: u64) -> bool {
        self.vmas.iter().any(|v| v.start < end && start < v.end)
    }

    /// Add a region (the range must be free).
    pub fn insert(&mut self, vma: Vma) {
        debug_assert!(!self.overlaps(vma.start, vma.end));
        let i = self.vmas.iter().position(|v| v.start > vma.start).unwrap_or(self.vmas.len());
        self.vmas.insert(i, vma);
    }

    /// Split regions so that none crosses `start` or `end`.
    fn split_at(&mut self, start: u64, end: u64) {
        let mut i = 0;
        while i < self.vmas.len() {
            for at in [start, end] {
                let v = &mut self.vmas[i];
                if v.start < at && at < v.end {
                    let tail = v.split_off(at);
                    self.vmas.insert(i + 1, tail);
                }
            }
            i += 1;
        }
    }

    /// Remove everything in `start..end` (page aligned): regions and pages.
    pub fn unmap(&mut self, start: u64, end: u64) {
        self.split_at(start, end);
        self.vmas.retain(|v| v.end <= start || v.start >= end);
        for addr in (start..end).step_by(PAGE_SIZE as usize) {
            if let Some(frame) = memory::unmap_user_page(self.l4, addr) {
                memory::release_frame(frame);
                memory::flush_user_page(addr);
            }
        }
    }

    /// mprotect: new permissions for `start..end`, which must be fully
    /// covered by regions (else ENOMEM, like Linux).
    pub fn protect(&mut self, start: u64, end: u64, prot: u64) -> Result<(), Errno> {
        let mut covered = start;
        for v in self.vmas.iter().filter(|v| v.end > start && v.start < end) {
            if v.start > covered {
                return Err(Errno::ENOMEM); // a hole
            }
            covered = v.end;
        }
        if covered < end {
            return Err(Errno::ENOMEM);
        }
        self.split_at(start, end);
        for v in self.vmas.iter_mut().filter(|v| v.start >= start && v.end <= end) {
            v.prot = prot;
            // Pages already there get the new flags. Private pages start out
            // read-only: the next write (if allowed) makes them writable
            // again, copying first if the frame is shared.
            for addr in (v.start..v.end).step_by(PAGE_SIZE as usize) {
                if memory::user_pte(self.l4, addr).is_some() {
                    memory::set_user_pte_flags(self.l4, addr, v.pte_flags(v.shared));
                    memory::flush_user_page(addr);
                }
            }
        }
        Ok(())
    }

    /// A free range of `len` bytes for mmap: the highest gap below MMAP_TOP
    /// (and above the heap).
    pub fn find_free(&self, len: u64) -> Option<u64> {
        let floor = page_up(self.brk).max(USER_START);
        let mut top = MMAP_TOP;
        for v in self.vmas.iter().rev().filter(|v| v.start < MMAP_TOP) {
            if v.end <= top && top - v.end >= len {
                return Some(top - len); // the gap above this region fits
            }
            top = top.min(v.start);
        }
        (top >= floor + len).then(|| top - len)
    }

    /// Is `start..end` inside user space and free (for MAP_FIXED / brk)?
    pub fn range_ok(start: u64, end: u64) -> bool {
        start >= USER_START && end <= MMAP_TOP && start < end
    }

    /// Fill the pages of `start..end` right away (shared anonymous memory).
    pub fn populate(&mut self, start: u64, end: u64) -> Result<(), Fault> {
        for addr in (start..end).step_by(PAGE_SIZE as usize) {
            self.handle_fault(addr, false, false)?;
        }
        Ok(())
    }

    /// brk(addr): move the end of the heap. Returns the new end (the old
    /// one if the request can't be met -- that's how Linux reports errors).
    pub fn set_brk(&mut self, new: u64) -> u64 {
        if new < self.brk_start {
            return self.brk; // brk(0) asks where the heap ends
        }
        let (old_end, new_end) = (page_up(self.brk), page_up(new));
        if new_end > old_end {
            if !Mm::range_ok(old_end, new_end) || self.overlaps(old_end, new_end) {
                return self.brk; // would run into a mapping
            }
            match self.vmas.iter_mut().find(|v| v.kind == Kind::Heap && v.end == old_end) {
                Some(heap) => heap.end = new_end,
                None => self.insert(Vma::new(old_end, new_end, PROT_READ | PROT_WRITE, Kind::Heap)),
            }
        } else if new_end < old_end {
            self.unmap(new_end, old_end);
        }
        self.brk = new;
        new
    }

    /// The page fault handler's real work (also used before the kernel
    /// touches user memory): make the page at `addr` present and allow the
    /// access, if its region allows it.
    pub fn handle_fault(&mut self, addr: u64, write: bool, exec: bool) -> Result<(), Fault> {
        let vma = self.find(addr).ok_or(Fault::Unmapped)?;
        if !vma.allows(write, exec) {
            return Err(Fault::Protection);
        }
        let page = addr & !(PAGE_SIZE - 1);
        match memory::user_pte(self.l4, page) {
            // Not there yet: demand paging.
            None => {
                let frame = memory::alloc_zeroed_frame().ok_or(Fault::OutOfMemory)?;
                if let Some((inode, offset)) = &vma.file {
                    let dst = memory::phys_to_virt(frame.start_address()).as_mut_ptr::<u8>();
                    let buf = unsafe { core::slice::from_raw_parts_mut(dst, PAGE_SIZE as usize) };
                    // A short read (end of file) leaves the rest zero.
                    if inode.read_at((offset + page - vma.start) as usize, buf).is_err() {
                        memory::free_frame(frame);
                        return Err(Fault::Io);
                    }
                }
                let flags = vma.pte_flags(true);
                if memory::map_user_frame(self.l4, page, frame, flags).is_none() {
                    memory::free_frame(frame);
                    return Err(Fault::OutOfMemory);
                }
            }
            // There, but read-only, and this is a write: copy-on-write.
            Some((frame, flags)) if write && !flags.contains(Flags::WRITABLE) => {
                if vma.shared || memory::frame_refcount(frame) == 1 {
                    // We're the only owner (the other one copied or exited):
                    // no copy needed, just allow writing again.
                    memory::set_user_pte_flags(self.l4, page, vma.pte_flags(true));
                } else {
                    let copy = memory::alloc_frame().ok_or(Fault::OutOfMemory)?;
                    unsafe {
                        core::ptr::copy_nonoverlapping(
                            memory::phys_to_virt(frame.start_address()).as_ptr::<u8>(),
                            memory::phys_to_virt(copy.start_address()).as_mut_ptr::<u8>(),
                            PAGE_SIZE as usize,
                        );
                    }
                    let flags = vma.pte_flags(true);
                    memory::replace_user_frame(self.l4, page, copy, flags).ok_or(Fault::OutOfMemory)?;
                    memory::release_frame(frame); // our share of the original
                }
            }
            // There and allowed by the region, but the PTE says no: after an
            // mprotect that widened the permissions. Refresh the flags.
            Some((frame, _)) => {
                let writable = vma.shared || memory::frame_refcount(frame) == 1;
                memory::set_user_pte_flags(self.l4, page, vma.pte_flags(writable && write));
            }
        }
        memory::flush_user_page(page);
        Ok(())
    }

    /// Make sure the program could access `addr..addr+len` (faulting pages
    /// in as needed), so the kernel can safely use the memory directly.
    pub fn prepare(&mut self, addr: u64, len: u64, write: bool) -> bool {
        let Some(end) = addr.checked_add(len) else { return false };
        if addr < USER_START || end > USER_END {
            return false;
        }
        let mut page = addr & !(PAGE_SIZE - 1);
        while page < end {
            let ok = match memory::user_pte(self.l4, page) {
                Some((_, f)) => f.contains(Flags::USER_ACCESSIBLE) && (!write || f.contains(Flags::WRITABLE)),
                None => false,
            };
            if !ok && self.handle_fault(page, write, false).is_err() {
                return false;
            }
            page += PAGE_SIZE;
        }
        true
    }

    /// Copy `data` to user address `addr` of this address space (which
    /// need not be the active one: we write through the physical-memory
    /// window). Used to build a new program's stack.
    pub fn write_bytes(&mut self, mut addr: u64, mut data: &[u8]) -> Result<(), Fault> {
        while !data.is_empty() {
            self.handle_fault(addr, true, false)?;
            let (frame, _) = memory::user_pte(self.l4, addr).ok_or(Fault::Unmapped)?;
            let offset = addr % PAGE_SIZE;
            let n = data.len().min((PAGE_SIZE - offset) as usize);
            let dst = memory::phys_to_virt(frame.start_address()) + offset;
            unsafe { core::ptr::copy_nonoverlapping(data.as_ptr(), dst.as_mut_ptr::<u8>(), n) };
            addr += n as u64;
            data = &data[n..];
        }
        Ok(())
    }

    /// fork: a copy of this address space sharing every frame. Private
    /// writable pages become read-only in *both* (copy-on-write); shared
    /// ones stay as they are. Our own TLB must forget the old writable
    /// entries, so the caller (the running process) flushes it.
    pub fn fork(&self) -> Option<Mm> {
        let mut child = Mm::new()?;
        child.vmas = self.vmas.clone();
        child.brk_start = self.brk_start;
        child.brk = self.brk;
        let vmas = &self.vmas;
        let ok = memory::clone_user_space(self.l4, child.l4, |addr, flags| {
            let shared = vmas.iter().any(|v| v.shared && (v.start..v.end).contains(&addr));
            if shared { flags } else { flags - Flags::WRITABLE }
        });
        flush_all();
        ok.map(|_| child) // on failure, dropping `child` frees what it got
    }

    /// execve: throw the whole user part away (regions, pages, tables).
    pub fn clear(&mut self) {
        memory::clear_user_space(self.l4);
        self.vmas.clear();
        self.brk_start = 0;
        self.brk = 0;
        flush_all();
    }

    /// Pages mapped right now (resident) and the total size of all regions.
    pub fn usage(&self) -> (u64, u64) {
        let mut resident = 0;
        let mut total = 0;
        for v in &self.vmas {
            total += v.end - v.start;
            for addr in (v.start..v.end).step_by(PAGE_SIZE as usize) {
                resident += memory::user_pte(self.l4, addr).is_some() as u64;
            }
        }
        (resident, total / PAGE_SIZE)
    }
}

impl Drop for Mm {
    /// The process is gone (its thread finished, nobody runs on these
    /// page tables): give every frame back.
    fn drop(&mut self) {
        unsafe { memory::free_address_space(self.l4) };
    }
}

/// Forget all of this CPU's user translations (reload CR3).
fn flush_all() {
    use x86_64::registers::control::Cr3;
    let (frame, flags) = Cr3::read();
    unsafe { Cr3::write(frame, flags) };
}

/// Syscalls: may the *current* process access `len` bytes at `addr`
/// (faulting the pages in, breaking copy-on-write for writes)? Never trust
/// a pointer from user space: without this check a program could make the
/// kernel read or overwrite kernel memory for it.
pub fn user_access_ok(addr: u64, len: u64, write: bool) -> bool {
    match crate::thread::current_process() {
        Some(p) => p.mm.lock().prepare(addr, len, write),
        None => false, // kernel threads have no user memory
    }
}