linkrs-storage 0.1.0

use crate::vertex::VertexRecord;
use linkrs_core::types::{LabelId, Timestamp, VertexId};
use std::sync::atomic::Ordering;

use super::super::GraphStorageContext;
use super::resolve::ExternalRef;

impl GraphStorageContext {
    /// Pre-allocate capacity for `additional` more vertices in the given label's table.
    /// Call before batch inserts to avoid repeated hash rehashing.
    pub fn reserve_vertex_capacity(&self, label: LabelId, additional: usize) {
        let _ = self
            .persistent
            .data_store
            .with_vertex_tables_mut(|vertex_tables| {
                if let Some(table) = vertex_tables.get(&label) {
                    table.reserve_id_capacity(additional);
                }
                Ok(())
            });
    }

    pub fn get_vertex(
        &self,
        label: LabelId,
        external_id: &str,
        ts: Timestamp,
    ) -> Option<VertexRecord> {
        if !self.persistent.is_open.load(Ordering::Acquire) {
            return None;
        }
        self.get_full_record(label, ExternalRef::Str(external_id), ts)
    }

    /// Fetch a vertex restricted to the given property projection, skipping
    /// the full-record cache so partial results never replace cached vertices.
    ///
    /// Pending-aware through [`Self::resolve_projected_on_table`].
    pub fn get_vertex_projected(
        &self,
        label: LabelId,
        external_id: &str,
        projection: &[std::sync::Arc<str>],
        ts: Timestamp,
    ) -> Option<VertexRecord> {
        if !self.persistent.is_open.load(Ordering::Acquire) {
            return None;
        }

        // ID-index hits are mapping-revalidated; the record itself is
        // pending-aware through `read_record`. A key with no internal id can
        // still be a staged insert of this transaction.
        let internal_id = self.persistent.data_store.with_vertex_tables(|tables| {
            let table = tables.get(&label)?;
            self.resolve_internal_id_rechecked(table, label, ExternalRef::Str(external_id), ts)
        });
        match internal_id {
            Some(internal_id) => self.read_record(label, internal_id, Some(projection), ts),
            None => self.staged_insert_record(
                label,
                &crate::vertex::IdKey::Text(external_id.to_string()),
                Some(projection),
            ),
        }
    }

    pub fn get_vertex_by_i64_projected(
        &self,
        label: LabelId,
        external_id: i64,
        projection: &[std::sync::Arc<str>],
        ts: Timestamp,
    ) -> Option<VertexRecord> {
        // ID-index hits are mapping-revalidated; the record itself is
        // pending-aware through `read_record`.
        if !self.persistent.is_open.load(Ordering::Acquire) {
            return None;
        }

        let internal_id = self.persistent.data_store.with_vertex_tables(|tables| {
            let table = tables.get(&label)?;
            self.resolve_internal_id_rechecked(table, label, ExternalRef::I64(external_id), ts)
        });
        match internal_id {
            Some(internal_id) => self.read_record(label, internal_id, Some(projection), ts),
            None => self.staged_insert_record(
                label,
                &crate::vertex::IdKey::Int(external_id),
                Some(projection),
            ),
        }
    }

    pub fn get_vertex_by_i64(
        &self,
        label: LabelId,
        external_id: i64,
        ts: Timestamp,
    ) -> Option<VertexRecord> {
        if !self.persistent.is_open.load(Ordering::Acquire) {
            return None;
        }

        self.get_full_record(label, ExternalRef::I64(external_id), ts)
    }

    pub fn get_vertex_by_internal_id(
        &self,
        label: LabelId,
        internal_id: u32,
        ts: Timestamp,
    ) -> Option<VertexRecord> {
        if !self.persistent.is_open.load(Ordering::Acquire) {
            return None;
        }

        self.get_full_record_by_internal_id(label, internal_id, ts)
    }

    pub fn get_external_vertex_id(
        &self,
        label: LabelId,
        internal_id: u32,
        ts: Timestamp,
    ) -> Option<VertexId> {
        self.persistent
            .data_store
            .with_vertex_tables(|vertex_tables| {
                let table = vertex_tables.get(&label)?;
                let key = table.get_external_id(internal_id, ts)?;
                match key {
                    crate::vertex::IdKey::Int(i) => VertexId::try_from_int64(i).ok(),
                    crate::vertex::IdKey::Text(s) => VertexId::try_from_string(&s).ok(),
                }
            })
    }

    pub fn get_external_id_by_internal_id(
        &self,
        label: LabelId,
        internal_id: u32,
    ) -> Option<VertexId> {
        self.persistent
            .data_store
            .with_vertex_tables(|vertex_tables| {
                let table = vertex_tables.get(&label)?;
                let key = table.get_external_id_raw(internal_id)?;
                match key {
                    crate::vertex::IdKey::Int(i) => VertexId::try_from_int64(i).ok(),
                    crate::vertex::IdKey::Text(s) => VertexId::try_from_string(&s).ok(),
                }
            })
    }
}