icydb_core/db/data/
store.rs1use 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
27pub struct DataStore {
37 map: BTreeMap<RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
38}
39
40impl DataStore {
41 #[must_use]
43 pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
44 Self {
45 map: BTreeMap::init(memory),
46 }
47 }
48
49 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 #[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 pub(in crate::db) fn remove(&mut self, key: &RawDataKey) -> Option<RawRow> {
66 self.map.remove(key)
67 }
68
69 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 #[must_use]
79 pub(in crate::db) fn contains(&self, key: &RawDataKey) -> bool {
80 self.map.contains_key(key)
81 }
82
83 #[cfg(test)]
85 pub(in crate::db) fn clear(&mut self) {
86 self.map.clear();
87 }
88
89 #[must_use]
91 pub(in crate::db) fn len(&self) -> u64 {
92 self.map.len()
93 }
94
95 #[must_use]
97 #[cfg(test)]
98 pub(in crate::db) fn is_empty(&self) -> bool {
99 self.map.is_empty()
100 }
101
102 pub(in crate::db) fn entries(
104 &self,
105 ) -> Iter<'_, RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
106 self.map.iter()
107 }
108
109 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 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 pub(in crate::db) fn memory_bytes(&self) -> u64 {
137 self.entries()
139 .map(|entry| entry.key().as_bytes().len() as u64 + entry.value().len() as u64)
140 .sum()
141 }
142
143 #[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}