semtree-store 0.5.0

use std::collections::HashMap;
use std::path::Path;
use std::sync::RwLock;

use async_trait::async_trait;
use semtree_embed::Embedding;
use usearch::{Index, IndexOptions, MetricKind, ScalarKind};

use crate::{Hit, Metric, StoreError, VectorStore};

/// Vectors a fresh index makes room for before it has to grow.
const INITIAL_CAPACITY: usize = 1024;

pub struct UsearchStore {
    dimensions: usize,
    metric: Metric,
    /// Behind a lock because `load` swaps the whole index in place, and a store
    /// is normally shared as `Arc<dyn VectorStore>` by the time it is reloaded.
    index: RwLock<Index>,
    /// maps usearch integer key → chunk id string
    id_map: RwLock<HashMap<u64, String>>,
    next_key: RwLock<u64>,
}

impl UsearchStore {
    /// A store using cosine similarity, the default for normalized text embeddings.
    pub fn new(dimensions: usize) -> Result<Self, StoreError> {
        Self::with_metric(dimensions, Metric::Cosine)
    }

    pub fn with_metric(dimensions: usize, metric: Metric) -> Result<Self, StoreError> {
        let index = Self::make_index(dimensions, metric)?;
        Ok(Self {
            dimensions,
            metric,
            index: RwLock::new(index),
            id_map: RwLock::new(HashMap::new()),
            next_key: RwLock::new(0),
        })
    }

    fn make_index(dimensions: usize, metric: Metric) -> Result<Index, StoreError> {
        let options = IndexOptions {
            dimensions,
            metric: match metric {
                Metric::Cosine => MetricKind::Cos,
                Metric::Euclidean => MetricKind::L2sq,
                Metric::DotProduct => MetricKind::IP,
            },
            quantization: ScalarKind::F32,
            ..Default::default()
        };
        let index = Index::new(&options).map_err(|e| StoreError::Init(e.to_string()))?;
        index
            .reserve(INITIAL_CAPACITY)
            .map_err(|e| StoreError::Init(e.to_string()))?;
        Ok(index)
    }

    /// Make room for one more vector.
    ///
    /// usearch refuses to add past its reserved capacity, and an index read
    /// back from disk is reserved for exactly what was saved. Without this, the
    /// first insert of an incremental pass over an existing index fails.
    fn reserve_one_more(&self) -> Result<(), StoreError> {
        if has_room(&self.index.read().unwrap()) {
            return Ok(());
        }

        // Growing reallocates, so it takes the write lock and locks out the
        // concurrent `add` calls that otherwise run under a read lock. Re-check
        // once we hold it: another writer may have grown the index meanwhile.
        //
        // `reserve` takes `&self`, so this reads as a read-only use of a write
        // lock; the exclusion is the whole point, not an oversight.
        #[allow(clippy::readonly_write_lock)]
        let index = self.index.write().unwrap();
        if has_room(&index) {
            return Ok(());
        }

        // Doubling keeps indexing a large tree from reserving once per file.
        let target = index.capacity().saturating_mul(2).max(INITIAL_CAPACITY);
        index
            .reserve(target)
            .map_err(|e| StoreError::Insert(e.to_string()))
    }

    /// Turn a usearch distance into a higher-is-better similarity score. Cosine
    /// and dot-product distances are `1 - similarity`; L2² has no upper bound,
    /// so it is folded into `(0, 1]` monotonically.
    fn distance_to_score(&self, distance: f32) -> f32 {
        match self.metric {
            Metric::Cosine | Metric::DotProduct => 1.0 - distance,
            Metric::Euclidean => 1.0 / (1.0 + distance),
        }
    }
}

#[async_trait]
impl VectorStore for UsearchStore {
    async fn insert(&self, id: &str, embedding: &Embedding) -> Result<(), StoreError> {
        self.reserve_one_more()?;

        let key = {
            let mut k = self.next_key.write().unwrap();
            let current = *k;
            *k += 1;
            current
        };
        self.index
            .read()
            .unwrap()
            .add(key, embedding)
            .map_err(|e| StoreError::Insert(e.to_string()))?;
        self.id_map.write().unwrap().insert(key, id.to_string());
        Ok(())
    }

    async fn search(&self, query: &Embedding, top_k: usize) -> Result<Vec<Hit>, StoreError> {
        let index = self.index.read().unwrap();
        if index.size() == 0 {
            return Ok(vec![]);
        }
        let results = index
            .search(query, top_k)
            .map_err(|e| StoreError::Search(e.to_string()))?;
        drop(index);

        let map = self.id_map.read().unwrap();
        let hits = results
            .keys
            .iter()
            .zip(results.distances.iter())
            .filter_map(|(key, dist)| {
                map.get(key).map(|id| Hit {
                    id: id.clone(),
                    score: self.distance_to_score(*dist),
                })
            })
            .collect();
        Ok(hits)
    }

    async fn delete(&self, id: &str) -> Result<(), StoreError> {
        let map = self.id_map.read().unwrap();
        if let Some((&key, _)) = map.iter().find(|(_, v)| v.as_str() == id) {
            drop(map);
            self.index
                .read()
                .unwrap()
                .remove(key)
                .map_err(|e| StoreError::Insert(e.to_string()))?;
            self.id_map.write().unwrap().remove(&key);
        }
        Ok(())
    }

    async fn clear(&self) -> Result<(), StoreError> {
        let index = Self::make_index(self.dimensions, self.metric)?;
        *self.index.write().unwrap() = index;
        self.id_map.write().unwrap().clear();
        *self.next_key.write().unwrap() = 0;
        Ok(())
    }

    fn save(&self, path: &Path) -> Result<(), StoreError> {
        // save vector index
        let index_path = path.join("index.usearch");
        self.index
            .read()
            .unwrap()
            .save(index_path.to_str().unwrap())
            .map_err(|e| StoreError::Init(e.to_string()))?;

        // save id_map + next_key as JSON
        let meta = serde_json::json!({
            "dimensions": self.dimensions,
            "metric": self.metric,
            "next_key": *self.next_key.read().unwrap(),
            "id_map": self.id_map.read().unwrap().iter()
                .map(|(k, v)| (k.to_string(), v.clone()))
                .collect::<HashMap<String, String>>(),
        });
        let meta_path = path.join("meta.json");
        std::fs::write(meta_path, serde_json::to_string(&meta).unwrap())
            .map_err(|e| StoreError::Init(e.to_string()))?;

        Ok(())
    }

    fn load(&self, path: &Path) -> Result<(), StoreError> {
        let meta_path = path.join("meta.json");
        let raw =
            std::fs::read_to_string(meta_path).map_err(|e| StoreError::Init(e.to_string()))?;
        let meta: serde_json::Value =
            serde_json::from_str(&raw).map_err(|e| StoreError::Init(e.to_string()))?;

        // Reject a reopen under a different metric: the persisted vectors were
        // ranked one way and cannot be re-ranked another without a rebuild.
        if let Some(saved) = meta.get("metric")
            && let Ok(saved) = serde_json::from_value::<Metric>(saved.clone())
            && saved != self.metric
        {
            return Err(StoreError::Init(format!(
                "index metric mismatch: on disk {saved}, requested {}",
                self.metric
            )));
        }

        let next_key: u64 = meta["next_key"].as_u64().unwrap_or(0);
        let id_map: HashMap<u64, String> = meta["id_map"]
            .as_object()
            .unwrap_or(&serde_json::Map::new())
            .iter()
            .filter_map(|(k, v)| Some((k.parse::<u64>().ok()?, v.as_str()?.to_string())))
            .collect();

        // Build the replacement index and read it off disk before touching any
        // live state, so a failed load leaves the store exactly as it was.
        let index = Self::make_index(self.dimensions, self.metric)?;
        let index_path = path.join("index.usearch");
        index
            .load(index_path.to_str().unwrap())
            .map_err(|e| StoreError::Init(e.to_string()))?;

        *self.index.write().unwrap() = index;
        *self.next_key.write().unwrap() = next_key;
        *self.id_map.write().unwrap() = id_map;

        Ok(())
    }

    fn len(&self) -> usize {
        self.index.read().unwrap().size()
    }

    fn metric(&self) -> Metric {
        self.metric
    }
}

fn has_room(index: &Index) -> bool {
    index.size() < index.capacity()
}

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

    #[test]
    fn score_is_higher_for_nearer_vectors() {
        // Cosine/dot: score = 1 - distance. Euclidean: closer (smaller L2²) must
        // still rank higher, i.e. a larger score.
        let cos = UsearchStore::with_metric(4, Metric::Cosine).unwrap();
        assert!(cos.distance_to_score(0.1) > cos.distance_to_score(0.9));

        let l2 = UsearchStore::with_metric(4, Metric::Euclidean).unwrap();
        assert!(l2.distance_to_score(0.1) > l2.distance_to_score(5.0));
        assert!(l2.distance_to_score(0.0) <= 1.0);
    }

    #[test]
    fn reload_rejects_metric_mismatch() {
        let dir = std::env::temp_dir().join("semtree_store_metric_test");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        let saved = UsearchStore::with_metric(4, Metric::Cosine).unwrap();
        saved.save(&dir).unwrap();

        // Same metric reloads fine...
        let same = UsearchStore::with_metric(4, Metric::Cosine).unwrap();
        assert!(same.load(&dir).is_ok());

        // ...a different metric is refused rather than silently mis-ranking.
        let other = UsearchStore::with_metric(4, Metric::Euclidean).unwrap();
        assert!(other.load(&dir).is_err());

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn inserts_keep_working_past_the_initial_capacity() {
        let store = UsearchStore::with_metric(4, Metric::Cosine).unwrap();
        for i in 0..(INITIAL_CAPACITY + 10) {
            store
                .insert(&format!("id{i}"), &vec![i as f32, 0.0, 0.0, 1.0])
                .await
                .expect("insert grows the index instead of failing");
        }
        assert_eq!(store.len(), INITIAL_CAPACITY + 10);
    }

    #[tokio::test]
    async fn a_reloaded_index_can_still_be_added_to() {
        // The incremental path: read an index back, then index a new file into
        // it. A store reloaded at exactly its saved capacity has no room left,
        // so this used to fail on the first insert.
        let dir = std::env::temp_dir().join("semtree_store_reload_insert_test");
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        let store = UsearchStore::with_metric(4, Metric::Cosine).unwrap();
        store.insert("a", &vec![1.0, 0.0, 0.0, 0.0]).await.unwrap();
        store.save(&dir).unwrap();

        let reopened = UsearchStore::with_metric(4, Metric::Cosine).unwrap();
        reopened.load(&dir).unwrap();
        assert_eq!(reopened.len(), 1, "the saved vector came back");

        reopened
            .insert("b", &vec![0.0, 1.0, 0.0, 0.0])
            .await
            .unwrap();
        assert_eq!(reopened.len(), 2, "and the new one was added on top");

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[tokio::test]
    async fn clear_empties_the_store() {
        let store = UsearchStore::with_metric(4, Metric::Cosine).unwrap();
        store.insert("a", &vec![1.0, 0.0, 0.0, 0.0]).await.unwrap();
        assert_eq!(store.len(), 1);

        store.clear().await.unwrap();
        assert_eq!(store.len(), 0, "a full rebuild starts from nothing");

        store.insert("b", &vec![0.0, 1.0, 0.0, 0.0]).await.unwrap();
        assert_eq!(store.len(), 1, "and can be filled again");
    }
}