holographic-memory 0.6.4

// Copyright 2024-2026 WritersLogic Contributors
// SPDX-License-Identifier: Apache-2.0

use std::collections::BinaryHeap;
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};

use crate::core::sync::RwLock;
use anyhow::Result;
use rayon::prelude::*;

use crate::core::entangled::EntangledHVec;
use crate::core::hash::{fx64, FastMap};
use crate::core::hopfield::{hopfield_query, HopfieldConfig};
use crate::core::index::inverted::{Accumulator, SparseInvertedIndex};
use crate::core::ivf::IVFIndex;
use crate::core::nsg::NSGIndex;
use crate::core::types::{QueryExplanation, RetrievalResult};

use super::router;

/// Per-shard state: vectors, registry, and indices.
/// Lock ordering within a shard: vectors -> ivf -> nsg.
pub(crate) struct Shard {
    pub vectors: RwLock<FastMap<String, EntangledHVec>>,
    pub registry: RwLock<Vec<String>>,
    slots: RwLock<FastMap<String, u32>>,
    free_slots: crate::core::sync::Mutex<Vec<u32>>,
    pub inverted: RwLock<SparseInvertedIndex>,
    pub accumulator: crate::core::sync::Mutex<Accumulator>,
    pub nsg: RwLock<Option<NSGIndex>>,
    pub ivf: RwLock<Option<IVFIndex>>,
    pub vector_count: AtomicU64,
}

impl Shard {
    pub fn new(dimensions: usize) -> Self {
        let m = (dimensions / 256).max(1);
        Self {
            vectors: RwLock::new(FastMap::default()),
            registry: RwLock::new(Vec::new()),
            slots: RwLock::new(FastMap::default()),
            free_slots: crate::core::sync::Mutex::new(Vec::new()),
            inverted: RwLock::new(SparseInvertedIndex::new(dimensions, m)),
            accumulator: crate::core::sync::Mutex::new(Accumulator::new(1024)),
            nsg: RwLock::new(None),
            ivf: RwLock::new(None),
            vector_count: AtomicU64::new(0),
        }
    }

    pub fn insert(&self, id: String, vector: EntangledHVec, _dimensions: usize) -> Result<()> {
        let mut vectors = self.vectors.write();
        let mut registry = self.registry.write();
        let mut slots = self.slots.write();
        let mut inverted = self.inverted.write();
        let old = vectors.insert(id.clone(), vector.clone());
        let slot = if let Some(&slot) = slots.get(&id) {
            if let Some(old) = &old {
                inverted.remove_doc(slot, old.indices());
            }
            slot
        } else if let Some(slot) = self.free_slots.lock().pop() {
            registry[slot as usize] = id.clone();
            slot
        } else {
            registry.push(id.clone());
            (registry.len() - 1) as u32
        };
        slots.insert(id.clone(), slot);
        inverted.add_doc(slot, vector.indices());
        self.vector_count
            .store(vectors.len() as u64, AtomicOrdering::Release);
        drop(inverted);
        drop(slots);
        drop(registry);
        drop(vectors);
        if old.is_some() {
            *self.ivf.write() = None;
            *self.nsg.write() = None;
        } else {
            let mut ivf = self.ivf.write();
            if let Some(index) = ivf.as_mut() {
                if let Err(error) = index.insert(&id, &vector) {
                    tracing::warn!(%error, "IVF invalidated after failed online insertion");
                    *ivf = None;
                }
            }
            let mut nsg = self.nsg.write();
            if let Some(index) = nsg.as_mut() {
                if let Err(error) = index.insert(&id, &vector) {
                    tracing::warn!(%error, "NSG invalidated after failed online insertion");
                    *nsg = None;
                }
            }
        }
        Ok(())
    }

    pub fn remove(&self, id: &str, _dimensions: usize) -> Result<bool> {
        let mut vectors = self.vectors.write();
        let Some(old) = vectors.remove(id) else {
            return Ok(false);
        };
        let mut registry = self.registry.write();
        if let Some(slot) = self.slots.write().remove(id) {
            self.inverted.write().remove_doc(slot, old.indices());
            registry[slot as usize].clear();
            self.free_slots.lock().push(slot);
        }
        self.vector_count
            .store(vectors.len() as u64, AtomicOrdering::Release);
        *self.ivf.write() = None;
        *self.nsg.write() = None;
        Ok(true)
    }

    pub fn rebuild_inverted_index(&self, dimensions: usize) -> Result<()> {
        let vectors = self.vectors.read();
        let registry = self.registry.read();
        let mut slots = self.slots.write();
        slots.clear();
        self.free_slots.lock().clear();
        let mut inv = self.inverted.write();

        *inv = SparseInvertedIndex::new(dimensions, (dimensions / 256).max(1));

        for (i, id) in registry.iter().enumerate() {
            if let Some(vec) = vectors.get(id) {
                slots.insert(id.clone(), i as u32);
                inv.add_doc(i as u32, &vec.indices);
            }
        }
        inv.finalize();

        let mut acc = self.accumulator.lock();
        *acc = Accumulator::new(registry.len().max(1024));

        Ok(())
    }

    pub fn load_all_vectors(&self) -> (Vec<String>, Vec<EntangledHVec>) {
        let vectors = self.vectors.read();
        let registry = self.registry.read();
        let mut ids = Vec::with_capacity(registry.len());
        let mut out_vectors = Vec::with_capacity(registry.len());

        for id in registry.iter() {
            if let Some(vec) = vectors.get(id) {
                ids.push(id.clone());
                out_vectors.push(vec.clone());
            }
        }

        (ids, out_vectors)
    }

    pub fn query(
        &self,
        query_vec: &EntangledHVec,
        k: u32,
        dimensions: usize,
    ) -> Vec<RetrievalResult> {
        let n = self.vector_count.load(AtomicOrdering::SeqCst) as usize;

        let planner =
            router::QueryPlanner::new(self.nsg_trained(), n > 0, self.ivf_trained(), n, dimensions);
        let plan = planner.plan(query_vec, k);

        let vectors = self.vectors.read();

        match plan.route {
            router::IndexRoute::NSG => {
                if let Some(ref nsg) = *self.nsg.read() {
                    let results: Vec<RetrievalResult> = nsg
                        .query(query_vec, k as usize, plan.ef_search)
                        .into_iter()
                        .filter(|r| vectors.contains_key(&r.id))
                        .collect();
                    if !results.is_empty() {
                        return results;
                    }
                }
            }
            router::IndexRoute::Inverted => {
                let mut acc = self.accumulator.lock();
                let results = self
                    .inverted
                    .read()
                    .query(&query_vec.indices, k as usize, &mut acc);

                let reg = self.registry.read();
                let mapped: Vec<RetrievalResult> = results
                    .into_iter()
                    .filter_map(|r| {
                        let id_str = reg.get(r.slot as usize)?;
                        if vectors.contains_key(id_str) {
                            Some(RetrievalResult {
                                id: id_str.clone(),
                                similarity: r.similarity,
                            })
                        } else {
                            None
                        }
                    })
                    .collect();

                if !mapped.is_empty() {
                    return mapped;
                }
            }
            router::IndexRoute::IVF => {
                if let Some(ref ivf) = *self.ivf.read() {
                    if let Ok(candidates) = ivf.query(query_vec, k as usize, plan.n_probe) {
                        let reranked =
                            Self::rerank_by_exact_similarity(&vectors, query_vec, &candidates);
                        if !reranked.is_empty() {
                            return reranked;
                        }
                    }
                }
            }
            router::IndexRoute::Hopfield => {
                let patterns: Vec<(&str, &EntangledHVec)> =
                    vectors.iter().map(|(id, vec)| (id.as_str(), vec)).collect();
                let config = HopfieldConfig::default();
                let results = hopfield_query(query_vec, &patterns, &config, k as usize);
                if !results.is_empty() {
                    return results;
                }
            }
            router::IndexRoute::BruteForce => {
                return Self::brute_force_scan(&vectors, query_vec, k as usize);
            }
        }

        Self::brute_force_scan(&vectors, query_vec, k as usize)
    }

    fn rerank_by_exact_similarity(
        vectors: &FastMap<String, EntangledHVec>,
        query_vec: &EntangledHVec,
        candidates: &[RetrievalResult],
    ) -> Vec<RetrievalResult> {
        let mut results: Vec<RetrievalResult> = candidates
            .iter()
            .filter_map(|c| {
                let vec = vectors.get(&c.id)?;
                Some(RetrievalResult {
                    id: c.id.clone(),
                    similarity: query_vec.similarity(vec),
                })
            })
            .collect();

        results.sort_unstable_by(|a, b| {
            b.similarity
                .partial_cmp(&a.similarity)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results
    }

    fn brute_force_scan(
        vectors: &FastMap<String, EntangledHVec>,
        query_vec: &EntangledHVec,
        k: usize,
    ) -> Vec<RetrievalResult> {
        if k == 0 || vectors.is_empty() {
            return Vec::new();
        }
        #[derive(PartialEq)]
        struct Candidate<'a> {
            similarity: f64,
            id: &'a str,
        }
        impl<'a> Eq for Candidate<'a> {}
        impl<'a> PartialOrd for Candidate<'a> {
            fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
                Some(self.cmp(other))
            }
        }
        impl<'a> Ord for Candidate<'a> {
            fn cmp(&self, other: &Self) -> std::cmp::Ordering {
                // Reversed on similarity so the max-heap behaves as a min-heap.
                other
                    .similarity
                    .partial_cmp(&self.similarity)
                    .unwrap_or(std::cmp::Ordering::Equal)
            }
        }

        let mut heap: BinaryHeap<Candidate<'_>> = BinaryHeap::with_capacity(k.min(vectors.len()));

        for (id, vec) in vectors.iter() {
            let sim = query_vec.similarity(vec);

            if heap.len() < k {
                heap.push(Candidate {
                    similarity: sim,
                    id: id.as_str(),
                });
            } else if let Some(top) = heap.peek() {
                if sim > top.similarity {
                    heap.pop();
                    heap.push(Candidate {
                        similarity: sim,
                        id: id.as_str(),
                    });
                }
            }
        }

        // The reversed `Ord` makes ascending order highest-similarity first.
        heap.into_sorted_vec()
            .into_iter()
            .map(|c| RetrievalResult {
                similarity: c.similarity,
                id: c.id.to_string(),
            })
            .collect()
    }

    pub fn nsg_trained(&self) -> bool {
        self.nsg.read().as_ref().is_some_and(|nsg| nsg.is_trained())
    }

    pub fn ivf_trained(&self) -> bool {
        self.ivf.read().as_ref().is_some_and(|ivf| ivf.is_trained())
    }

    pub fn count(&self) -> u64 {
        self.vector_count.load(AtomicOrdering::SeqCst)
    }
}

/// Multi-shard coordinator. Distributes vectors by consistent hash of ID.
pub(crate) struct ShardManager {
    pub shards: Vec<Shard>,
}

impl ShardManager {
    pub fn new(shard_count: usize, dimensions: usize) -> Self {
        assert!(shard_count >= 2, "ShardManager requires at least 2 shards");
        let shards = (0..shard_count).map(|_| Shard::new(dimensions)).collect();
        Self { shards }
    }

    pub fn shard_for(&self, id: &str) -> usize {
        (fx64(id.as_bytes()) as usize) % self.shards.len()
    }

    pub fn shard(&self, id: &str) -> &Shard {
        &self.shards[self.shard_for(id)]
    }

    pub fn query(
        &self,
        query_vec: &EntangledHVec,
        k: u32,
        dimensions: usize,
    ) -> Vec<RetrievalResult> {
        let per_shard: Vec<Vec<RetrievalResult>> = self
            .shards
            .par_iter()
            .map(|shard| shard.query(query_vec, k, dimensions))
            .collect();

        let candidates = per_shard
            .iter()
            .map(Vec::len)
            .fold(0usize, usize::saturating_add);
        let mut heap: BinaryHeap<RetrievalResult> =
            BinaryHeap::with_capacity((k as usize).saturating_add(1).min(candidates));
        for results in per_shard {
            for r in results {
                heap.push(r);
                if heap.len() > k as usize {
                    heap.pop();
                }
            }
        }
        // `RetrievalResult`'s reversed `Ord` makes ascending order highest-similarity first.
        heap.into_sorted_vec()
    }

    pub fn total_count(&self) -> u64 {
        self.shards.iter().map(|s| s.count()).sum()
    }

    pub fn nsg_trained(&self) -> bool {
        self.shards.iter().all(|s| s.nsg_trained())
    }

    pub fn ivf_trained(&self) -> bool {
        self.shards.iter().all(|s| s.ivf_trained())
    }
}

/// Wrapper that transparently handles single vs multi-shard operation.
pub(crate) enum ShardSet {
    Single(Box<Shard>),
    Multi(ShardManager),
}

impl ShardSet {
    pub fn insert(&self, id: String, vector: EntangledHVec, dimensions: usize) -> Result<()> {
        match self {
            ShardSet::Single(shard) => shard.insert(id, vector, dimensions),
            ShardSet::Multi(mgr) => mgr.shard(&id).insert(id, vector, dimensions),
        }
    }

    pub fn remove(&self, id: &str, dimensions: usize) -> Result<bool> {
        match self {
            ShardSet::Single(shard) => shard.remove(id, dimensions),
            ShardSet::Multi(mgr) => mgr.shard(id).remove(id, dimensions),
        }
    }

    pub fn query(
        &self,
        query_vec: &EntangledHVec,
        k: u32,
        dimensions: usize,
    ) -> Vec<RetrievalResult> {
        match self {
            ShardSet::Single(shard) => shard.query(query_vec, k, dimensions),
            ShardSet::Multi(mgr) => mgr.query(query_vec, k, dimensions),
        }
    }

    pub fn explain_query(
        &self,
        query_vec: &EntangledHVec,
        k: u32,
        dimensions: usize,
    ) -> QueryExplanation {
        let count = self.count() as usize;
        let nsg = self.nsg_trained();
        let ivf = self.ivf_trained();
        let plan =
            router::QueryPlanner::new(nsg, count > 0, ivf, count, dimensions).plan(query_vec, k);
        QueryExplanation {
            route: plan.route.as_str().to_string(),
            vector_count: count.min(u32::MAX as usize) as u32,
            query_nonzero_dimensions: query_vec.indices.len().min(u32::MAX as usize) as u32,
            requested_results: k,
            ef_search: plan.ef_search as u32,
            n_probe: plan.n_probe as u32,
            nsg_available: nsg,
            ivf_available: ivf,
            rationale: plan.rationale.to_string(),
        }
    }

    pub fn count(&self) -> u64 {
        match self {
            ShardSet::Single(shard) => shard.count(),
            ShardSet::Multi(mgr) => mgr.total_count(),
        }
    }

    pub fn nsg_trained(&self) -> bool {
        match self {
            ShardSet::Single(shard) => shard.nsg_trained(),
            ShardSet::Multi(mgr) => mgr.nsg_trained(),
        }
    }

    pub fn ivf_trained(&self) -> bool {
        match self {
            ShardSet::Single(shard) => shard.ivf_trained(),
            ShardSet::Multi(mgr) => mgr.ivf_trained(),
        }
    }

    pub fn get_vector(&self, id: &str) -> Option<EntangledHVec> {
        match self {
            ShardSet::Single(shard) => shard.vectors.read().get(id).cloned(),
            ShardSet::Multi(mgr) => mgr.shard(id).vectors.read().get(id).cloned(),
        }
    }

    pub fn for_each_shard<F: FnMut(&Shard)>(&self, mut f: F) {
        match self {
            ShardSet::Single(shard) => f(shard),
            ShardSet::Multi(mgr) => {
                for shard in &mgr.shards {
                    f(shard);
                }
            }
        }
    }

    pub fn collect_all_patterns(&self) -> Vec<(String, EntangledHVec)> {
        let mut patterns = Vec::new();
        self.for_each_shard(|shard| {
            let vectors = shard.vectors.read();
            patterns.extend(vectors.iter().map(|(id, vec)| (id.clone(), vec.clone())));
        });
        patterns
    }

    /// Calls `f` with every stored pattern borrowed in place, in `collect_all_patterns`
    /// order (shard order, then each shard's map order), with every shard's vectors read
    /// lock held for the call.
    pub fn with_patterns<R>(&self, f: impl FnOnce(&[(&str, &EntangledHVec)]) -> R) -> R {
        let guards: Vec<_> = match self {
            ShardSet::Single(shard) => vec![shard.vectors.read()],
            ShardSet::Multi(mgr) => mgr.shards.iter().map(|s| s.vectors.read()).collect(),
        };
        let mut patterns = Vec::with_capacity(guards.iter().map(|g| g.len()).sum());
        for vectors in &guards {
            patterns.extend(vectors.iter().map(|(id, vec)| (id.as_str(), vec)));
        }
        f(&patterns)
    }

    pub fn try_for_each_shard<F: FnMut(&Shard) -> Result<()>>(&self, mut f: F) -> Result<()> {
        match self {
            ShardSet::Single(shard) => f(shard),
            ShardSet::Multi(mgr) => {
                for shard in &mgr.shards {
                    f(shard)?;
                }
                Ok(())
            }
        }
    }

    /// Number of shards backing this set (1 for the single-shard case).
    pub fn shard_count(&self) -> usize {
        match self {
            ShardSet::Single(_) => 1,
            ShardSet::Multi(mgr) => mgr.shards.len(),
        }
    }

    /// Like `try_for_each_shard`, but also passes each shard's index so callers
    /// can address per-shard on-disk artifacts.
    pub fn try_for_each_shard_indexed<F: FnMut(usize, &Shard) -> Result<()>>(
        &self,
        mut f: F,
    ) -> Result<()> {
        match self {
            ShardSet::Single(shard) => f(0, shard),
            ShardSet::Multi(mgr) => {
                for (i, shard) in mgr.shards.iter().enumerate() {
                    f(i, shard)?;
                }
                Ok(())
            }
        }
    }
}

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

    fn fingerprint(results: &[RetrievalResult]) -> Vec<(String, u64)> {
        results
            .iter()
            .map(|r| (r.id.clone(), r.similarity.to_bits()))
            .collect()
    }

    /// Borrowing the stored patterns ranks exactly as the owned snapshot does, ties included:
    /// same ids, same order, same weight bits, for every entmax family and the iterative rule.
    #[test]
    fn borrowed_patterns_rank_as_the_owned_snapshot() -> Result<()> {
        let dim = 2048;
        for set in [
            ShardSet::Single(Box::new(Shard::new(dim))),
            ShardSet::Multi(ShardManager::new(4, dim)),
        ] {
            for i in 0..300u64 {
                // Every fifth pattern repeats an earlier vector under a new id, so weights tie.
                let seed = if i % 5 == 4 { i - 4 } else { i };
                set.insert(
                    format!("p{i}"),
                    EntangledHVec::new_deterministic(dim, seed),
                    dim,
                )?;
            }
            let owned = set.collect_all_patterns();
            for (alpha, beta, max_iter) in [
                (1.0, 40.0, 1),
                (1.5, 60.0, 1),
                (2.0, 100.0, 1),
                (2.0, 100.0, 3),
                (1.5, 30.0, 4),
            ] {
                let config = HopfieldConfig {
                    beta,
                    alpha,
                    max_iter,
                };
                for seed in [0u64, 8, 9, 151, 9_999] {
                    let query = EntangledHVec::new_deterministic(dim, seed);
                    let expected = crate::core::hopfield::hopfield_query_iterative(
                        &query, &owned, &config, 50,
                    );
                    let borrowed = set.with_patterns(|patterns| {
                        crate::core::hopfield::hopfield_query_iterative(
                            &query, patterns, &config, 50,
                        )
                    });
                    assert!(!expected.is_empty());
                    assert_eq!(fingerprint(&expected), fingerprint(&borrowed));
                }
            }
        }
        Ok(())
    }
}