//! Loaded parameter sets, kept between decodes.
//!
//! A decode through the integer operator set lowers its parameter set to `WHG1` graphs
//! and loads each one (`extensions/integer-operators.md` §10.5), the bounds pass
//! included (load step 8, §8). For a large set that costs more than the decode itself,
//! and it is the same work for every payload the set decodes. A sink may keep the
//! verdict of a set it has loaded, by the set's digest (§8.6). This module keeps the
//! loaded graphs themselves, for every integer decode this crate runs from a set's
//! bytes through `int_payload::decode_int` and `int_payload::declared_cost`, so also for
//! the neural sink and the browser crate, which decode through them.
//! (`int_synth::decode_from_spec`, which the byte-equality crate shares, loads its set
//! every time.)
//!
//! - **Keyed by the set's bytes, never by a name.** An entry is found by its [`SetKey`]:
//!   the capability, the lowering, the profile the graphs were loaded under, and the
//!   SHA-256 of the listing of the members the decode reads, each listed by its own
//!   SHA-256 (`pin::listing_from`). That is the digest a pin in the form
//!   `wai.prior.set` over exactly those members names (SPEC §3.1). Two sets whose
//!   members have the same names and whose bytes differ in one byte have two keys.
//! - **Only what loaded.** An entry is made from what a load returned, after it
//!   returned. A [`wai_ops::Loaded`] graph is made only by `wai_ops::graph::load`, after
//!   every load step, so nothing reaches the store that did not pass the bounds pass. A
//!   refusal is not kept: a set that is refused is lowered and loaded again on its next
//!   decode, and refused again with the same refusal.
//! - **Nothing else moves.** A decode makes every check before the load as it did, in
//!   integer-payloads §7's order. A hit stands in for the load step alone, whose outcome
//!   for those bytes it is. The graphs run as they would after a fresh load, so the
//!   output is the same bytes.
//! - **Bounded.** The entries ([`Resident`], the shared allocation that holds each, and
//!   [`LoadedSets::INDEX_BYTES_PER_ENTRY`]) are kept within a budget:
//!   [`LoadedSets::DEFAULT_BUDGET`] unless the sink sets one
//!   ([`LoadedSets::set_budget`]). The least recently used entry goes first, and a set
//!   larger than the whole budget is not kept. A budget of 0 turns the
//!   store off; the set is then not even hashed.
//! - **Memory and work.** A decode's declared cost (integer-payloads §7) counts what the
//!   decode holds beyond the payload and its parameter set. A set lowered and loaded is
//!   the set's, the same for every payload, so an entry here is outside every decode's
//!   declared working bytes, and the budget is what bounds it.
//!
//! There is one store per process ([`LoadedSets::global`]), behind a lock. On `wasm32`
//! there is one thread, and the lock is never contended. A load runs outside the lock,
//! so two threads that miss the same set at once both load it, and one entry is kept.

use std::any::Any;
use std::collections::BTreeMap;
use std::sync::{Arc, Mutex, MutexGuard, PoisonError};

use wai_ops::{Loaded, Profile};

/// What a loaded set holds on the heap, beyond the value itself.
pub trait Resident {
    fn heap_bytes(&self) -> u64;
}

impl Resident for Loaded {
    fn heap_bytes(&self) -> u64 {
        self.graph().heap_bytes()
    }
}

/// The conditional codec's two networks (`int_hyper_synth::ops::Nets`). The impl is here,
/// not beside the type, because the byte-equality crate includes `int_hyper_synth.rs`
/// and not this module.
impl Resident for crate::int_hyper_synth::ops::Nets {
    fn heap_bytes(&self) -> u64 {
        self.h_s.graph().heap_bytes() + self.g_s.graph().heap_bytes()
    }
}

/// The MLIC++-topology codec's three networks (`int_mlicpp::ops::Nets`), here for the same
/// reason.
impl Resident for crate::int_mlicpp::ops::Nets {
    fn heap_bytes(&self) -> u64 {
        self.h_s.graph().heap_bytes() + self.entropy.graph().heap_bytes() + self.g_s.graph().heap_bytes()
    }
}

/// `wai.image.jpegai`'s scale path (`jpegai_ops::ScalePath`): its two scale graphs and the
/// tables the entropy decodes read.
impl Resident for crate::jpegai_ops::ScalePath {
    fn heap_bytes(&self) -> u64 {
        crate::jpegai_ops::ScalePath::heap_bytes(self)
    }
}

/// `wai.image.jpegai`'s other integer networks (`jpegai_ops::Networks`): its four graphs.
impl Resident for crate::jpegai_ops::Networks {
    fn heap_bytes(&self) -> u64 {
        crate::jpegai_ops::Networks::heap_bytes(self)
    }
}

/// The key of a loaded set: the capability whose driver loaded it, the lowering
/// (integer-operators §13), the profile its graphs were loaded under (§11), and the
/// SHA-256 of the members it was lowered from, as [`SetKey::of`] forms it.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
pub struct SetKey {
    pub capability: &'static str,
    pub lowering: &'static str,
    pub profile: Profile,
    pub digest: [u8; 32],
}

impl SetKey {
    /// The key of the set whose members a decode reads, `members` (name → bytes): the
    /// capability, the lowering and the profile, and [`SetKey::listing_digest`] of the
    /// members. `None` where the names cannot make a listing; no set this crate decodes
    /// has such a name, and such a set is then loaded and not kept.
    pub fn of(capability: &'static str, lowering: &'static str, profile: Profile, members: &BTreeMap<String, &[u8]>) -> Option<SetKey> {
        Some(SetKey { capability, lowering, profile, digest: SetKey::listing_digest(members)? })
    }

    /// The SHA-256 of the listing of `members` (`pin::listing_from`, each member by its
    /// own SHA-256): the digest a pin in the form `wai.prior.set` over exactly those
    /// members names. It is the digest of a [`SetKey`].
    pub fn listing_digest(members: &BTreeMap<String, &[u8]>) -> Option<[u8; 32]> {
        crate::pin::listing_from(members).ok().map(|l| crate::pin::sha256(&l))
    }
}

/// What a store has done, and what it holds.
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
#[non_exhaustive]
pub struct Stats {
    /// Lookups that found their set loaded.
    pub hits: u64,
    /// Lookups that did not, and loaded it.
    pub misses: u64,
    /// Sets kept after they loaded.
    pub kept: u64,
    /// Entries removed to make room, or to fit a smaller budget.
    pub evicted: u64,
    /// Sets that loaded and were not kept, being larger than the whole budget.
    pub too_large: u64,
    /// The entries held now.
    pub entries: u64,
    /// Their bytes, as counted against the budget.
    pub resident_bytes: u64,
    /// The budget.
    pub budget: u64,
}

/// A store of loaded parameter sets (the module docs): an entry per [`SetKey`], within a
/// byte budget, the least recently used going first.
pub struct LoadedSets(Mutex<Inner>);

struct Inner {
    budget: u64,
    resident: u64,
    clock: u64,
    sets: BTreeMap<SetKey, Entry>,
    by_use: BTreeMap<u64, SetKey>,
    stats: Stats,
}

struct Entry {
    set: Arc<dyn Any + Send + Sync>,
    bytes: u64,
    used: u64,
}

static GLOBAL: LoadedSets = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);

impl LoadedSets {
    /// The budget a process starts with: 256 MiB. A set holds its weights as `i64` once
    /// lowered, four times their bytes in a `WHM1` file, so this keeps a set of about
    /// 60 MB of `WHM1` networks with room beside it for small ones.
    pub const DEFAULT_BUDGET: u64 = 256 << 20;

    /// What an entry is counted for beyond its set: its share of the store's two indexes
    /// (by key, and by last use). It is more than either index's node for one entry, so
    /// the count is never under what the store holds.
    pub const INDEX_BYTES_PER_ENTRY: u64 = 4096;

    /// An empty store with the budget `budget` (bytes; 0 is off).
    pub const fn new(budget: u64) -> LoadedSets {
        let stats = Stats { hits: 0, misses: 0, kept: 0, evicted: 0, too_large: 0, entries: 0, resident_bytes: 0, budget: 0 };
        LoadedSets(Mutex::new(Inner { budget, resident: 0, clock: 0, sets: BTreeMap::new(), by_use: BTreeMap::new(), stats }))
    }

    /// The process's store, which every decode in this crate uses. It starts with
    /// [`LoadedSets::DEFAULT_BUDGET`].
    pub fn global() -> &'static LoadedSets {
        &GLOBAL
    }

    fn lock(&self) -> MutexGuard<'_, Inner> {
        // Nothing panics while the lock is held; a poisoned lock still holds a whole store.
        self.0.lock().unwrap_or_else(PoisonError::into_inner)
    }

    /// The budget, in bytes.
    pub fn budget(&self) -> u64 {
        self.lock().budget
    }

    /// Set the budget, in bytes, evicting the least recently used entries until the rest
    /// fit. 0 empties the store and turns it off: a decode then loads its set every time,
    /// as it would with no store.
    pub fn set_budget(&self, bytes: u64) {
        let mut i = self.lock();
        i.budget = bytes;
        i.evict_to(bytes);
    }

    /// Remove every entry, keeping the budget and the counts.
    pub fn clear(&self) {
        let mut i = self.lock();
        i.sets.clear();
        i.by_use.clear();
        i.resident = 0;
    }

    /// Set the counts of [`Stats`] to 0.
    pub fn reset_stats(&self) {
        self.lock().stats = Stats::default();
    }

    /// The counts, and what the store holds now.
    pub fn stats(&self) -> Stats {
        let i = self.lock();
        Stats { entries: i.sets.len() as u64, resident_bytes: i.resident, budget: i.budget, ..i.stats }
    }

    /// Whether the set `key` is held.
    pub fn contains(&self, key: &SetKey) -> bool {
        self.lock().sets.contains_key(key)
    }

    /// The set `key()` names, from the store when it is held (a hit), or else `load()`
    /// (a miss), kept when it loads and fits the budget. A refusal is returned as it is
    /// and keeps nothing. With a budget of 0, `load()` alone: `key` is not called.
    ///
    /// `load` must lower and load the set from exactly the bytes `key` was formed from,
    /// by the lowering and under the profile it names.
    pub fn get_or_load<T, E>(&self, key: impl FnOnce() -> Option<SetKey>, load: impl FnOnce() -> Result<T, E>) -> Result<Arc<T>, E>
    where
        T: Resident + Send + Sync + 'static,
    {
        if self.budget() == 0 {
            return load().map(Arc::new);
        }
        let Some(key) = key() else { return load().map(Arc::new) };
        {
            let mut i = self.lock();
            if let Some(set) = i.touch(&key).and_then(|s| s.downcast::<T>().ok()) {
                i.stats.hits += 1;
                return Ok(set);
            }
            i.stats.misses += 1;
        }
        let set = Arc::new(load()?);
        let bytes = entry_bytes(&*set);
        self.lock().keep(key, set.clone(), bytes);
        Ok(set)
    }
}

/// The bytes an entry is counted for: the shared allocation that holds the set (two
/// counts, then the value), what the value holds on the heap, and its share of the
/// indexes.
fn entry_bytes<T: Resident>(t: &T) -> u64 {
    use std::mem::{align_of, size_of};
    let align = align_of::<T>().max(align_of::<usize>());
    let shared = ((2 * size_of::<usize>()).next_multiple_of(align_of::<T>()) + size_of::<T>()).next_multiple_of(align);
    shared as u64 + t.heap_bytes() + LoadedSets::INDEX_BYTES_PER_ENTRY
}

impl Inner {
    fn tick(&mut self) -> u64 {
        self.clock += 1;
        self.clock
    }

    /// The entry `key`, marked as used now.
    fn touch(&mut self, key: &SetKey) -> Option<Arc<dyn Any + Send + Sync>> {
        let now = self.tick();
        let Inner { sets, by_use, .. } = self;
        let e = sets.get_mut(key)?;
        by_use.remove(&e.used);
        by_use.insert(now, *key);
        e.used = now;
        Some(e.set.clone())
    }

    /// Keep `set`, of `bytes`, under `key`, evicting the least recently used entries to
    /// make room; a set larger than the budget is not kept.
    fn keep(&mut self, key: SetKey, set: Arc<dyn Any + Send + Sync>, bytes: u64) {
        if self.budget == 0 {
            return;
        }
        if bytes > self.budget {
            self.stats.too_large += 1;
            return;
        }
        if let Some(old) = self.sets.remove(&key) {
            self.by_use.remove(&old.used);
            self.resident -= old.bytes;
        }
        self.evict_to(self.budget - bytes);
        let used = self.tick();
        self.by_use.insert(used, key);
        self.sets.insert(key, Entry { set, bytes, used });
        self.resident += bytes;
        self.stats.kept += 1;
    }

    /// Evict the least recently used entries until the rest hold at most `limit` bytes.
    fn evict_to(&mut self, limit: u64) {
        while self.resident > limit {
            let Some((_, key)) = self.by_use.pop_first() else { break };
            if let Some(e) = self.sets.remove(&key) {
                self.resident -= e.bytes;
                self.stats.evicted += 1;
            }
        }
    }
}

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

    /// A set of a given size, and a set of another type.
    struct Blob(u64);
    impl Resident for Blob {
        fn heap_bytes(&self) -> u64 {
            self.0
        }
    }
    struct Other;
    impl Resident for Other {
        fn heap_bytes(&self) -> u64 {
            0
        }
    }

    fn key(d: u8) -> SetKey {
        SetKey { capability: "c", lowering: "l", profile: Profile::Acc64, digest: [d; 32] }
    }

    fn get(s: &LoadedSets, d: u8, size: u64) -> Arc<Blob> {
        s.get_or_load(|| Some(key(d)), || Ok::<_, ()>(Blob(size))).unwrap()
    }

    /// The bytes one entry of `size` heap bytes is counted for.
    fn one(size: u64) -> u64 {
        entry_bytes(&Blob(size))
    }

    #[test]
    fn a_second_lookup_is_a_hit_and_returns_the_kept_set() {
        let s = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
        let a = get(&s, 1, 100);
        let loads = Cell::new(0);
        let b = s.get_or_load(|| Some(key(1)), || { loads.set(loads.get() + 1); Ok::<_, ()>(Blob(100)) }).unwrap();
        assert!(Arc::ptr_eq(&a, &b) && loads.get() == 0, "the second lookup is answered from the store");
        let st = s.stats();
        assert_eq!((st.hits, st.misses, st.kept, st.entries, st.resident_bytes), (1, 1, 1, 1, one(100)));
    }

    #[test]
    fn a_refusal_keeps_nothing() {
        let s = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
        for _ in 0..2 {
            assert_eq!(s.get_or_load(|| Some(key(1)), || Err::<Blob, _>("refused")).err(), Some("refused"));
        }
        let st = s.stats();
        assert_eq!((st.hits, st.misses, st.kept, st.entries), (0, 2, 0, 0), "a refused set is loaded, and refused, each time");
    }

    #[test]
    fn a_key_differs_by_its_digest_and_a_type_never_crosses() {
        let s = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
        get(&s, 1, 10);
        get(&s, 2, 10);
        assert_eq!(s.stats().entries, 2);
        // A set of another type under a held key is loaded, never cast.
        let other = s.get_or_load(|| Some(key(1)), || Ok::<_, ()>(Other)).map(|_| ());
        assert!(other.is_ok());
        assert_eq!(s.stats().misses, 3);
    }

    #[test]
    fn the_least_recently_used_entry_goes_first() {
        let s = LoadedSets::new(2 * one(10));
        get(&s, 1, 10);
        get(&s, 2, 10);
        get(&s, 1, 10); // 1 is now the most recently used
        get(&s, 3, 10);
        assert!(s.contains(&key(1)) && s.contains(&key(3)) && !s.contains(&key(2)));
        let st = s.stats();
        assert_eq!((st.evicted, st.entries, st.resident_bytes), (1, 2, 2 * one(10)));
        // A smaller budget evicts down to it; a set past the whole budget is not kept.
        s.set_budget(one(10));
        assert!(s.contains(&key(3)) && !s.contains(&key(1)));
        get(&s, 4, 11);
        assert!(!s.contains(&key(4)) && s.contains(&key(3)));
        assert_eq!(s.stats().too_large, 1);
    }

    #[test]
    fn a_budget_of_zero_neither_keeps_nor_hashes() {
        let s = LoadedSets::new(LoadedSets::DEFAULT_BUDGET);
        get(&s, 1, 10);
        s.set_budget(0);
        assert_eq!(s.stats().entries, 0, "a budget of 0 empties the store");
        let keyed = Cell::new(false);
        let r = s.get_or_load(|| { keyed.set(true); Some(key(1)) }, || Ok::<_, ()>(Blob(10)));
        assert!(r.is_ok() && !keyed.get(), "no key is formed");
        assert_eq!((s.stats().misses, s.stats().entries), (1, 0));
    }
}
