mace-kv 0.2.0

use crate::must_true;
pub(crate) mod adapter;
pub(crate) mod buffer;
pub(crate) mod cache;
pub(crate) mod data;
pub(crate) mod evictor;
pub(crate) mod flow;
pub(crate) mod publish;

pub(crate) mod flush;
pub mod table;

use std::sync::Arc;

use dashmap::{DashMap, Entry};
use rustc_hash::FxHasher;
use std::hash::BuildHasherDefault;

use crate::{
    Options,
    cc::context::Context,
    map::buffer::Pool,
    static_assert,
    types::{
        refbox::{BoxRef, BoxView},
        traits::{IHeader, ILoader},
    },
    utils::{
        Handle, MutRef,
        data::{GroupPositions, Position},
        lru::{CachePriority, ShardPriorityLru},
    },
};

pub(crate) type PagesMap = DashMap<u64, BoxRef, BuildHasherDefault<FxHasher>>;
pub(crate) type JunksMap = DashMap<u64, RetiredChain, BuildHasherDefault<FxHasher>>;
pub(crate) type Node = crate::types::node::Node<Loader>;
pub(crate) type Page = crate::types::page::Page<Loader>;
type NodePins = DashMap<u64, BoxRef, BuildHasherDefault<FxHasher>>;

#[derive(Clone, Copy)]
struct CacheKey(u128);

impl CacheKey {
    const fn new(bucket_id: u64, addr: u64) -> Self {
        Self((bucket_id as u128) << 64 | addr as u128)
    }

    const fn raw(self) -> u128 {
        self.0
    }
}

#[derive(Default, Clone)]
pub(crate) struct SparseFrontier {
    mask: u128,
    items: Vec<(u8, Position)>,
}

impl SparseFrontier {
    pub(crate) fn is_empty(&self) -> bool {
        self.items.is_empty()
    }

    pub(crate) fn merge_group(&mut self, group: u8, lsn: Position) {
        static_assert!(Options::MAX_CONCURRENT_WRITE <= 128);
        let bit = 1u128 << group;
        if self.mask & bit == 0 {
            self.mask |= bit;
            self.items.push((group, lsn));
            return;
        }
        for (g, pos) in self.items.iter_mut() {
            if *g == group {
                *pos = (*pos).max(lsn);
                return;
            }
        }
        unreachable!("frontier mask/item mismatch for group {}", group);
    }

    pub(crate) fn merge_sparse(&mut self, other: &SparseFrontier) {
        for &(group, lsn) in &other.items {
            self.merge_group(group, lsn);
        }
    }

    pub(crate) fn apply_to(&self, dst: &mut GroupPositions) {
        for &(group, lsn) in &self.items {
            let idx = group as usize;
            must_true!(idx < dst.len());
            dst[idx] = dst[idx].max(lsn);
        }
    }
}

#[derive(Default, Clone)]
pub(crate) struct RetiredChain {
    // append-only retired logical addresses carried with a base lineage
    //
    // duplicates may appear when shared history pages or inherited lineage
    // contribute the same logical addr more than once
    // checkpoint snapshot normalizes them before metadata updates
    pub(crate) addrs: Vec<u64>,
    pub(crate) frontier: SparseFrontier,
}

pub struct Loader {
    pub(crate) pool: Handle<Pool>,
    pub(crate) ctx: Handle<Context>,
    pub(crate) lru: Handle<ShardPriorityLru<BoxRef>>,
    pub(crate) node_pins: MutRef<NodePins>,
    pub(crate) bucket_id: u64,
    pub(crate) reader: Arc<dyn IDataReader>,
}

impl Drop for Loader {
    fn drop(&mut self) {}
}

impl Loader {
    pub const PIN_CAP: usize = 64;
    const PIN_SHARDS: usize = 16;

    #[inline]
    fn cache_key(&self, addr: u64) -> u128 {
        // logical addr is only bucket-local, while this LRU is shared by all bucket contexts.
        CacheKey::new(self.bucket_id, addr).raw()
    }

    pub(crate) fn new_node_pins(cap: usize) -> NodePins {
        DashMap::with_capacity_and_hasher_and_shard_amount(
            cap,
            BuildHasherDefault::<FxHasher>::default(),
            Self::PIN_SHARDS,
        )
    }
}

impl ILoader for Loader {
    fn copy_detached(&self) -> Self {
        Self {
            pool: self.pool,
            ctx: self.ctx,
            lru: self.lru,
            node_pins: MutRef::new(Self::new_node_pins(Self::PIN_CAP)),
            bucket_id: self.bucket_id,
            reader: self.reader.clone(),
        }
    }

    fn copy(&self) -> Self {
        Self {
            pool: self.pool,
            ctx: self.ctx,
            lru: self.lru,
            node_pins: self.node_pins.clone(),
            bucket_id: self.bucket_id,
            reader: self.reader.clone(),
        }
    }

    fn pin(&self, data: BoxRef) {
        self.node_pins.insert(data.header().addr, data);
    }

    fn load_pinned(&self, addr: u64) -> BoxView {
        if let Some(p) = self.node_pins.get(&addr) {
            return p.view();
        }
        let x = if let Some(x) = self.pool.get_dirty_page(addr) {
            x
        } else {
            self.reader.load_data(self.bucket_id, addr, &|_| {})
        };
        let e = self.node_pins.entry(addr);
        match e {
            Entry::Occupied(o) => o.get().view(),
            Entry::Vacant(v) => {
                let r = x.view();
                v.insert(x);
                r
            }
        }
    }

    fn load_sibling(&self, addr: u64) -> BoxRef {
        let key = self.cache_key(addr);
        if let Some(x) = self.lru.get(key) {
            return x.clone();
        }
        if let Some(x) = self.pool.get_dirty_page(addr) {
            return x;
        }
        self.reader.load_data(self.bucket_id, addr, &|b| {
            self.lru
                .add(CachePriority::High, key, b.header().total_size as usize, b);
        })
    }

    fn load_blob(&self, addr: u64, cache: bool) -> BoxRef {
        let key = self.cache_key(addr);
        if cache && let Some(x) = self.lru.get(key) {
            return x.clone();
        }
        if let Some(x) = self.pool.get_dirty_page(addr) {
            return x;
        }
        self.reader.load_blob(self.bucket_id, addr, &|b| {
            if cache {
                self.lru.add(
                    CachePriority::Low,
                    key,
                    b.header().total_size as usize,
                    b.to_box(),
                );
            }
        })
    }
}

pub(crate) enum SharedState {
    Quit,
    Evict(u64),
}

pub trait IDataReader: Send + Sync {
    fn load_data(&self, bucket_id: u64, addr: u64, cache: &dyn Fn(BoxRef)) -> BoxRef;

    fn load_blob(&self, bucket_id: u64, addr: u64, cache: &dyn Fn(BoxView)) -> BoxRef;
}