icydb-core 0.159.3

//! Module: data::store
//! Responsibility: stable BTreeMap-backed row persistence.
//! Does not own: key/row validation policy beyond type boundaries.
//! Boundary: commit/executor call into this layer after prevalidation.

use crate::{
    db::data::{CanonicalRow, DataKey, RawDataKey, RawRow},
    types::EntityTag,
};
use canic_cdk::structures::{BTreeMap, DefaultMemoryImpl, btreemap::Iter, memory::VirtualMemory};
#[cfg(feature = "diagnostics")]
use std::cell::Cell;
use std::ops::{Bound, RangeBounds};

#[cfg(feature = "diagnostics")]
thread_local! {
    static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
}

#[cfg(feature = "diagnostics")]
fn record_data_store_get_call() {
    DATA_STORE_GET_CALL_COUNT.with(|count| {
        count.set(count.get().saturating_add(1));
    });
}

///
/// DataStore
///
/// Thin persistence wrapper over one stable BTreeMap.
///
/// Invariant: callers provide already-validated `RawDataKey` and canonical row bytes.
/// This type intentionally does not enforce commit-phase ordering.
///

pub struct DataStore {
    map: BTreeMap<RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
}

impl DataStore {
    /// Initialize a data store with the provided backing memory.
    #[must_use]
    pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
        Self {
            map: BTreeMap::init(memory),
        }
    }

    /// Insert or replace one row by raw key.
    pub(in crate::db) fn insert(&mut self, key: RawDataKey, row: CanonicalRow) -> Option<RawRow> {
        self.map.insert(key, row.into_raw_row())
    }

    /// Insert one raw row directly for corruption-focused test setup only.
    #[cfg(test)]
    pub(in crate::db) fn insert_raw_for_test(
        &mut self,
        key: RawDataKey,
        row: RawRow,
    ) -> Option<RawRow> {
        self.map.insert(key, row)
    }

    /// Remove one row by raw key.
    pub(in crate::db) fn remove(&mut self, key: &RawDataKey) -> Option<RawRow> {
        self.map.remove(key)
    }

    /// Load one row by raw key.
    pub(in crate::db) fn get(&self, key: &RawDataKey) -> Option<RawRow> {
        #[cfg(feature = "diagnostics")]
        record_data_store_get_call();

        self.map.get(key)
    }

    /// Return whether one raw key exists without cloning the row payload.
    #[must_use]
    pub(in crate::db) fn contains(&self, key: &RawDataKey) -> bool {
        self.map.contains_key(key)
    }

    /// Clear all stored rows from the data store.
    #[cfg(test)]
    pub(in crate::db) fn clear(&mut self) {
        self.map.clear();
    }

    /// Return the number of stored rows without exposing the backing map.
    #[must_use]
    pub(in crate::db) fn len(&self) -> u64 {
        self.map.len()
    }

    /// Return whether the data store currently contains no rows.
    #[must_use]
    #[cfg(test)]
    pub(in crate::db) fn is_empty(&self) -> bool {
        self.map.is_empty()
    }

    /// Return raw row entries in canonical storage order.
    pub(in crate::db) fn entries(
        &self,
    ) -> Iter<'_, RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
        self.map.iter()
    }

    /// Iterate over raw row entries whose keys belong to the provided storage range.
    pub(in crate::db) fn range(
        &self,
        key_range: impl RangeBounds<RawDataKey>,
    ) -> Iter<'_, RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
        self.map.range(key_range)
    }

    /// Iterate over raw row entries for one entity using compact prefix bounds.
    pub(in crate::db) fn range_for_entity(
        &self,
        entity: EntityTag,
    ) -> Iter<'_, RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
        let range = DataKey::raw_entity_prefix_range_for(entity);
        match range.upper_exclusive() {
            Some(upper) => self.map.range((
                Bound::Included(range.lower_inclusive().clone()),
                Bound::Excluded(upper.clone()),
            )),
            None => self.map.range((
                Bound::Included(range.lower_inclusive().clone()),
                Bound::Unbounded,
            )),
        }
    }

    /// Sum of bytes used by all stored rows.
    pub(in crate::db) fn memory_bytes(&self) -> u64 {
        // Report map footprint as key bytes + row bytes per entry.
        self.entries()
            .map(|entry| entry.key().as_bytes().len() as u64 + entry.value().len() as u64)
            .sum()
    }

    /// Return the monotonic perf-only count of stable row fetches seen by this process.
    #[cfg(feature = "diagnostics")]
    pub(in crate::db) fn current_get_call_count() -> u64 {
        DATA_STORE_GET_CALL_COUNT.with(Cell::get)
    }
}