//! Large output buffers backed by transparent huge pages where the OS offers
//! them.
//!
//! A fresh multi-megabyte `Vec` is mapped lazily by the kernel: the first
//! write to each 4 KiB page takes a page fault, and the kernel zeroes the page
//! before handing it over. For a 64 MiB read that is 16384 faults, and they
//! cost far more than the copy that fills the buffer — single-threaded
//! contiguous reads ran at about a quarter of h5py's speed because of them.
//! numpy (so h5py) avoids this by asking for transparent huge pages
//! (`madvise(MADV_HUGEPAGE)`) on every allocation of 4 MiB or more, which
//! turns 512 faults into one; this module does the same.
//!
//! The advice only changes how the pages are backed, never their contents, so
//! it is harmless when it cannot be honoured (THP disabled, not Linux, a
//! region that is part of the heap): the buffer is then exactly what it would
//! have been without it.

#[cfg(not(feature = "std"))]
use alloc::vec::Vec;

/// Buffers smaller than this are left alone (numpy uses the same threshold).
#[cfg(any(target_os = "linux", test))]
pub(crate) const HUGE_PAGE_THRESHOLD: usize = 4 << 20;

/// Advise the kernel to back `[ptr, ptr + len)` with transparent huge pages,
/// when `len` is large enough to benefit. Call it before the first write so
/// the faults happen at huge-page granularity.
#[inline]
pub(crate) fn advise_huge_pages(ptr: *const u8, len: usize) {
    #[cfg(target_os = "linux")]
    if len >= HUGE_PAGE_THRESHOLD {
        const PAGE: usize = 4096;
        let start = (ptr as usize).next_multiple_of(PAGE);
        let end = (ptr as usize + len) & !(PAGE - 1);
        if end > start {
            // SAFETY: `[start, end)` lies inside an allocation of `len` bytes
            // at `ptr` that the caller owns, and is page aligned as madvise
            // requires. MADV_HUGEPAGE does not change the memory's contents or
            // validity; on failure (EINVAL when THP is compiled out, etc.) the
            // region is simply left as it was, so the result is ignored.
            unsafe {
                libc::madvise(start as *mut libc::c_void, end - start, libc::MADV_HUGEPAGE);
            }
        }
    }
    #[cfg(not(target_os = "linux"))]
    let _ = (ptr, len);
}

/// `Vec::with_capacity(count)` for a buffer about to be filled in bulk, with
/// huge-page advice when it is large (see the module docs).
#[inline]
pub(crate) fn vec_for_bulk<T>(count: usize) -> Vec<T> {
    let v: Vec<T> = Vec::with_capacity(count);
    advise_huge_pages(
        v.as_ptr().cast::<u8>(),
        v.capacity().saturating_mul(core::mem::size_of::<T>()),
    );
    v
}

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

    #[test]
    fn bulk_vec_is_an_ordinary_vec() {
        for count in [0usize, 1, 1000, HUGE_PAGE_THRESHOLD / 4 + 3] {
            let mut v: Vec<u32> = vec_for_bulk(count);
            assert!(v.capacity() >= count);
            v.extend((0..count as u32).map(|i| i.wrapping_mul(2654435761)));
            assert!(
                v.iter()
                    .enumerate()
                    .all(|(i, &x)| x == (i as u32).wrapping_mul(2654435761))
            );
        }
    }
}