use crate::db::data::{CanonicalRow, DataKey, RawDataKey, RawRow};
use canic_cdk::structures::{BTreeMap, DefaultMemoryImpl, btreemap::Iter, memory::VirtualMemory};
#[cfg(feature = "diagnostics")]
use std::cell::Cell;
use std::ops::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));
});
}
pub struct DataStore {
map: BTreeMap<RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
}
impl DataStore {
#[must_use]
pub fn init(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
Self {
map: BTreeMap::init(memory),
}
}
pub(in crate::db) fn insert(&mut self, key: RawDataKey, row: CanonicalRow) -> Option<RawRow> {
self.map.insert(key, row.into_raw_row())
}
#[cfg(test)]
pub(in crate::db) fn insert_raw_for_test(
&mut self,
key: RawDataKey,
row: RawRow,
) -> Option<RawRow> {
self.map.insert(key, row)
}
pub(in crate::db) fn remove(&mut self, key: &RawDataKey) -> Option<RawRow> {
self.map.remove(key)
}
pub(in crate::db) fn get(&self, key: &RawDataKey) -> Option<RawRow> {
#[cfg(feature = "diagnostics")]
record_data_store_get_call();
self.map.get(key)
}
#[must_use]
pub(in crate::db) fn contains(&self, key: &RawDataKey) -> bool {
self.map.contains_key(key)
}
#[cfg(test)]
pub(in crate::db) fn clear(&mut self) {
self.map.clear();
}
#[must_use]
pub(in crate::db) fn len(&self) -> u64 {
self.map.len()
}
#[must_use]
#[cfg(test)]
pub(in crate::db) fn is_empty(&self) -> bool {
self.map.is_empty()
}
pub(in crate::db) fn entries(
&self,
) -> Iter<'_, RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
self.map.iter()
}
pub(in crate::db) fn range(
&self,
key_range: impl RangeBounds<RawDataKey>,
) -> Iter<'_, RawDataKey, RawRow, VirtualMemory<DefaultMemoryImpl>> {
self.map.range(key_range)
}
pub(in crate::db) fn memory_bytes(&self) -> u64 {
self.entries()
.map(|entry| DataKey::STORED_SIZE_BYTES + entry.value().len() as u64)
.sum()
}
#[cfg(feature = "diagnostics")]
pub(in crate::db) fn current_get_call_count() -> u64 {
DATA_STORE_GET_CALL_COUNT.with(Cell::get)
}
}