obikmer/src/obikidxcache/src/index_cache.rs

use std::collections::HashMap;
use std::sync::Arc;

use obikindex::{KmerIndex, layer::KmerLayer};
use obikseq::CanonicalKmer;

use crate::meta_cache::MetaCache;

pub struct IndexCache {
    /// Owns a strong reference to the `KmerIndex` it was opened from — the
    /// cache cannot outlive its index (the `Arc` keeps it alive at least as
    /// long as the cache itself, enforced at runtime by refcounting rather
    /// than by a borrow-checked lifetime). `Arc`, not `&'a KmerIndex`,
    /// because `IndexCache` must itself be `'static`-capable to be shared
    /// across `obipipeline`'s `thread::spawn`-based workers (see
    /// `obikquery::query_layer`) — a plain borrow can never satisfy that,
    /// regardless of how it's threaded through. `index()` is a genuine
    /// accessor, but the field's real job is this ownership guarantee, kept
    /// even on paths that never call `index()`.
    raw_index: Arc<KmerIndex>,
    meta: MetaCache,

    /// Keyed by each partition's own index number, not renumbered —
    /// `partitions: Some(vec![5])` still reads back as partition `5`
    /// everywhere (`get_layer(5, ..)`, `find_in_partition(5, ..)`, …), not
    /// `0`. A `HashMap` rather than a flat `Vec` because the set of cached
    /// partitions doesn't have to be `0..n_partitions`, or even contiguous
    /// — see [`new`](Self::new).
    layer_cache: HashMap<usize, Vec<KmerLayer>>,
}

impl IndexCache {
    /// Opens every layer of the given `partitions` (every partition of the
    /// index, if `None`) once, up front, and keeps them alive for the
    /// cache's lifetime — each partition's own layer count is read
    /// independently (`KmerPartition::n_layers`), so partitions don't need
    /// a uniform layer count across each other. That matters for a caller
    /// like `obikmerge`, which only ever needs one specific partition of a
    /// destination index still being merged into elsewhere — the other
    /// partitions can legitimately have a different layer count at that
    /// moment (built in parallel, mid-run), without that affecting a cache
    /// scoped to just the one partition being read.
    ///
    /// Panics if any requested partition fails to open (e.g. an
    /// out-of-range partition number, or a layer that fails to open): a
    /// cache can only be built on partitions that actually exist and are
    /// complete, so any failure here is a caller error, not a case to
    /// design around with a `Result`.
    pub fn new(index: Arc<KmerIndex>, partitions: Option<Vec<usize>>) -> Self {
        let selected = partitions.unwrap_or_else(|| (0..index.n_partitions()).collect());

        let mut layer_cache = HashMap::with_capacity(selected.len());
        for p in selected {
            let partition = index
                .partition(p)
                .unwrap_or_else(|e| panic!("IndexCache::new: failed to open partition {p}: {e}"));
            let n_layer = partition.n_layers();
            let mut layers = Vec::with_capacity(n_layer);
            for l in 0..n_layer {
                let layer = partition.layer(l).and_then(KmerLayer::open).unwrap_or_else(|e| {
                    panic!(
                        "IndexCache::new: failed to open layer (partition {p}, layer {l}): {e} — partition is not fully built"
                    )
                });
                layers.push(layer);
            }
            layer_cache.insert(p, layers);
        }

        let meta = MetaCache::from_meta(&index.meta());
        IndexCache {
            raw_index: index,
            meta,
            layer_cache,
        }
    }

    /// The index this cache was opened from.
    #[inline]
    pub fn index(&self) -> &KmerIndex {
        &self.raw_index
    }

    /// The index's metadata, snapshotted once at construction — see
    /// [`MetaCache`]. Read-only
    #[inline]
    pub fn meta(&self) -> &MetaCache {
        &self.meta
    }

    /// Number of layers cached for `partition` — `None` if `partition`
    /// isn't part of this cache.
    #[inline]
    pub fn n_layer(&self, partition: usize) -> Option<usize> {
        Some(self.layer_cache.get(&partition)?.len())
    }

    /// Number of partitions cached — not necessarily `raw_index.n_partitions()`
    /// when this cache was built on a subset (see [`new`](Self::new)).
    #[inline]
    pub fn n_partition(&self) -> usize {
        self.layer_cache.len()
    }

    /// Which partitions this cache holds, in no particular order.
    pub fn partitions(&self) -> impl Iterator<Item = usize> + '_ {
        self.layer_cache.keys().copied()
    }

    #[inline]
    pub fn get_layer(&self, partition: usize, layer: usize) -> Option<&KmerLayer> {
        self.layer_cache.get(&partition)?.get(layer)
    }

    pub fn iter(&self) -> impl Iterator<Item = &KmerLayer> {
        self.layer_cache.values().flatten()
    }

    pub fn iter_indexed(&self) -> impl Iterator<Item = (usize, usize, &KmerLayer)> {
        self.layer_cache.iter().flat_map(|(&p, layers)| {
            layers
                .iter()
                .enumerate()
                .map(move |(l, layer)| (p, l, layer))
        })
    }

    /// Iterate over the unitigs of every cached layer, chained — the
    /// partition/index level of `KmerLayer::iter_unitigs`'s layer/partition/
    /// index chain, covering whichever scope this cache was opened with
    /// (one partition, several, or the whole index — see [`new`](Self::new)).
    /// No extra opening: every layer here is already open, this only
    /// streams what [`iter`](Self::iter) already gives out.
    pub fn iter_unitigs(&self) -> impl Iterator<Item = obikseq::Unitig> + '_ {
        self.iter().flat_map(|l| l.iter_unitigs())
    }

    #[inline]
    pub fn hash(&self, partition: usize, layer: usize, kmer: CanonicalKmer) -> Option<usize> {
        Some(self.get_layer(partition, layer)?.hash(kmer))
    }

    pub fn n_kmer(&self, partition: usize, layer: usize) -> Option<usize> {
        Some(self.get_layer(partition, layer)?.n())
    }

    /// Membership-checked lookup in one `(partition, layer)` — `None` if
    /// that layer doesn't carry `kmer`. Unlike [`hash`](Self::hash), this
    /// goes through `KmerLayer::find_slot`'s evidence check, not a blind
    /// MPHF lookup.
    #[inline]
    pub fn find_in_layer(
        &self,
        partition: usize,
        layer: usize,
        kmer: CanonicalKmer,
    ) -> Option<usize> {
        self.get_layer(partition, layer)?.find_slot(kmer)
    }

    /// Membership-checked lookup across every cached layer of one
    /// partition — first layer that carries `kmer` wins. `None` if
    /// `partition` isn't cached or no layer of it carries `kmer`.
    pub fn find_in_partition(
        &self,
        partition: usize,
        kmer: CanonicalKmer,
    ) -> Option<(usize, usize, usize)> {
        let n_layer = self.n_layer(partition)?;
        for l in 0..n_layer {
            if let Some(slot) = self.find_in_layer(partition, l, kmer) {
                return Some((partition, l, slot));
            }
        }
        None
    }

    /// Membership-checked lookup — routes `kmer` straight to its one
    /// possible partition (`CanonicalKmer::partition`, the same
    /// minimizer-hash rule every k-mer was scattered by when the index was
    /// built) instead of scanning every cached partition: a k-mer can only
    /// ever live in the partition it hashes to, so anywhere else is not
    /// worth even asking. Uses `self.index().n_partitions()` — the index's
    /// real total, not [`n_partition`](Self::n_partition) (this cache's own
    /// count, which can be a subset — see [`new`](Self::new)). `None` if
    /// that partition isn't cached or doesn't carry `kmer`.
    pub fn find(&self, kmer: CanonicalKmer) -> Option<(usize, usize, usize)> {
        let partition = kmer.partition(self.raw_index.n_partitions());
        self.find_in_partition(partition, kmer)
    }
}