//! Scope-tagged counting allocator (`SMH-SPEC-SPEC0001`, Memory accounting).
//!
//! [`Counting`] wraps [`System`](std::alloc::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`].
//!
//! The only unsafe code is the [`GlobalAlloc`](std::alloc::GlobalAlloc)
//! implementation and its header arithmetic, kept in one private module;
//! the public surface is safe.

#![deny(unsafe_code)]

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 _: () = assert!(
    SCOPES.len() <= u8::MAX as usize,
    "scope tags must fit in one byte"
);
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] }) };
    // Kept separately so allocation hot paths do not copy `Stack` from its `Cell`.
    static CURRENT_SCOPE: Cell<Scope> = const { Cell::new(Scope::Process) };
    // Zero means inactive; scopes use their discriminant plus one.
    static ACTIVE_MEASUREMENT: Cell<u8> = const { Cell::new(0) };
    static MEASUREMENT: Cell<Measurement> = const { Cell::new(Measurement::EMPTY) };
}

#[derive(Clone, Copy)]
struct Measurement {
    live: u64,
    peak: u64,
}

impl Measurement {
    const EMPTY: Self = Self { live: 0, peak: 0 };
}

const _: () = assert!(
    std::mem::size_of::<Measurement>() == 16,
    "measurement state must remain two u64 words"
);

#[inline]
fn measures(scope: u8) -> bool {
    ACTIVE_MEASUREMENT
        .try_with(|active| active.get() == scope + 1)
        .unwrap_or(false)
}

#[inline]
fn measure_allocation(size: u64) {
    // TLS teardown must not make allocator accounting panic or abort allocation.
    let _measurement_update: std::result::Result<(), std::thread::AccessError> = MEASUREMENT
        .try_with(|cell| {
            let mut measurement = cell.get();
            measurement.live += size;
            measurement.peak = measurement.peak.max(measurement.live);
            cell.set(measurement);
        });
}

#[inline]
fn measure_free(size: u64) {
    // TLS teardown must not make allocator accounting panic or abort deallocation.
    let _measurement_update: std::result::Result<(), std::thread::AccessError> = MEASUREMENT
        .try_with(|cell| {
            let mut measurement = cell.get();
            measurement.live = measurement.live.saturating_sub(size);
            cell.set(measurement);
        });
}

#[inline]
fn current() -> Scope {
    CURRENT_SCOPE.try_with(Cell::get).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 current = 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);
                let current = stack
                    .items
                    .get(stack.len.wrapping_sub(1))
                    .copied()
                    .unwrap_or(Scope::Process);
                cell.set(stack);
                current
            })
            .unwrap_or(Scope::Process);
        // Scope restoration is best-effort during thread-local destruction.
        let _scope_update: std::result::Result<(), std::thread::AccessError> =
            CURRENT_SCOPE.try_with(|cell| cell.set(current));
    }
}

/// 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, current) = 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,
                stack
                    .items
                    .get(stack.len.wrapping_sub(1))
                    .copied()
                    .unwrap_or(Scope::Process),
            )
        })
        .unwrap_or((0, Scope::Process));
    // Scope setup is best-effort if thread-local state is already being torn down.
    let _scope_update: std::result::Result<(), std::thread::AccessError> =
        CURRENT_SCOPE.try_with(|cell| cell.set(current));
    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),
    }
}

struct MeasurementGuard {
    scope: Option<Guard>,
    active: u8,
    measurement: Measurement,
}

impl Drop for MeasurementGuard {
    fn drop(&mut self) {
        drop(self.scope.take());
        ACTIVE_MEASUREMENT.set(self.active);
        MEASUREMENT.set(self.measurement);
    }
}

/// Run `work` charged to `scope` and report what it cost.
///
/// `allocations`, `bytes`, and `peak` describe this invocation only.
/// `live` is the scope's process-wide value after it returns.
/// Measurements on the same scope must not overlap on one thread.
pub fn measure<T>(scope: Scope, work: impl FnOnce() -> T) -> (T, Stats) {
    let before = snapshot(scope);
    let guard = MeasurementGuard {
        scope: Some(enter(scope)),
        active: ACTIVE_MEASUREMENT.replace(scope as u8 + 1),
        measurement: MEASUREMENT.replace(Measurement::EMPTY),
    };
    let out = work();
    let measurement = MEASUREMENT.get();
    drop(guard);
    let after = snapshot(scope);
    (
        out,
        Stats {
            allocations: after.allocations - before.allocations,
            bytes: after.bytes - before.bytes,
            live: after.live,
            peak: measurement.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);
    if measures(scope as u8) {
        measure_allocation(size as u64);
    }
}

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);
        if measures(tag) {
            measure_free(size as u64);
        }
    }
}

/// Counting wrapper around [`System`](std::alloc::System); install with
/// `#[global_allocator]`.
pub struct Counting;

/// Block layout: `System` holds `pad = align` header bytes in front of
/// each user block, the last header byte being the scope tag.
#[expect(
    unsafe_code,
    reason = "GlobalAlloc has no safe interface; implementing it is the only way to attribute every allocation, third-party ones included, to a scope"
)]
mod raw {
    use super::{Counting, current, record_alloc, record_free};
    use std::alloc::{GlobalAlloc, Layout, System};

    /// The `System` layout for `layout` and its header size; `pad >= 1`
    /// because every alignment is at least one.
    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))
    }

    /// A fresh tagged block for `layout`, or null.
    fn tagged(layout: Layout, zeroed: bool) -> *mut u8 {
        let Some((outer, pad)) = padded(layout) else {
            return std::ptr::null_mut();
        };
        let base = if zeroed {
            // SAFETY: `outer` is `layout` grown by `pad >= 1`, so its size is non-zero.
            unsafe { System.alloc_zeroed(outer) }
        } else {
            // SAFETY: `outer` is `layout` grown by `pad >= 1`, so its size is non-zero.
            unsafe { System.alloc(outer) }
        };
        // SAFETY: `base` is null or a fresh `System` block of `outer`,
        // which spans `pad + layout.size()` bytes.
        unsafe { stamp(base, 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 live, writable block spanning `pad + size`
    /// bytes, aligned to `pad`, with `pad >= 1`.
    unsafe fn stamp(base: *mut u8, pad: usize, size: usize) -> *mut u8 {
        if base.is_null() {
            return base;
        }
        let scope = current();
        // `pad` is the alignment, so the user pointer stays aligned.
        let user = base.wrapping_add(pad);
        // SAFETY: `1 <= pad`, so the tag byte `user - 1` is the last header
        // byte, inside the block the caller hands over.
        unsafe { user.wrapping_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, and the user block keeps the size and alignment of `layout`.
    // Header and base are reached from the caller's pointer, keeping its
    // provenance: a `Box` protects the pointer it frees, so any other
    // pointer to the block would be a foreign access (Tree Borrows).
    unsafe impl GlobalAlloc for Counting {
        unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
            tagged(layout, false)
        }

        unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
            tagged(layout, true)
        }

        unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
            let Some((outer, pad)) = padded(layout) else {
                return;
            };
            // SAFETY: the caller passes a block this allocator made with
            // `layout`, so its tag byte lies one byte before `ptr`.
            let tag = unsafe { ptr.wrapping_sub(1).read() };
            record_free(tag, layout.size());
            // SAFETY: the `System` base lies `pad` bytes before `ptr` and was
            // allocated with `outer = padded(layout)`.
            unsafe { System.dealloc(ptr.wrapping_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 in `dealloc`, the tag byte lies one byte before `ptr`.
            let tag = unsafe { ptr.wrapping_sub(1).read() };
            // SAFETY: the base `pad` bytes before `ptr` is a live `System`
            // block of `outer`; `new_outer` is a valid layout of the same
            // alignment and non-zero size.
            let base = unsafe { System.realloc(ptr.wrapping_sub(pad), outer, new_outer.size()) };
            // On success `System` moved the header with the data and the old
            // block is gone, so its charge is released first.
            if !base.is_null() {
                record_free(tag, layout.size());
            }
            // SAFETY: `base` is null or a `System` block of `new_outer`,
            // which spans `pad + new_size` bytes aligned to `pad`.
            unsafe { stamp(base, pad, new_size) }
        }
    }
}

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

    #[global_allocator]
    static ALLOC: Counting = Counting;

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

    #[cfg(target_pointer_width = "64")]
    #[test]
    fn active_measurement_state_fits_two_words() {
        assert_eq!(std::mem::size_of::<super::Measurement>(), 16);
    }

    #[test]
    fn foreign_free_tags_are_ignored_during_and_outside_measurement() {
        let active = super::ACTIVE_MEASUREMENT.replace(0);
        let measurement = super::MEASUREMENT.replace(super::Measurement { live: 7, peak: 11 });
        let before = snapshot(Scope::Session);
        for tag in 6..=u8::MAX {
            super::record_free(tag, 13);
        }
        assert_eq!(snapshot(Scope::Session), before);
        assert_eq!(
            (super::MEASUREMENT.get().live, super::MEASUREMENT.get().peak),
            (7, 11)
        );

        super::ACTIVE_MEASUREMENT.set(Scope::IoRequest as u8 + 1);
        for tag in 6..=u8::MAX {
            super::record_free(tag, 13);
        }
        assert_eq!(snapshot(Scope::Session), before);
        assert_eq!(
            (super::MEASUREMENT.get().live, super::MEASUREMENT.get().peak),
            (7, 11)
        );

        super::ACTIVE_MEASUREMENT.set(active);
        super::MEASUREMENT.set(measurement);
    }

    #[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 measure_nested_scope_restores_outer_peak() {
        let (_, outer) = measure(Scope::Turn, || {
            let first = std::hint::black_box(vec![0u8; 4]);
            let (_, inner) = measure(Scope::Process, || std::hint::black_box(vec![0u8; 8]));
            assert_eq!(inner.peak, 8);
            drop(first);
            std::hint::black_box(vec![0u8; 16])
        });
        assert_eq!((outer.allocations, outer.bytes, outer.peak), (2, 20, 16));
    }

    #[test]
    fn panicking_measure_restores_enclosing_measurement() {
        let (_, outer) = measure(Scope::Turn, || {
            let result = std::panic::catch_unwind(|| {
                let _ = measure(Scope::Process, || {
                    let _block = std::hint::black_box(vec![0u8; 8]);
                    std::panic::resume_unwind(Box::new(()));
                });
            });
            assert!(result.is_err());
            std::hint::black_box(vec![0u8; 16])
        });
        assert_eq!(outer.peak, 16);
    }

    #[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 zeroed_blocks_are_zero_after_a_dirty_free() {
        drop(std::hint::black_box(vec![0xAAu8; 64]));
        let zeroed = std::hint::black_box(vec![0u8; 64]);
        assert!(zeroed.iter().all(|byte| *byte == 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]
    #[expect(
        unsafe_code,
        reason = "drives GlobalAlloc directly: no safe caller can request a size past the header capacity"
    )]
    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` has non-zero size.
        let block = unsafe { Counting.alloc(layout) };
        assert!(!block.is_null());
        // SAFETY: `block` is live with `layout`; `new_size` is non-zero and,
        // rounded to the alignment, fits `isize`.
        assert!(unsafe { Counting.realloc(block, layout, new_size) }.is_null());
        // SAFETY: the failed `realloc` kept `block`, still live with `layout`.
        unsafe { 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);
    }
}