icydb_core/db/data/
store.rs1use 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
30pub struct DataStore {
40 map: BTreeMap<RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
41}
42
43impl DataStore {
44 #[must_use]
46 pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
47 Self {
48 map: BTreeMap::init(memory),
49 }
50 }
51
52 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 #[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 pub(in crate::db) fn remove(&mut self, key: &RawDataStoreKey) -> Option<RawRow> {
73 self.map.remove(key)
74 }
75
76 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 #[must_use]
86 pub(in crate::db) fn contains(&self, key: &RawDataStoreKey) -> bool {
87 self.map.contains_key(key)
88 }
89
90 #[cfg(test)]
92 pub(in crate::db) fn clear(&mut self) {
93 self.map.clear();
94 }
95
96 #[must_use]
98 pub(in crate::db) fn len(&self) -> u64 {
99 self.map.len()
100 }
101
102 #[must_use]
104 #[cfg(test)]
105 pub(in crate::db) fn is_empty(&self) -> bool {
106 self.map.is_empty()
107 }
108
109 pub(in crate::db) fn entries(
111 &self,
112 ) -> Iter<'_, RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
113 self.map.iter()
114 }
115
116 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 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 pub(in crate::db) fn memory_bytes(&self) -> u64 {
135 self.entries()
137 .map(|entry| entry.key().as_bytes().len() as u64 + entry.value().len() as u64)
138 .sum()
139 }
140
141 #[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}