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