//! Scope-tagged counting allocator (`SMH-SPEC-SPEC0001`, Memory accounting).
//!
//! [`Counting`] wraps [`System`] and attributes every allocation, including
//! third-party ones, to the innermost [`Scope`] entered on the allocating
//! thread. Each block carries its scope tag in a header, so a free is charged
//! to the scope that allocated it and live bytes stay exact per scope.
//! Scopes do not propagate to spawned threads; those start at
//! [`Scope::Process`].

use std::alloc::{GlobalAlloc, Layout, System};
use std::cell::Cell;
use std::marker::PhantomData;
use std::sync::atomic::{AtomicU64, Ordering::Relaxed};

/// Logical lifetime an allocation is charged to.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum Scope {
    /// Everything outside an explicit scope, including pre-main.
    Process,
    /// One session.
    Session,
    /// One agent turn.
    Turn,
    /// One kernel call.
    KernelCall,
    /// One rendered frame.
    Frame,
    /// One I/O request.
    IoRequest,
}

const SCOPES: [Scope; 6] = [
    Scope::Process,
    Scope::Session,
    Scope::Turn,
    Scope::KernelCall,
    Scope::Frame,
    Scope::IoRequest,
];

/// Counters recorded for one scope since process start.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct Stats {
    /// Allocations, reallocations included.
    pub allocations: u64,
    /// Bytes requested by those allocations.
    pub bytes: u64,
    /// Bytes allocated in this scope and not yet freed.
    pub live: u64,
    /// Highest `live` observed.
    pub peak: u64,
}

struct Counters {
    allocations: AtomicU64,
    bytes: AtomicU64,
    live: AtomicU64,
    peak: AtomicU64,
}

static COUNTERS: [Counters; SCOPES.len()] = [const {
    Counters {
        allocations: AtomicU64::new(0),
        bytes: AtomicU64::new(0),
        live: AtomicU64::new(0),
        peak: AtomicU64::new(0),
    }
}; SCOPES.len()];

const DEPTH: usize = 16;

#[derive(Clone, Copy)]
struct Stack {
    len: usize,
    items: [Scope; DEPTH],
}

thread_local! {
    // Const-initialized and drop-free: access never allocates or registers a destructor.
    static STACK: Cell<Stack> = const { Cell::new(Stack { len: 0, items: [Scope::Process; DEPTH] }) };
}

fn current() -> Scope {
    STACK
        .try_with(|stack| {
            let stack = stack.get();
            stack.items.get(stack.len.wrapping_sub(1)).copied()
        })
        .ok()
        .flatten()
        .unwrap_or(Scope::Process)
}

/// Restores the enclosing scope when dropped; bound to its thread.
#[must_use = "the scope ends when the guard drops"]
pub struct Guard {
    previous: usize,
    _thread_bound: PhantomData<*const ()>,
}

impl Drop for Guard {
    fn drop(&mut self) {
        let _ = STACK.try_with(|cell| {
            let mut stack = cell.get();
            // `min` tolerates non-LIFO drops, such as a `Vec` of guards.
            stack.len = stack.len.min(self.previous);
            cell.set(stack);
        });
    }
}

/// Charge this thread's allocations to `scope` until the guard drops.
///
/// Nesting deeper than 16 keeps charging the outer scope (soft limit).
pub fn enter(scope: Scope) -> Guard {
    let previous = STACK
        .try_with(|cell| {
            let mut stack = cell.get();
            let previous = stack.len;
            if let Some(slot) = stack.items.get_mut(previous) {
                *slot = scope;
                stack.len += 1;
                cell.set(stack);
            }
            previous
        })
        .unwrap_or(0);
    Guard {
        previous,
        _thread_bound: PhantomData,
    }
}

/// Counters recorded for `scope`.
pub fn snapshot(scope: Scope) -> Stats {
    let counters = &COUNTERS[scope as usize];
    Stats {
        allocations: counters.allocations.load(Relaxed),
        bytes: counters.bytes.load(Relaxed),
        live: counters.live.load(Relaxed),
        peak: counters.peak.load(Relaxed),
    }
}

/// Run `work` charged to `scope` and report what it cost: `allocations`
/// and `bytes` are the deltas across the run, `live` and `peak` the scope's
/// values right after it.
///
/// Counters are process-wide; concurrent measurements of one scope mix.
pub fn measure<T>(scope: Scope, work: impl FnOnce() -> T) -> (T, Stats) {
    let before = snapshot(scope);
    let guard = enter(scope);
    let out = work();
    drop(guard);
    let after = snapshot(scope);
    (
        out,
        Stats {
            allocations: after.allocations - before.allocations,
            bytes: after.bytes - before.bytes,
            live: after.live,
            peak: after.peak,
        },
    )
}

/// Every scope with its counters.
pub fn snapshot_all() -> [(Scope, Stats); SCOPES.len()] {
    SCOPES.map(|scope| (scope, snapshot(scope)))
}

fn record_alloc(scope: Scope, size: usize) {
    let counters = &COUNTERS[scope as usize];
    counters.allocations.fetch_add(1, Relaxed);
    counters.bytes.fetch_add(size as u64, Relaxed);
    let now = counters.live.fetch_add(size as u64, Relaxed) + size as u64;
    counters.peak.fetch_max(now, Relaxed);
}

fn record_free(tag: u8, size: usize) {
    // A tag outside the scope range means a foreign block; nothing to charge.
    if let Some(counters) = COUNTERS.get(usize::from(tag)) {
        counters.live.fetch_sub(size as u64, Relaxed);
    }
}

/// Header of `align` bytes in front of each block; its last byte is the scope tag.
fn padded(layout: Layout) -> Option<(Layout, usize)> {
    let pad = layout.align();
    let size = layout.size().checked_add(pad)?;
    Layout::from_size_align(size, pad)
        .ok()
        .map(|outer| (outer, pad))
}

/// Counting wrapper around [`System`]; install with `#[global_allocator]`.
pub struct Counting;

impl Counting {
    /// # Safety
    /// `raw` must allocate a block for the layout it is given, like [`GlobalAlloc::alloc`].
    unsafe fn tagged(layout: Layout, raw: impl FnOnce(Layout) -> *mut u8) -> *mut u8 {
        let Some((outer, pad)) = padded(layout) else {
            return std::ptr::null_mut();
        };
        // SAFETY: `raw` returned null or a block of `outer`.
        unsafe { Self::stamp(raw(outer), pad, layout.size()) }
    }

    /// Tag a fresh `System` block with the current scope and charge `size`
    /// to it; null passes through.
    ///
    /// # Safety
    /// `base` is null or a block spanning `pad + size` bytes, `pad >= 1`.
    unsafe fn stamp(base: *mut u8, pad: usize, size: usize) -> *mut u8 {
        if base.is_null() {
            return base;
        }
        let scope = current();
        // SAFETY: the user pointer and its tag byte lie inside the block;
        // `pad` is the alignment, so the user pointer stays aligned.
        let user = unsafe { base.add(pad) };
        unsafe { user.sub(1).write(scope as u8) };
        record_alloc(scope, size);
        user
    }
}

// SAFETY: every block is a `System` block of `padded(layout)`; the user
// pointer sits `pad` bytes in, so `System` receives back exactly what it gave.
unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        // SAFETY: forwards the caller's non-zero layout, grown by the header.
        unsafe { Self::tagged(layout, |outer| System.alloc(outer)) }
    }

    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
        // SAFETY: as `alloc`.
        unsafe { Self::tagged(layout, |outer| System.alloc_zeroed(outer)) }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let Some((outer, pad)) = padded(layout) else {
            return;
        };
        // SAFETY: `ptr` came from `tagged` with this layout, so the tag byte
        // and the `System` base lie `1` and `pad` bytes before it.
        unsafe {
            record_free(ptr.sub(1).read(), layout.size());
            System.dealloc(ptr.sub(pad), outer);
        }
    }

    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
        // The grown outer layout must itself be valid, as `System.realloc` requires.
        let new_layout = Layout::from_size_align(new_size, layout.align()).ok();
        let (Some((outer, pad)), Some((new_outer, _))) =
            (padded(layout), new_layout.and_then(padded))
        else {
            return std::ptr::null_mut();
        };
        // SAFETY: as `dealloc`; on success `System` moved the header with the
        // data and the old block is gone, so its charge is released first.
        unsafe {
            let tag = ptr.sub(1).read();
            let base = System.realloc(ptr.sub(pad), outer, new_outer.size());
            if !base.is_null() {
                record_free(tag, layout.size());
            }
            Self::stamp(base, pad, new_size)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::{Counting, Scope, enter, snapshot};

    #[global_allocator]
    static ALLOC: Counting = Counting;

    // Each test owns a distinct scope: counters are process-wide and libtest runs tests in parallel.

    #[test]
    fn free_outside_scope_is_charged_to_allocating_scope() {
        let guard = enter(Scope::Frame);
        let block = std::hint::black_box(vec![0u8; 1000]);
        drop(guard);
        let inside = snapshot(Scope::Frame);
        assert_eq!(
            (inside.allocations, inside.bytes, inside.live),
            (1, 1000, 1000)
        );
        drop(block);
        let after = snapshot(Scope::Frame);
        assert_eq!((after.live, after.peak), (0, 1000));
    }

    #[test]
    fn nested_scope_restores_outer_and_realloc_retags() {
        let outer = enter(Scope::KernelCall);
        let mut grown: Vec<u8> = Vec::with_capacity(8);
        {
            let _inner = enter(Scope::IoRequest);
            grown.reserve_exact(64);
        }
        let tail = std::hint::black_box(vec![1u8; 16]);
        drop(outer);
        let kernel = snapshot(Scope::KernelCall);
        let io = snapshot(Scope::IoRequest);
        assert_eq!((kernel.allocations, kernel.bytes, kernel.live), (2, 24, 16));
        assert_eq!((io.allocations, io.bytes, io.live), (1, 64, 64));
        drop((grown, tail));
        assert_eq!(snapshot(Scope::IoRequest).live, 0);
    }

    #[test]
    fn overaligned_blocks_stay_aligned() {
        #[repr(align(4096))]
        struct Page([u8; 4096]);
        let page = Box::new(Page([0; 4096]));
        assert_eq!(std::ptr::from_ref(&*page).addr() % 4096, 0);
        assert_eq!(page.0[4095], 0);
    }

    #[test]
    fn realloc_past_header_capacity_fails_and_keeps_block() {
        use std::alloc::{GlobalAlloc, Layout};
        let layout = Layout::new::<u64>();
        // Valid for the caller, but no room left for the header in front.
        let new_size = isize::MAX.cast_unsigned() - 7;
        // SAFETY: `layout` is non-zero; `new_size` rounded to its alignment
        // fits `isize`; the block is freed with the layout it was made with.
        unsafe {
            let block = Counting.alloc(layout);
            assert!(!block.is_null());
            assert!(Counting.realloc(block, layout, new_size).is_null());
            Counting.dealloc(block, layout);
        }
    }

    #[test]
    fn spawned_threads_start_at_process_scope() {
        let _turn = enter(Scope::Turn);
        let seen = std::thread::spawn(super::current)
            .join()
            .map_err(|_| "join");
        assert_eq!(seen, Ok(Scope::Process));
    }

    #[test]
    fn overflowing_depth_keeps_outer_scope() {
        let _session = enter(Scope::Session);
        let guards: Vec<_> = (0..20).map(|_| enter(Scope::Process)).collect();
        assert_eq!(super::current(), Scope::Process);
        drop(guards);
        assert_eq!(super::current(), Scope::Session);
    }
}