armdb 0.11.2

use std::sync::Arc;

use quick_cache::Weighter;
use quick_cache::sync::Cache;

use crate::config::CacheConfig;
use crate::io::aligned_buf::AlignedBuf;

/// Cache key: identifies a 4096-byte aligned block on disk.
/// Includes `shard_id` because file_ids are assigned independently per shard
/// (each shard starts at file_id=1), so (file_id, block_offset) alone is ambiguous.
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
pub struct BlockKey {
    pub shard_id: u8,
    pub file_id: u32,
    pub block_offset: u64,
}

#[derive(Clone)]
struct BlockWeighter;

impl Weighter<BlockKey, Arc<AlignedBuf>> for BlockWeighter {
    fn weight(&self, _key: &BlockKey, _value: &Arc<AlignedBuf>) -> u64 {
        4096
    }
}

/// Block cache for `VarTree`, backed by `quick_cache` with S3-FIFO eviction.
/// Caches 4096-byte aligned disk blocks — the natural unit of O_DIRECT reads.
pub struct BlockCache {
    inner: Option<Cache<BlockKey, Arc<AlignedBuf>, BlockWeighter>>,
}

impl BlockCache {
    pub fn new(config: &CacheConfig) -> Self {
        if config.max_size == 0 {
            return Self { inner: None };
        }
        let cache = Cache::with_weighter(config.estimated_items, config.max_size, BlockWeighter);
        Self { inner: Some(cache) }
    }

    pub fn get(&self, key: &BlockKey) -> Option<Arc<AlignedBuf>> {
        self.inner.as_ref()?.get(key)
    }

    pub fn insert(&self, key: BlockKey, block: Arc<AlignedBuf>) {
        if let Some(cache) = &self.inner {
            cache.insert(key, block);
        }
    }

    /// Drop every cached block belonging to a retired file.
    ///
    /// This is RAM hygiene, not a correctness measure: file ids are never
    /// reused (`Shard::allocate_file_id` strictly increments and compaction
    /// outputs get fresh ids), so a block cached under a retired id can never
    /// be mistaken for a live one — the value cache's verify-on-hit depends on
    /// exactly that. Blocks left behind would only waste cache budget until
    /// eviction reclaims them.
    pub fn invalidate_file(&self, shard_id: u8, file_id: u32, total_bytes: u64) {
        if let Some(cache) = &self.inner {
            let num_blocks = total_bytes.div_ceil(4096);
            for i in 0..num_blocks {
                cache.remove(&BlockKey {
                    shard_id,
                    file_id,
                    block_offset: i * 4096,
                });
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::CacheConfig;
    use crate::io::aligned_buf::AlignedBuf;

    fn make_marked_buf(marker: u8) -> Arc<AlignedBuf> {
        let mut buf = AlignedBuf::zeroed(4096);
        buf[0] = marker;
        Arc::new(buf)
    }

    #[test]
    fn block_cache_disambiguates_file_ids_across_u16_boundary() {
        let cfg = CacheConfig {
            max_size: 1 << 20,
            estimated_items: 32,
        };
        let cache = BlockCache::new(&cfg);

        // Build two distinguishable 4 KiB buffers.
        let buf_a = make_marked_buf(0xAA);
        let buf_b = make_marked_buf(0xBB);

        let key_low = BlockKey {
            shard_id: 0,
            file_id: 65_535,
            block_offset: 0,
        };
        let key_high = BlockKey {
            shard_id: 0,
            file_id: 65_536,
            block_offset: 0,
        };

        cache.insert(key_low, buf_a.clone());
        cache.insert(key_high, buf_b.clone());

        let got_low = cache.get(&key_low).expect("low present");
        let got_high = cache.get(&key_high).expect("high present");

        assert_eq!(got_low[0], 0xAA);
        assert_eq!(got_high[0], 0xBB);
        assert!(!Arc::ptr_eq(&got_low, &got_high));
    }

    /// Characterises `invalidate_file` ahead of moving its call site out of
    /// compaction's critical section 2: every block of the named file goes,
    /// and nothing belonging to another file or shard is touched.
    #[test]
    fn invalidate_file_drops_only_that_file_s_blocks() {
        let cfg = CacheConfig {
            max_size: 1 << 20,
            estimated_items: 32,
        };
        let cache = BlockCache::new(&cfg);

        let retired = |offset| BlockKey {
            shard_id: 0,
            file_id: 7,
            block_offset: offset,
        };
        // Same file id, different shard — shard_id is part of the key.
        let other_shard = BlockKey {
            shard_id: 1,
            file_id: 7,
            block_offset: 0,
        };
        let other_file = BlockKey {
            shard_id: 0,
            file_id: 8,
            block_offset: 0,
        };

        cache.insert(retired(0), make_marked_buf(0x01));
        cache.insert(retired(4096), make_marked_buf(0x02));
        cache.insert(other_shard, make_marked_buf(0x03));
        cache.insert(other_file, make_marked_buf(0x04));

        cache.invalidate_file(0, 7, 8192);

        assert!(cache.get(&retired(0)).is_none());
        assert!(cache.get(&retired(4096)).is_none());
        assert_eq!(cache.get(&other_shard).expect("other shard kept")[0], 0x03);
        assert_eq!(cache.get(&other_file).expect("other file kept")[0], 0x04);
    }
}