Skip to main content

icydb_core/db/data/
store.rs

1//! Module: data::store
2//! Responsibility: stable BTreeMap-backed row persistence.
3//! Does not own: key/row validation policy beyond type boundaries.
4//! Boundary: commit/executor call into this layer after prevalidation.
5
6use crate::{
7    db::{
8        data::{CanonicalRow, RawDataStoreKey, RawRow},
9        key_taxonomy::RawDataStoreKeyRange,
10    },
11    types::EntityTag,
12};
13use canic_cdk::structures::{BTreeMap, DefaultMemoryImpl, btreemap::Iter, memory::VirtualMemory};
14#[cfg(feature = "diagnostics")]
15use std::cell::Cell;
16use std::ops::RangeBounds;
17
18#[cfg(feature = "diagnostics")]
19thread_local! {
20    static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
21}
22
23#[cfg(feature = "diagnostics")]
24fn record_data_store_get_call() {
25    DATA_STORE_GET_CALL_COUNT.with(|count| {
26        count.set(count.get().saturating_add(1));
27    });
28}
29
30///
31/// DataStore
32///
33/// Thin persistence wrapper over one stable BTreeMap.
34///
35/// Invariant: callers provide already-validated `RawDataStoreKey` and canonical row bytes.
36/// This type intentionally does not enforce commit-phase ordering.
37///
38
39pub struct DataStore {
40    map: BTreeMap<RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
41}
42
43impl DataStore {
44    /// Initialize a data store with the provided backing memory.
45    #[must_use]
46    pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
47        Self {
48            map: BTreeMap::init(memory),
49        }
50    }
51
52    /// Insert or replace one row by raw key.
53    pub(in crate::db) fn insert(
54        &mut self,
55        key: RawDataStoreKey,
56        row: CanonicalRow,
57    ) -> Option<RawRow> {
58        self.map.insert(key, row.into_raw_row())
59    }
60
61    /// Insert one raw row directly for corruption-focused test setup only.
62    #[cfg(test)]
63    pub(in crate::db) fn insert_raw_for_test(
64        &mut self,
65        key: RawDataStoreKey,
66        row: RawRow,
67    ) -> Option<RawRow> {
68        self.map.insert(key, row)
69    }
70
71    /// Remove one row by raw key.
72    pub(in crate::db) fn remove(&mut self, key: &RawDataStoreKey) -> Option<RawRow> {
73        self.map.remove(key)
74    }
75
76    /// Load one row by raw key.
77    pub(in crate::db) fn get(&self, key: &RawDataStoreKey) -> Option<RawRow> {
78        #[cfg(feature = "diagnostics")]
79        record_data_store_get_call();
80
81        self.map.get(key)
82    }
83
84    /// Return whether one raw key exists without cloning the row payload.
85    #[must_use]
86    pub(in crate::db) fn contains(&self, key: &RawDataStoreKey) -> bool {
87        self.map.contains_key(key)
88    }
89
90    /// Clear all stored rows from the data store.
91    #[cfg(test)]
92    pub(in crate::db) fn clear(&mut self) {
93        self.map.clear();
94    }
95
96    /// Return the number of stored rows without exposing the backing map.
97    #[must_use]
98    pub(in crate::db) fn len(&self) -> u64 {
99        self.map.len()
100    }
101
102    /// Return whether the data store currently contains no rows.
103    #[must_use]
104    #[cfg(test)]
105    pub(in crate::db) fn is_empty(&self) -> bool {
106        self.map.is_empty()
107    }
108
109    /// Return raw row entries in canonical storage order.
110    pub(in crate::db) fn entries(
111        &self,
112    ) -> Iter<'_, RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
113        self.map.iter()
114    }
115
116    /// Iterate over raw row entries whose keys belong to the provided storage range.
117    pub(in crate::db) fn range(
118        &self,
119        key_range: impl RangeBounds<RawDataStoreKey>,
120    ) -> Iter<'_, RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
121        self.map.range(key_range)
122    }
123
124    /// Iterate over raw row entries for one entity using compact prefix bounds.
125    pub(in crate::db) fn range_for_entity(
126        &self,
127        entity: EntityTag,
128    ) -> Iter<'_, RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
129        let range = RawDataStoreKeyRange::entity_prefix(entity);
130        self.map.range(RawDataStoreKey::store_range_bounds(&range))
131    }
132
133    /// Sum of bytes used by all stored rows.
134    pub(in crate::db) fn memory_bytes(&self) -> u64 {
135        // Report map footprint as key bytes + row bytes per entry.
136        self.entries()
137            .map(|entry| entry.key().as_bytes().len() as u64 + entry.value().len() as u64)
138            .sum()
139    }
140
141    /// Return the monotonic perf-only count of stable row fetches seen by this process.
142    #[cfg(feature = "diagnostics")]
143    pub(in crate::db) fn current_get_call_count() -> u64 {
144        DATA_STORE_GET_CALL_COUNT.with(Cell::get)
145    }
146}