// SPDX-License-Identifier: GPL-2.0

//! Kernel page allocation and management.

use crate::{
    alloc::{AllocError, Flags},
    bindings,
    error::code::*,
    error::Result,
    uaccess::UserSliceReader,
};
use core::{
    marker::PhantomData,
    mem::ManuallyDrop,
    ops::Deref,
    ptr::{self, NonNull},
};

/// A bitwise shift for the page size.
pub const PAGE_SHIFT: usize = bindings::PAGE_SHIFT as usize;

/// The number of bytes in a page.
pub const PAGE_SIZE: usize = bindings::PAGE_SIZE;

/// A bitmask that gives the page containing a given address.
pub const PAGE_MASK: usize = !(PAGE_SIZE - 1);

/// Rounds up to the next multiple of [`PAGE_SIZE`].
///
/// Returns [`None`] on integer overflow.
///
/// # Examples
///
/// ```
/// use kernel::page::{
///     page_align,
///     PAGE_SIZE,
/// };
///
/// // Requested address is already aligned.
/// assert_eq!(page_align(0x0), Some(0x0));
/// assert_eq!(page_align(PAGE_SIZE), Some(PAGE_SIZE));
///
/// // Requested address needs alignment up.
/// assert_eq!(page_align(0x1), Some(PAGE_SIZE));
/// assert_eq!(page_align(PAGE_SIZE + 1), Some(2 * PAGE_SIZE));
///
/// // Requested address causes overflow (returns `None`).
/// let overflow_addr = usize::MAX - (PAGE_SIZE / 2);
/// assert_eq!(page_align(overflow_addr), None);
/// ```
#[inline(always)]
pub const fn page_align(addr: usize) -> Option<usize> {
    let Some(sum) = addr.checked_add(PAGE_SIZE - 1) else {
        return None;
    };
    Some(sum & PAGE_MASK)
}

/// Representation of a non-owning reference to a [`Page`].
///
/// This type provides a borrowed version of a [`Page`] that is owned by some other entity, e.g. a
/// [`Vmalloc`] allocation such as [`VBox`].
///
/// # Example
///
/// ```
/// # use kernel::{bindings, prelude::*};
/// use kernel::page::{BorrowedPage, Page, PAGE_SIZE};
/// # use core::{mem::MaybeUninit, ptr, ptr::NonNull };
///
/// fn borrow_page<'a>(vbox: &'a mut VBox<MaybeUninit<[u8; PAGE_SIZE]>>) -> BorrowedPage<'a> {
///     let ptr = ptr::from_ref(&**vbox);
///
///     // SAFETY: `ptr` is a valid pointer to `Vmalloc` memory.
///     let page = unsafe { bindings::vmalloc_to_page(ptr.cast()) };
///
///     // SAFETY: `vmalloc_to_page` returns a valid pointer to a `struct page` for a valid
///     // pointer to `Vmalloc` memory.
///     let page = unsafe { NonNull::new_unchecked(page) };
///
///     // SAFETY:
///     // - `self.0` is a valid pointer to a `struct page`.
///     // - `self.0` is valid for the entire lifetime of `self`.
///     unsafe { BorrowedPage::from_raw(page) }
/// }
///
/// let mut vbox = VBox::<[u8; PAGE_SIZE]>::new_uninit(GFP_KERNEL)?;
/// let page = borrow_page(&mut vbox);
///
/// // SAFETY: There is no concurrent read or write to this page.
/// unsafe { page.fill_zero_raw(0, PAGE_SIZE)? };
/// # Ok::<(), Error>(())
/// ```
///
/// # Invariants
///
/// The borrowed underlying pointer to a `struct page` is valid for the entire lifetime `'a`.
///
/// [`VBox`]: kernel::alloc::VBox
/// [`Vmalloc`]: kernel::alloc::allocator::Vmalloc
pub struct BorrowedPage<'a>(ManuallyDrop<Page>, PhantomData<&'a Page>);

impl<'a> BorrowedPage<'a> {
    /// Constructs a [`BorrowedPage`] from a raw pointer to a `struct page`.
    ///
    /// # Safety
    ///
    /// - `ptr` must point to a valid `bindings::page`.
    /// - `ptr` must remain valid for the entire lifetime `'a`.
    pub unsafe fn from_raw(ptr: NonNull<bindings::page>) -> Self {
        let page = Page { page: ptr };

        // INVARIANT: The safety requirements guarantee that `ptr` is valid for the entire lifetime
        // `'a`.
        Self(ManuallyDrop::new(page), PhantomData)
    }
}

impl<'a> Deref for BorrowedPage<'a> {
    type Target = Page;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

/// Trait to be implemented by types which provide an [`Iterator`] implementation of
/// [`BorrowedPage`] items, such as [`VmallocPageIter`](kernel::alloc::allocator::VmallocPageIter).
pub trait AsPageIter {
    /// The [`Iterator`] type, e.g. [`VmallocPageIter`](kernel::alloc::allocator::VmallocPageIter).
    type Iter<'a>: Iterator<Item = BorrowedPage<'a>>
    where
        Self: 'a;

    /// Returns an [`Iterator`] of [`BorrowedPage`] items over all pages owned by `self`.
    fn page_iter(&mut self) -> Self::Iter<'_>;
}

/// A pointer to a page that owns the page allocation.
///
/// # Invariants
///
/// The pointer is valid, and has ownership over the page.
pub struct Page {
    page: NonNull<bindings::page>,
}

// SAFETY: Pages have no logic that relies on them staying on a given thread, so moving them across
// threads is safe.
unsafe impl Send for Page {}

// SAFETY: Pages have no logic that relies on them not being accessed concurrently, so accessing
// them concurrently is safe.
unsafe impl Sync for Page {}

impl Page {
    /// Allocates a new page.
    ///
    /// # Examples
    ///
    /// Allocate memory for a page.
    ///
    /// ```
    /// use kernel::page::Page;
    ///
    /// let page = Page::alloc_page(GFP_KERNEL)?;
    /// # Ok::<(), kernel::alloc::AllocError>(())
    /// ```
    ///
    /// Allocate memory for a page and zero its contents.
    ///
    /// ```
    /// use kernel::page::Page;
    ///
    /// let page = Page::alloc_page(GFP_KERNEL | __GFP_ZERO)?;
    /// # Ok::<(), kernel::alloc::AllocError>(())
    /// ```
    #[inline]
    pub fn alloc_page(flags: Flags) -> Result<Self, AllocError> {
        // SAFETY: Depending on the value of `gfp_flags`, this call may sleep. Other than that, it
        // is always safe to call this method.
        let page = unsafe { bindings::alloc_pages(flags.as_raw(), 0) };
        let page = NonNull::new(page).ok_or(AllocError)?;
        // INVARIANT: We just successfully allocated a page, so we now have ownership of the newly
        // allocated page. We transfer that ownership to the new `Page` object.
        Ok(Self { page })
    }

    /// Returns a raw pointer to the page.
    pub fn as_ptr(&self) -> *mut bindings::page {
        self.page.as_ptr()
    }

    /// Get the node id containing this page.
    pub fn nid(&self) -> i32 {
        // SAFETY: Always safe to call with a valid page.
        unsafe { bindings::page_to_nid(self.as_ptr()) }
    }

    /// Runs a piece of code with this page mapped to an address.
    ///
    /// The page is unmapped when this call returns.
    ///
    /// # Using the raw pointer
    ///
    /// It is up to the caller to use the provided raw pointer correctly. The pointer is valid for
    /// `PAGE_SIZE` bytes and for the duration in which the closure is called. The pointer might
    /// only be mapped on the current thread, and when that is the case, dereferencing it on other
    /// threads is UB. Other than that, the usual rules for dereferencing a raw pointer apply: don't
    /// cause data races, the memory may be uninitialized, and so on.
    ///
    /// If multiple threads map the same page at the same time, then they may reference with
    /// different addresses. However, even if the addresses are different, the underlying memory is
    /// still the same for these purposes (e.g., it's still a data race if they both write to the
    /// same underlying byte at the same time).
    fn with_page_mapped<T>(&self, f: impl FnOnce(*mut u8) -> T) -> T {
        // SAFETY: `page` is valid due to the type invariants on `Page`.
        let mapped_addr = unsafe { bindings::kmap_local_page(self.as_ptr()) };

        let res = f(mapped_addr.cast());

        // This unmaps the page mapped above.
        //
        // SAFETY: Since this API takes the user code as a closure, it can only be used in a manner
        // where the pages are unmapped in reverse order. This is as required by `kunmap_local`.
        //
        // In other words, if this call to `kunmap_local` happens when a different page should be
        // unmapped first, then there must necessarily be a call to `kmap_local_page` other than the
        // call just above in `with_page_mapped` that made that possible. In this case, it is the
        // unsafe block that wraps that other call that is incorrect.
        unsafe { bindings::kunmap_local(mapped_addr) };

        res
    }

    /// Runs a piece of code with a raw pointer to a slice of this page, with bounds checking.
    ///
    /// If `f` is called, then it will be called with a pointer that points at `off` bytes into the
    /// page, and the pointer will be valid for at least `len` bytes. The pointer is only valid on
    /// this task, as this method uses a local mapping.
    ///
    /// If `off` and `len` refers to a region outside of this page, then this method returns
    /// [`EINVAL`] and does not call `f`.
    ///
    /// # Using the raw pointer
    ///
    /// It is up to the caller to use the provided raw pointer correctly. The pointer is valid for
    /// `len` bytes and for the duration in which the closure is called. The pointer might only be
    /// mapped on the current thread, and when that is the case, dereferencing it on other threads
    /// is UB. Other than that, the usual rules for dereferencing a raw pointer apply: don't cause
    /// data races, the memory may be uninitialized, and so on.
    ///
    /// If multiple threads map the same page at the same time, then they may reference with
    /// different addresses. However, even if the addresses are different, the underlying memory is
    /// still the same for these purposes (e.g., it's still a data race if they both write to the
    /// same underlying byte at the same time).
    fn with_pointer_into_page<T>(
        &self,
        off: usize,
        len: usize,
        f: impl FnOnce(*mut u8) -> Result<T>,
    ) -> Result<T> {
        let bounds_ok = off <= PAGE_SIZE && len <= PAGE_SIZE && (off + len) <= PAGE_SIZE;

        if bounds_ok {
            self.with_page_mapped(move |page_addr| {
                // SAFETY: The `off` integer is at most `PAGE_SIZE`, so this pointer offset will
                // result in a pointer that is in bounds or one off the end of the page.
                f(unsafe { page_addr.add(off) })
            })
        } else {
            Err(EINVAL)
        }
    }

    /// Maps the page and reads from it into the given buffer.
    ///
    /// This method will perform bounds checks on the page offset. If `offset .. offset+len` goes
    /// outside of the page, then this call returns [`EINVAL`].
    ///
    /// # Safety
    ///
    /// * Callers must ensure that `dst` is valid for writing `len` bytes.
    /// * Callers must ensure that this call does not race with a write to the same page that
    ///   overlaps with this read.
    pub unsafe fn read_raw(&self, dst: *mut u8, offset: usize, len: usize) -> Result {
        self.with_pointer_into_page(offset, len, move |src| {
            // SAFETY: If `with_pointer_into_page` calls into this closure, then
            // it has performed a bounds check and guarantees that `src` is
            // valid for `len` bytes.
            //
            // There caller guarantees that there is no data race.
            unsafe { ptr::copy_nonoverlapping(src, dst, len) };
            Ok(())
        })
    }

    /// Maps the page and writes into it from the given buffer.
    ///
    /// This method will perform bounds checks on the page offset. If `offset .. offset+len` goes
    /// outside of the page, then this call returns [`EINVAL`].
    ///
    /// # Safety
    ///
    /// * Callers must ensure that `src` is valid for reading `len` bytes.
    /// * Callers must ensure that this call does not race with a read or write to the same page
    ///   that overlaps with this write.
    pub unsafe fn write_raw(&self, src: *const u8, offset: usize, len: usize) -> Result {
        self.with_pointer_into_page(offset, len, move |dst| {
            // SAFETY: If `with_pointer_into_page` calls into this closure, then it has performed a
            // bounds check and guarantees that `dst` is valid for `len` bytes.
            //
            // There caller guarantees that there is no data race.
            unsafe { ptr::copy_nonoverlapping(src, dst, len) };
            Ok(())
        })
    }

    /// Maps the page and zeroes the given slice.
    ///
    /// This method will perform bounds checks on the page offset. If `offset .. offset+len` goes
    /// outside of the page, then this call returns [`EINVAL`].
    ///
    /// # Safety
    ///
    /// Callers must ensure that this call does not race with a read or write to the same page that
    /// overlaps with this write.
    pub unsafe fn fill_zero_raw(&self, offset: usize, len: usize) -> Result {
        self.with_pointer_into_page(offset, len, move |dst| {
            // SAFETY: If `with_pointer_into_page` calls into this closure, then it has performed a
            // bounds check and guarantees that `dst` is valid for `len` bytes.
            //
            // There caller guarantees that there is no data race.
            unsafe { ptr::write_bytes(dst, 0u8, len) };
            Ok(())
        })
    }

    /// Copies data from userspace into this page.
    ///
    /// This method will perform bounds checks on the page offset. If `offset .. offset+len` goes
    /// outside of the page, then this call returns [`EINVAL`].
    ///
    /// Like the other `UserSliceReader` methods, data races are allowed on the userspace address.
    /// However, they are not allowed on the page you are copying into.
    ///
    /// # Safety
    ///
    /// Callers must ensure that this call does not race with a read or write to the same page that
    /// overlaps with this write.
    pub unsafe fn copy_from_user_slice_raw(
        &self,
        reader: &mut UserSliceReader,
        offset: usize,
        len: usize,
    ) -> Result {
        self.with_pointer_into_page(offset, len, move |dst| {
            // SAFETY: If `with_pointer_into_page` calls into this closure, then it has performed a
            // bounds check and guarantees that `dst` is valid for `len` bytes. Furthermore, we have
            // exclusive access to the slice since the caller guarantees that there are no races.
            reader.read_raw(unsafe { core::slice::from_raw_parts_mut(dst.cast(), len) })
        })
    }
}

impl Drop for Page {
    #[inline]
    fn drop(&mut self) {
        // SAFETY: By the type invariants, we have ownership of the page and can free it.
        unsafe { bindings::__free_pages(self.page.as_ptr(), 0) };
    }
}