linkrs-storage 0.1.0

use std::collections::{HashMap, HashSet};
use std::sync::Arc;

use crate::index::helpers::{effective_index_values, merged_included_columns, vertex_entity_ref};
use crate::index::key_codec::key_builder::normalize_int_value;
use crate::index::key_codec::key_types::SecondaryIndexKey;
use crate::index::key_codec::{KeyBuilder, KeyParser};
use crate::index::traits::VertexIndexOps;
use crate::index::types::{IndexIdentity, IndexRecord};
use crate::index::{shard_runtime::IndexMaps, IndexDataManagerImpl};
use linkrs_core::types::{Index, IndexType, Timestamp, MAX_TIMESTAMP};
use linkrs_core::value::ordered_codec::OrderedCodec;
use linkrs_core::{StorageError, StorageResult, Value};

fn add_entry(
    per_shard: &mut HashMap<u32, IndexMaps>,
    shard_id: u32,
    fwd_key: SecondaryIndexKey,
    rev_key: SecondaryIndexKey,
    entry: IndexRecord,
) {
    let (ref mut fwd_map, ref mut rev_map) = per_shard.entry(shard_id).or_default();
    fwd_map.insert(fwd_key, entry.clone());
    rev_map.insert(rev_key, entry);
}

impl VertexIndexOps for IndexDataManagerImpl {
    fn update_vertex_indexes_mvcc(
        &self,
        space_id: u64,
        vertex_id: &Value,
        index_name: &str,
        props: &[(Arc<str>, Value)],
        write_ts: Timestamp,
    ) -> Result<(), StorageError> {
        let Some(index_id) = self.index_alias(space_id, index_name) else {
            return Ok(());
        };
        let identity = IndexIdentity { space_id, index_id };
        let runtime = self.runtime(space_id, index_id)?;

        let delta = {
            let catalog = self
                .manifest_catalog(space_id, index_id)
                .ok_or_else(|| StorageError::not_found("Index manifest catalog is unavailable"))?;
            let manifest = catalog.acquire();
            let index_definition = self.index_definitions.read().get(&identity).cloned();

            let covering = index_definition.as_ref().is_some_and(|idx| idx.covering);
            let new_values = effective_index_values(index_definition.as_ref(), props, Vec::new());

            let latest_gen = runtime.generation(manifest.manifest().generation);
            let reverse_prefix =
                KeyBuilder::build_vertex_reverse_key_v2(space_id, vertex_id, index_name)?;
            let reverse_end = KeyBuilder::build_range_end(&reverse_prefix);
            let mut existing_values = Vec::new();
            let mut existing_encoded: HashSet<Vec<u8>> = HashSet::new();
            let mut existing_columns = Vec::new();
            let mut covering_populated = false;
            if let Some(latest_gen) = latest_gen {
                for shard in latest_gen.shards() {
                    if !shard.reverse_may_have_range(&reverse_prefix.0, &reverse_end.0) {
                        continue;
                    }
                    for (_suffix, record) in shard.reverse_range_suffix_visible(
                        &reverse_prefix.0,
                        &reverse_end.0,
                        write_ts,
                    ) {
                        if let Ok(encoded) = KeyParser::extract_value_from_reverse_suffix(&_suffix)
                        {
                            if existing_encoded.insert(encoded.clone()) {
                                if let Ok(value) = OrderedCodec::new().decode(&encoded) {
                                    existing_values.push(normalize_int_value(&value));
                                }
                            }
                        }
                        if covering && !covering_populated {
                            if let Some(cols) = &record.included_columns {
                                existing_columns.clone_from(cols);
                                covering_populated = true;
                            }
                        }
                    }
                }
            }
            // also account for entries still awaiting generation publication.
            let pending_guard = self.pending_deltas.lock();
            if let Some(pending) = pending_guard.get(&identity) {
                let mut scan = crate::index::index_data_manager::PendingExistingScan {
                    existing_values: &mut existing_values,
                    existing_encoded: &mut existing_encoded,
                    existing_columns: &mut existing_columns,
                    covering_populated: &mut covering_populated,
                };
                crate::index::index_data_manager::merge_pending_existing_values(
                    pending,
                    &reverse_prefix.0,
                    &reverse_end.0,
                    write_ts,
                    false,
                    &mut scan,
                );
            }
            drop(pending_guard);

            let values_to_remove: Vec<&Value> = existing_values
                .iter()
                .filter(|v| !new_values.contains(v))
                .collect();
            let values_to_add: Vec<&Value> = new_values
                .iter()
                .filter(|v| !existing_values.contains(v))
                .collect();
            let values_to_update: Vec<&Value> = new_values
                .iter()
                .filter(|v| existing_values.contains(v))
                .collect();

            let included_columns = if covering {
                merged_included_columns(index_definition.as_ref(), existing_columns, props)
            } else {
                Vec::new()
            };

            let mut per_shard: HashMap<u32, IndexMaps> = HashMap::new();

            let route = |forward: &[u8]| -> StorageResult<u32> {
                manifest
                    .manifest()
                    .route_key(forward)
                    .map(|s| s.shard_id)
                    .ok_or_else(|| {
                        StorageError::invalid_operation(
                            "Index manifest does not cover the ordered key",
                        )
                    })
            };

            for value in &values_to_update {
                let forward =
                    KeyBuilder::build_vertex_index_key(space_id, index_name, value, vertex_id)?;
                let reverse = KeyBuilder::build_vertex_reverse_key_with_value(
                    space_id, vertex_id, index_name, value,
                )?;
                let shard_id = route(&forward.0)?;
                let mut entry = if covering {
                    IndexRecord::new_with_columns(write_ts, included_columns.clone())
                } else {
                    IndexRecord::new(write_ts)
                };
                if let Some(entity) = vertex_entity_ref(vertex_id) {
                    entry = entry.with_entity_ref(entity);
                }
                add_entry(&mut per_shard, shard_id, forward.0, reverse.0, entry);
            }

            for value in &values_to_remove {
                let forward =
                    KeyBuilder::build_vertex_index_key(space_id, index_name, value, vertex_id)?;
                let reverse = KeyBuilder::build_vertex_reverse_key_with_value(
                    space_id, vertex_id, index_name, value,
                )?;
                let shard_id = route(&forward.0)?;
                let mut entry = if covering {
                    IndexRecord::new_with_columns(write_ts, included_columns.clone())
                } else {
                    IndexRecord::new(write_ts)
                };
                entry.mark_deleted(write_ts);
                if let Some(entity) = vertex_entity_ref(vertex_id) {
                    entry = entry.with_entity_ref(entity);
                }
                add_entry(&mut per_shard, shard_id, forward.0, reverse.0, entry);
            }

            for value in &values_to_add {
                let forward =
                    KeyBuilder::build_vertex_index_key(space_id, index_name, value, vertex_id)?;
                let reverse = KeyBuilder::build_vertex_reverse_key_with_value(
                    space_id, vertex_id, index_name, value,
                )?;
                let shard_id = route(&forward.0)?;
                let mut entry = if covering {
                    IndexRecord::new_with_columns(write_ts, included_columns.clone())
                } else {
                    IndexRecord::new(write_ts)
                };
                if let Some(entity) = vertex_entity_ref(vertex_id) {
                    entry = entry.with_entity_ref(entity);
                }
                add_entry(&mut per_shard, shard_id, forward.0, reverse.0, entry);
            }

            per_shard
        };

        if !delta.is_empty() {
            self.accumulate_delta(identity, delta, write_ts)?;
        }
        Ok(())
    }

    fn delete_vertex_indexes_mvcc(
        &self,
        space_id: u64,
        vertex_id: &Value,
        index_names: &[String],
        write_ts: Timestamp,
    ) -> Result<(), StorageError> {
        for index_name in index_names {
            self.clear_vertex_entity(space_id, vertex_id, index_name, write_ts)?;
        }
        Ok(())
    }

    fn lookup_tag_index_mvcc(
        &self,
        space_id: u64,
        index: &Index,
        value: &Value,
        read_ts: Timestamp,
    ) -> Result<Vec<Value>, StorageError> {
        self.lookup_tag_index_internal(space_id, index, value, read_ts, false)
    }

    fn lookup_tag_index_pending_aware_mvcc(
        &self,
        space_id: u64,
        index: &Index,
        value: &Value,
        read_ts: Timestamp,
    ) -> Result<Vec<Value>, StorageError> {
        self.lookup_tag_index_internal(space_id, index, value, read_ts, true)
    }

    fn clear_tag_index(&self, space_id: u64, index_name: &str) -> Result<(), StorageError> {
        let Some(index_id) = self.index_alias(space_id, index_name) else {
            return Ok(());
        };
        self.clear_index(
            index_id,
            space_id,
            index_name,
            IndexType::TagIndex,
            MAX_TIMESTAMP,
        )
    }
}

impl IndexDataManagerImpl {
    /// Shared forward-chain lookup. When `merge_pending` is set, pending
    /// (unpublished) deltas are treated as the newest overlay and merged
    /// in-memory; otherwise any pending delta is published first so the
    /// generation chain scan observes it.
    fn lookup_tag_index_internal(
        &self,
        space_id: u64,
        index: &Index,
        value: &Value,
        read_ts: Timestamp,
        merge_pending: bool,
    ) -> Result<Vec<Value>, StorageError> {
        let Some(index_id) = self.index_alias(space_id, &index.name) else {
            return Ok(Vec::new());
        };
        let identity = IndexIdentity { space_id, index_id };
        if !merge_pending {
            // make any pending (unpublished) writes visible to this read.
            self.publish_pending_delta(identity)?;
        }
        let runtime = self.runtime(space_id, index_id)?;
        let catalog = self
            .manifest_catalog(space_id, index_id)
            .ok_or_else(|| StorageError::not_found(format!("Index {index_id} has no manifest")))?;
        let handle = catalog.acquire();
        let manifest = handle.manifest();
        let chain = runtime.generation_chain_until(manifest.generation)?;
        let prefix = KeyBuilder::build_vertex_index_value_prefix(space_id, &index.name, value)?;
        let end = KeyBuilder::build_range_end(&prefix);
        let mut seen = std::collections::HashSet::new();
        let mut tombstoned = std::collections::HashSet::new();
        let mut results = Vec::new();

        // pending holds the newest writes. Scan it first so its per-entry
        // state wins over older published generations, mirroring the
        // newest-first semantics of the published generation chain.
        if merge_pending {
            let pending_guard = self.pending_deltas.lock();
            if let Some(pending) = pending_guard.get(&identity) {
                for (forward, _) in pending.per_shard.values() {
                    for (key, entry) in forward.range((
                        std::ops::Bound::Included(prefix.0.clone()),
                        std::ops::Bound::Excluded(end.0.clone()),
                    )) {
                        if let Ok(vertex_id) = KeyParser::parse_vertex_id_from_key(key) {
                            if entry.created_ts > read_ts {
                                continue;
                            }
                            if entry.deleted_ts.is_some_and(|d| d <= read_ts) {
                                tombstoned.insert(vertex_id);
                                continue;
                            }
                            if !seen.contains(&vertex_id) {
                                seen.insert(vertex_id.clone());
                                results.push(vertex_id);
                            }
                        }
                    }
                }
            }
            drop(pending_guard);
        }

        for generation in &chain {
            for shard in manifest
                .shards
                .iter()
                .filter_map(|s| generation.shard(s.shard_id))
            {
                for (key, entry) in shard.forward_range(&prefix.0, &end.0) {
                    if let Ok(vertex_id) = KeyParser::parse_vertex_id_from_key(&key) {
                        if tombstoned.contains(&vertex_id) {
                            continue;
                        }
                        if seen.contains(&vertex_id) {
                            continue;
                        }
                        if entry.created_ts > read_ts {
                            continue;
                        }
                        if entry.deleted_ts.is_some_and(|d| d <= read_ts) {
                            tombstoned.insert(vertex_id);
                            continue;
                        }
                        seen.insert(vertex_id.clone());
                        results.push(vertex_id);
                    }
                }
            }
        }
        Ok(results)
    }
}