//
// Syd: rock-solid application kernel
// benches/sys/mmap.rs: mmap microbenchmarks
//
// Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
// Based in part upon gVisor's mapping_benchmark.cc which is:
//   Copyright 2020 The gVisor Authors.
//   SPDX-License-Identifier: Apache-2.0
//
// SPDX-License-Identifier: GPL-3.0

// A micro-benchmark that approximates the gVisor mmap micro-benchmarks.
// We replicate the following tests:
//   1) MapUnmap
//   2) MapTouchUnmap
//   3) MapTouchMany
//   4) PageFault
//
// Notes/Caveats:
//   - We map in pages, typically 4KB each on Linux.
//   - The original code uses Google Benchmark's Range(1, 1<<17) etc. That can
//     be large (~512 MB of memory). Adjust the arrays below if needed.
//   - The PageFault test in the original code tries to stress max_map_count
//     (~64k VMAs). We replicate the same logic in a single pass. This can be
//     quite large and may require sufficient memory/swap.

use std::{num::NonZeroUsize, ptr::NonNull, time::Duration};

use brunch::{benches, Bench};
use libc::c_void;
use nix::sys::mman::{madvise, mmap_anonymous, mprotect, munmap, MapFlags, MmapAdvise, ProtFlags};

// Typical page size on Linux.
const K_PAGE_SIZE: usize = 4096;
// A conservative limit on the number of VMAs, per the original snippet.
const K_MAX_VMAS: usize = 64001;

//------------------------------------------------------------------------------
// 1) BM_MapUnmap: Map then unmap `pages` pages without touching them.
//------------------------------------------------------------------------------
fn map_unmap(pages: usize) {
    let length = pages * K_PAGE_SIZE;
    let length_nonzero = match NonZeroUsize::new(length) {
        Some(v) => v,
        None => return, // If pages=0, just no-op.
    };

    // Map (anonymous).
    // The original snippet: MAP_PRIVATE | MAP_ANONYMOUS, PROT_READ|PROT_WRITE
    // We'll replicate that with nix's mmap_anonymous wrapper.
    let addr = unsafe {
        mmap_anonymous(
            None,
            length_nonzero,
            ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            MapFlags::MAP_PRIVATE,
        )
    }
    .expect("mmap failed in map_unmap");

    // Unmap.
    unsafe { munmap(addr, length) }.expect("munmap failed in map_unmap");
}

//------------------------------------------------------------------------------
// 2) BM_MapTouchUnmap: Map, touch each page, then unmap.
//------------------------------------------------------------------------------
fn map_touch_unmap(pages: usize) {
    let length = pages * K_PAGE_SIZE;
    let length_nonzero = match NonZeroUsize::new(length) {
        Some(v) => v,
        None => return,
    };

    let addr = unsafe {
        mmap_anonymous(
            None,
            length_nonzero,
            ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
            MapFlags::MAP_PRIVATE,
        )
    }
    .expect("mmap failed in map_touch_unmap");

    // Touch each page.
    unsafe {
        let mut cursor = addr.as_ptr() as *mut u8;
        let end = cursor.add(length);
        while cursor < end {
            // Write something to cause a demand-fault.
            *cursor = 42;
            cursor = cursor.add(K_PAGE_SIZE);
        }
    }

    // Unmap.
    unsafe { munmap(addr, length) }.expect("munmap failed in map_touch_unmap");
}

//------------------------------------------------------------------------------
// 3) BM_MapTouchMany: Map and touch many single-page mappings, unmapping all
//    at once. This replicates the loop-based approach of the original:
//    "Map each page, write to it, store the pointer, then unmap all in a batch."
//------------------------------------------------------------------------------
fn map_touch_many(page_count: usize) {
    // We'll store the individual addresses in a Vec.
    let mut pages_vec = Vec::with_capacity(page_count);

    // Map each page separately, PROT_READ|PROT_WRITE, MAP_PRIVATE|ANONYMOUS.
    for _ in 0..page_count {
        let addr = unsafe {
            mmap_anonymous(
                None,
                NonZeroUsize::new(K_PAGE_SIZE).unwrap(),
                ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
                MapFlags::MAP_PRIVATE,
            )
        }
        .expect("mmap failed in map_touch_many");

        // Touch the page.
        unsafe {
            *(addr.as_ptr() as *mut u8) = 42;
        }

        pages_vec.push(addr);
    }

    // Unmap all of them.
    for &addr in &pages_vec {
        unsafe { munmap(addr, K_PAGE_SIZE) }.expect("munmap failed in map_touch_many");
    }
}

//------------------------------------------------------------------------------
// 4) BM_PageFault: Large region mapped with PROT_READ, half its pages changed
//    to PROT_NONE, so we alternate reading and forcing page faults. We also
//    use MADV_DONTNEED to reset the region after we cross all mapped pages.
//------------------------------------------------------------------------------
fn page_fault_bench() {
    let mut test_pages = K_MAX_VMAS; // ~64k
                                     // Ensure it's odd: the snippet says "Ensure test_pages is odd..."
    if test_pages % 2 == 0 {
        test_pages -= 1;
    }
    let region_len = test_pages * K_PAGE_SIZE;
    let region_len_nonzero = match NonZeroUsize::new(region_len) {
        Some(v) => v,
        None => return,
    };

    // Map the region with PROT_READ, MAP_SHARED|MAP_POPULATE
    // The snippet uses MmapAnon in gVisor, which is effectively an
    // anonymous + shared. We'll replicate that with nix:
    let map_flags = MapFlags::MAP_SHARED | MapFlags::MAP_POPULATE | MapFlags::MAP_ANONYMOUS;
    let addr = unsafe {
        // Because nix's `mmap_anonymous` sets MAP_ANONYMOUS|MAP_PRIVATE by default,
        // we'll do a raw mmap to match MAP_SHARED.
        nix::sys::mman::mmap_anonymous(None, region_len_nonzero, ProtFlags::PROT_READ, map_flags)
    }
    .expect("mmap failed in page_fault_bench");

    // Convert to NonNull for further usage.
    let map_base = unsafe { NonNull::new_unchecked(addr.as_ptr() as *mut c_void) };

    // Mark every other page as PROT_NONE to force distinct VMAs.
    // i.e., pages: [mapped, none, mapped, none, ...]
    for i in 0..(test_pages / 2) {
        let page_addr = unsafe { map_base.as_ptr().add((2 * i + 1) * K_PAGE_SIZE) };
        let page_addr_nn = unsafe { NonNull::new_unchecked(page_addr as *mut c_void) };
        unsafe { mprotect(page_addr_nn, K_PAGE_SIZE, ProtFlags::PROT_NONE) }
            .expect("mprotect failed setting PROT_NONE");
    }

    // We'll do a quick read loop across these pages. Once we pass all "mapped" pages,
    // we do a MADV_DONTNEED to reset them and start over, simulating the repeated
    // page faults from the snippet.
    let mapped_pages = test_pages / 2 + 1;
    let mut cur_page = mapped_pages; // start beyond end -> triggers madvise first.

    // In a single call (since brunch calls us multiple times), we'll do ~2*mapped_pages
    // touches just to demonstrate the page fault cycle. Adjust if you want more/less.
    let iterations = 2 * mapped_pages;
    for _ in 0..iterations {
        if cur_page >= mapped_pages {
            // Reset the entire region, so we get faults again.
            unsafe { madvise(map_base, region_len, MmapAdvise::MADV_DONTNEED) }
                .expect("madvise(DONTNEED) failed");
            cur_page = 0;
        }
        let read_addr = unsafe { map_base.as_ptr().add((2 * cur_page) * K_PAGE_SIZE) };
        // Force a read fault. If it's PROT_NONE, that page won't be touched
        // in the snippet anyway. The snippet only allowed half the pages as READ,
        // but we'll read them in sequence. Some are NO-OP, some cause an actual read.
        let val = unsafe { *(read_addr as *const u8) };
        std::hint::black_box(val);
        cur_page += 1;
    }

    // Cleanup: unmap the entire region.
    unsafe { munmap(map_base, region_len) }.expect("munmap failed at page_fault_bench cleanup");
}

fn main() {
    benches!(
        inline:

        // 1) MapUnmap
        Bench::new("MapUnmap(1 pages)").run(|| {
            map_unmap(1);
        }),
        Bench::new("MapUnmap(256 pages)").run(|| {
            map_unmap(256);
        }),
        Bench::new("MapUnmap(65536 pages)").run(|| {
            map_unmap(65536);
        }),
        Bench::new("MapUnmap(131072 pages)").run(|| {
            map_unmap(131072);
        }),

        // 2) MapTouchUnmap
        Bench::new("MapTouchUnmap(1 pages)").run(|| {
            map_touch_unmap(1);
        }),
        Bench::new("MapTouchUnmap(256 pages)").run(|| {
            map_touch_unmap(256);
        }),
        Bench::new("MapTouchUnmap(65536 pages)").run(|| {
            map_touch_unmap(65536);
        }),
        Bench::new("MapTouchUnmap(131072 pages)").run(|| {
            map_touch_unmap(131072);
        }),

        // 3) MapTouchMany
        Bench::new("MapTouchMany(1 pages)").run(|| {
            map_touch_many(1);
        }),
        Bench::new("MapTouchMany(16 pages)").run(|| {
            map_touch_many(16);
        }),
        Bench::new("MapTouchMany(256 pages)").run(|| {
            map_touch_many(256);
        }),
        Bench::new("MapTouchMany(4096 pages)").run(|| {
            map_touch_many(4096);
        }),

        // 4) PageFault
        Bench::new("PageFault")
            .with_timeout(Duration::from_secs(30))
            .run(|| {
                page_fault_bench();
            }),
    );
}
