icydb_core/db/data/
store.rs

1//! Module: data::store
2//! Responsibility: journaled-or-heap row storage behind the data-store boundary.
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        direction::Direction,
10        ordered_overlay::{OrderedOverlayEntry, OrderedOverlayVisit, visit_ordered_overlay},
11    },
12    types::EntityTag,
13};
14use ic_stable_structures::{
15    BTreeMap as StableBTreeMap, DefaultMemoryImpl, memory_manager::VirtualMemory,
16};
17#[cfg(all(feature = "sql", feature = "diagnostics"))]
18use std::cell::Cell;
19use std::collections::{BTreeMap as HeapBTreeMap, BTreeSet};
20use std::convert::Infallible;
21use std::ops::{Bound, RangeBounds};
22
23#[cfg(all(feature = "sql", feature = "diagnostics"))]
24thread_local! {
25    static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
26}
27
28#[cfg(all(feature = "sql", feature = "diagnostics"))]
29fn record_data_store_get_call() {
30    DATA_STORE_GET_CALL_COUNT.with(|count| {
31        count.set(count.get().saturating_add(1));
32    });
33}
34
35///
36/// DataStore
37///
38/// Thin persistence wrapper over one journaled or heap BTreeMap.
39///
40/// Invariant: callers provide already-validated `RawDataStoreKey` and canonical row bytes.
41/// This type intentionally does not enforce commit-phase ordering.
42///
43
44pub struct DataStore {
45    backend: DataStoreBackend,
46    generation: u64,
47    entity_cardinality: EntityCardinality,
48}
49
50enum DataStoreBackend {
51    Heap(HeapBTreeMap<RawDataStoreKey, RawRow>),
52    Journaled {
53        canonical: StableBTreeMap<RawDataStoreKey, RawRow, VirtualMemory<DefaultMemoryImpl>>,
54        live: HeapBTreeMap<RawDataStoreKey, RawRow>,
55        tombstones: BTreeSet<RawDataStoreKey>,
56    },
57}
58
59/// One visible row read that borrows heap/live state and owns stable state.
60///
61/// Callers that only need selected fields can evaluate a borrowed row while
62/// the store handle is active. Stable-structure reads remain owned because
63/// that backend cannot expose a value reference beyond its storage call.
64pub(in crate::db) enum StoredRowRead<'a> {
65    Missing,
66    Borrowed(&'a RawRow),
67    Owned(RawRow),
68}
69
70impl StoredRowRead<'_> {
71    /// Borrow the visible row regardless of its physical backing.
72    #[must_use]
73    pub(in crate::db) const fn as_row(&self) -> Option<&RawRow> {
74        match self {
75            Self::Missing => None,
76            Self::Borrowed(row) => Some(row),
77            Self::Owned(row) => Some(row),
78        }
79    }
80
81    /// Convert the visible row into the existing owned get contract.
82    #[must_use]
83    fn into_owned(self) -> Option<RawRow> {
84        match self {
85            Self::Missing => None,
86            Self::Borrowed(row) => Some(row.clone()),
87            Self::Owned(row) => Some(row),
88        }
89    }
90}
91
92/// Control-flow result for store traversal visitors.
93#[derive(Clone, Copy, Debug, Eq, PartialEq)]
94pub(in crate::db) enum StoreVisit {
95    Continue,
96    Stop,
97}
98
99impl StoreVisit {
100    const fn should_stop(self) -> bool {
101        matches!(self, Self::Stop)
102    }
103}
104
105impl DataStore {
106    /// Initialize a volatile heap-backed data store.
107    #[must_use]
108    pub const fn init_heap() -> Self {
109        Self {
110            backend: DataStoreBackend::Heap(HeapBTreeMap::new()),
111            generation: 0,
112            entity_cardinality: EntityCardinality::empty(),
113        }
114    }
115
116    /// Initialize a journaled cached-stable data store.
117    ///
118    /// Normal writes update only the live projection. The canonical stable map
119    /// is the future fold target and is not mutated by this wrapper's write
120    /// methods.
121    #[must_use]
122    pub fn init_journaled(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
123        let mut store = Self {
124            backend: DataStoreBackend::Journaled {
125                canonical: StableBTreeMap::init(memory),
126                live: HeapBTreeMap::new(),
127                tombstones: BTreeSet::new(),
128            },
129            generation: 0,
130            entity_cardinality: EntityCardinality::empty(),
131        };
132        store.rebuild_entity_cardinality_from_entries();
133        store
134    }
135
136    /// Insert or replace one row by raw key.
137    pub(in crate::db) fn insert(
138        &mut self,
139        key: RawDataStoreKey,
140        row: CanonicalRow,
141    ) -> Option<RawRow> {
142        let row = row.into_raw_row();
143        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
144            self.get(&key)
145        } else {
146            None
147        };
148        let cardinality_key = key.clone();
149        let previous = match &mut self.backend {
150            DataStoreBackend::Heap(map) => map.insert(key, row),
151            DataStoreBackend::Journaled {
152                live, tombstones, ..
153            } => {
154                tombstones.remove(&key);
155                live.insert(key, row);
156                previous_journaled
157            }
158        };
159        self.entity_cardinality
160            .apply_insert(&cardinality_key, previous.as_ref());
161        self.bump_generation();
162        previous
163    }
164
165    /// Insert one raw row directly for corruption-focused test setup only.
166    #[cfg(test)]
167    pub(in crate::db) fn insert_raw_for_test(
168        &mut self,
169        key: RawDataStoreKey,
170        row: RawRow,
171    ) -> Option<RawRow> {
172        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
173            self.get(&key)
174        } else {
175            None
176        };
177        let cardinality_key = key.clone();
178        let previous = match &mut self.backend {
179            DataStoreBackend::Heap(map) => map.insert(key, row),
180            DataStoreBackend::Journaled {
181                live, tombstones, ..
182            } => {
183                tombstones.remove(&key);
184                live.insert(key, row);
185                previous_journaled
186            }
187        };
188        self.entity_cardinality
189            .apply_insert(&cardinality_key, previous.as_ref());
190        self.bump_generation();
191        previous
192    }
193
194    /// Remove one row by raw key.
195    pub(in crate::db) fn remove(&mut self, key: &RawDataStoreKey) -> Option<RawRow> {
196        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
197            self.get(key)
198        } else {
199            None
200        };
201        let previous = match &mut self.backend {
202            DataStoreBackend::Heap(map) => map.remove(key),
203            DataStoreBackend::Journaled {
204                live, tombstones, ..
205            } => {
206                live.remove(key);
207                tombstones.insert(key.clone());
208                previous_journaled
209            }
210        };
211        self.entity_cardinality.apply_remove(key, previous.as_ref());
212        self.bump_generation();
213        previous
214    }
215
216    /// Reset the volatile projection for journaled recovery without mutating
217    /// the canonical stable base.
218    pub(in crate::db) fn reset_journaled_live_projection(
219        &mut self,
220    ) -> Result<(), crate::error::InternalError> {
221        let DataStoreBackend::Journaled {
222            live, tombstones, ..
223        } = &mut self.backend
224        else {
225            return Err(crate::error::InternalError::store_invariant());
226        };
227
228        live.clear();
229        tombstones.clear();
230        self.rebuild_entity_cardinality_from_entries();
231        self.bump_generation();
232
233        Ok(())
234    }
235
236    /// Apply one recovered journal row put into the volatile projection.
237    pub(in crate::db) fn apply_recovered_journal_put(
238        &mut self,
239        key: RawDataStoreKey,
240        row: RawRow,
241    ) -> Result<Option<RawRow>, crate::error::InternalError> {
242        let DataStoreBackend::Journaled {
243            canonical,
244            live,
245            tombstones,
246        } = &mut self.backend
247        else {
248            return Err(crate::error::InternalError::store_invariant());
249        };
250
251        let previous = if tombstones.contains(&key) {
252            None
253        } else {
254            live.get(&key).cloned().or_else(|| canonical.get(&key))
255        };
256        tombstones.remove(&key);
257        let cardinality_key = key.clone();
258        live.insert(key, row);
259        self.entity_cardinality
260            .apply_insert(&cardinality_key, previous.as_ref());
261        self.bump_generation();
262
263        Ok(previous)
264    }
265
266    /// Apply one recovered journal row delete into the volatile projection.
267    pub(in crate::db) fn apply_recovered_journal_delete(
268        &mut self,
269        key: &RawDataStoreKey,
270    ) -> Result<Option<RawRow>, crate::error::InternalError> {
271        let DataStoreBackend::Journaled {
272            canonical,
273            live,
274            tombstones,
275        } = &mut self.backend
276        else {
277            return Err(crate::error::InternalError::store_invariant());
278        };
279
280        let previous = if tombstones.contains(key) {
281            None
282        } else {
283            live.get(key).cloned().or_else(|| canonical.get(key))
284        };
285        live.remove(key);
286        tombstones.insert(key.clone());
287        self.entity_cardinality.apply_remove(key, previous.as_ref());
288        self.bump_generation();
289
290        Ok(previous)
291    }
292
293    /// Apply one folded journal row put into the canonical stable base.
294    pub(in crate::db) fn fold_recovered_journal_put(
295        &mut self,
296        key: RawDataStoreKey,
297        row: RawRow,
298    ) -> Result<Option<RawRow>, crate::error::InternalError> {
299        let DataStoreBackend::Journaled {
300            canonical,
301            live,
302            tombstones,
303        } = &mut self.backend
304        else {
305            return Err(crate::error::InternalError::store_invariant());
306        };
307
308        let visible = !live.contains_key(&key) && !tombstones.contains(&key);
309        let cardinality_key = key.clone();
310        let previous = canonical.insert(key, row);
311        if visible {
312            self.entity_cardinality
313                .apply_insert(&cardinality_key, previous.as_ref());
314        }
315        self.bump_generation();
316
317        Ok(previous)
318    }
319
320    /// Apply one folded journal row delete into the canonical stable base.
321    pub(in crate::db) fn fold_recovered_journal_delete(
322        &mut self,
323        key: &RawDataStoreKey,
324    ) -> Result<Option<RawRow>, crate::error::InternalError> {
325        let DataStoreBackend::Journaled {
326            canonical,
327            live,
328            tombstones,
329        } = &mut self.backend
330        else {
331            return Err(crate::error::InternalError::store_invariant());
332        };
333
334        let visible = !live.contains_key(key) && !tombstones.contains(key);
335        let previous = canonical.remove(key);
336        if visible {
337            self.entity_cardinality.apply_remove(key, previous.as_ref());
338        }
339        self.bump_generation();
340
341        Ok(previous)
342    }
343
344    /// Load one row by raw key.
345    pub(in crate::db) fn get(&self, key: &RawDataStoreKey) -> Option<RawRow> {
346        self.read(key).into_owned()
347    }
348
349    /// Read one visible row without cloning heap/live payloads.
350    pub(in crate::db) fn read<'a>(&'a self, key: &RawDataStoreKey) -> StoredRowRead<'a> {
351        #[cfg(all(feature = "sql", feature = "diagnostics"))]
352        record_data_store_get_call();
353
354        match &self.backend {
355            DataStoreBackend::Heap(map) => map
356                .get(key)
357                .map_or(StoredRowRead::Missing, StoredRowRead::Borrowed),
358            DataStoreBackend::Journaled {
359                canonical,
360                live,
361                tombstones,
362            } => {
363                if tombstones.contains(key) {
364                    StoredRowRead::Missing
365                } else if let Some(row) = live.get(key) {
366                    StoredRowRead::Borrowed(row)
367                } else {
368                    canonical
369                        .get(key)
370                        .map_or(StoredRowRead::Missing, StoredRowRead::Owned)
371                }
372            }
373        }
374    }
375
376    /// Return whether one raw key exists without cloning the row payload.
377    #[must_use]
378    pub(in crate::db) fn contains(&self, key: &RawDataStoreKey) -> bool {
379        match &self.backend {
380            DataStoreBackend::Heap(map) => map.contains_key(key),
381            DataStoreBackend::Journaled {
382                canonical,
383                live,
384                tombstones,
385            } => {
386                !tombstones.contains(key)
387                    && (live.contains_key(key) || canonical.get(key).is_some())
388            }
389        }
390    }
391
392    /// Return the number of stored rows without exposing the backing map.
393    #[must_use]
394    pub(in crate::db) fn len(&self) -> u64 {
395        match &self.backend {
396            DataStoreBackend::Heap(map) => u64::try_from(map.len()).unwrap_or(u64::MAX),
397            DataStoreBackend::Journaled { .. } => {
398                let mut count = 0_u64;
399                let _: Result<(), Infallible> = self.visit_entries(|_key, _row| {
400                    count = count.saturating_add(1);
401                    Ok(StoreVisit::Continue)
402                });
403                count
404            }
405        }
406    }
407
408    /// Return the row-store generation used to prove index metadata freshness.
409    #[must_use]
410    pub(in crate::db) const fn generation(&self) -> u64 {
411        self.generation
412    }
413
414    /// Return an exact current row count for one entity when store metadata is valid.
415    #[must_use]
416    pub(in crate::db) fn exact_entity_count(&self, entity: EntityTag) -> Option<u64> {
417        self.entity_cardinality.exact_count(entity)
418    }
419
420    /// Visit raw row entries in canonical storage order.
421    pub(in crate::db) fn visit_entries<E>(
422        &self,
423        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
424    ) -> Result<(), E> {
425        match &self.backend {
426            DataStoreBackend::Heap(map) => {
427                for (key, row) in map {
428                    if visitor(key, row)?.should_stop() {
429                        break;
430                    }
431                }
432            }
433            DataStoreBackend::Journaled {
434                canonical: _,
435                live: _,
436                tombstones: _,
437            } => Self::visit_journaled_entries_in_bounds(
438                &self.backend,
439                (Bound::Unbounded, Bound::Unbounded),
440                visitor,
441            )?,
442        }
443
444        Ok(())
445    }
446
447    /// Visit raw row entries whose keys belong to the provided storage range.
448    pub(in crate::db) fn visit_range<E>(
449        &self,
450        key_range: impl RangeBounds<RawDataStoreKey>,
451        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
452    ) -> Result<(), E> {
453        let bounds = Self::owned_range_bounds(&key_range);
454        match &self.backend {
455            DataStoreBackend::Heap(map) => {
456                for (key, row) in map.range((bounds.0.clone(), bounds.1)) {
457                    if visitor(key, row)?.should_stop() {
458                        break;
459                    }
460                }
461            }
462            DataStoreBackend::Journaled {
463                canonical: _,
464                live: _,
465                tombstones: _,
466            } => Self::visit_journaled_entries_in_bounds(&self.backend, bounds, visitor)?,
467        }
468
469        Ok(())
470    }
471
472    /// Visit one ascending row range while allowing the caller to stop after
473    /// seeing the key but before a stable row payload is materialized.
474    ///
475    /// Mixed journal overlays retain the ordinary key-then-point-read path;
476    /// one physical backing can keep the range iterator open for the complete
477    /// scan and avoid a second tree lookup per row.
478    pub(in crate::db) fn try_visit_range_with_row_preflight<E>(
479        &self,
480        key_range: impl RangeBounds<RawDataStoreKey>,
481        mut preflight: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
482        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
483    ) -> Result<Option<bool>, E> {
484        let bounds = Self::owned_range_bounds(&key_range);
485        let mut stopped = false;
486        match &self.backend {
487            DataStoreBackend::Heap(map) => {
488                for (key, row) in map.range((bounds.0.clone(), bounds.1)) {
489                    if preflight(key)?.should_stop() {
490                        stopped = true;
491                        break;
492                    }
493                    if visitor(key, row)?.should_stop() {
494                        stopped = true;
495                        break;
496                    }
497                }
498            }
499            DataStoreBackend::Journaled {
500                canonical,
501                live,
502                tombstones,
503            } if canonical.is_empty() => {
504                for (key, row) in live.range((bounds.0.clone(), bounds.1)) {
505                    if tombstones.contains(key) {
506                        continue;
507                    }
508                    if preflight(key)?.should_stop() {
509                        stopped = true;
510                        break;
511                    }
512                    if visitor(key, row)?.should_stop() {
513                        stopped = true;
514                        break;
515                    }
516                }
517            }
518            DataStoreBackend::Journaled {
519                canonical,
520                live,
521                tombstones,
522            } if live.is_empty() && tombstones.is_empty() => {
523                for entry in canonical.range((bounds.0.clone(), bounds.1)) {
524                    if preflight(entry.key())?.should_stop() {
525                        stopped = true;
526                        break;
527                    }
528                    if visitor(entry.key(), &entry.value())?.should_stop() {
529                        stopped = true;
530                        break;
531                    }
532                }
533            }
534            DataStoreBackend::Journaled { .. } => return Ok(None),
535        }
536
537        Ok(Some(!stopped))
538    }
539
540    /// Visit only raw keys in storage order without fetching row payloads.
541    ///
542    /// Primary-key access streams use this boundary to discover candidate
543    /// identities before the terminal row runtime decides whether the payload
544    /// is needed. Journaled traversal merges canonical and live keys while
545    /// preserving live overrides and tombstones without reading stable values.
546    pub(in crate::db) fn visit_key_range<E>(
547        &self,
548        key_range: impl RangeBounds<RawDataStoreKey>,
549        visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
550    ) -> Result<(), E> {
551        self.visit_keys_in_bounds(Self::owned_range_bounds(&key_range), false, visitor)
552    }
553
554    /// Visit only raw keys in reverse storage order without fetching row payloads.
555    pub(in crate::db) fn visit_key_range_rev<E>(
556        &self,
557        key_range: impl RangeBounds<RawDataStoreKey>,
558        visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
559    ) -> Result<(), E> {
560        self.visit_keys_in_bounds(Self::owned_range_bounds(&key_range), true, visitor)
561    }
562
563    /// Sum of bytes used by all stored rows.
564    pub(in crate::db) fn memory_bytes(&self) -> u64 {
565        // Report map footprint as key bytes + row bytes per entry.
566        let mut bytes = 0u64;
567        let _: Result<(), Infallible> = self.visit_entries(|key, row| {
568            bytes = bytes.saturating_add(key.as_bytes().len() as u64 + row.len() as u64);
569            Ok(StoreVisit::Continue)
570        });
571        bytes
572    }
573
574    const fn bump_generation(&mut self) {
575        self.generation = self.generation.saturating_add(1);
576    }
577
578    fn rebuild_entity_cardinality_from_entries(&mut self) {
579        let mut cardinality = EntityCardinality::empty();
580        let _: Result<(), Infallible> = self.visit_entries(|key, _row| {
581            cardinality.apply_present_key(key);
582            Ok(StoreVisit::Continue)
583        });
584        self.entity_cardinality = cardinality;
585    }
586
587    /// Return the monotonic perf-only count of stable row fetches seen by this process.
588    #[cfg(all(feature = "sql", feature = "diagnostics"))]
589    pub(in crate::db) fn current_get_call_count() -> u64 {
590        DATA_STORE_GET_CALL_COUNT.with(Cell::get)
591    }
592
593    fn owned_range_bounds(
594        key_range: &impl RangeBounds<RawDataStoreKey>,
595    ) -> (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>) {
596        let lower = match key_range.start_bound() {
597            Bound::Included(key) => Bound::Included(key.clone()),
598            Bound::Excluded(key) => Bound::Excluded(key.clone()),
599            Bound::Unbounded => Bound::Unbounded,
600        };
601        let upper = match key_range.end_bound() {
602            Bound::Included(key) => Bound::Included(key.clone()),
603            Bound::Excluded(key) => Bound::Excluded(key.clone()),
604            Bound::Unbounded => Bound::Unbounded,
605        };
606
607        (lower, upper)
608    }
609
610    fn visit_keys_in_bounds<E>(
611        &self,
612        bounds: (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>),
613        reverse: bool,
614        mut visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
615    ) -> Result<(), E> {
616        match &self.backend {
617            DataStoreBackend::Heap(map) => {
618                if reverse {
619                    for (key, _row) in map.range(bounds).rev() {
620                        if visitor(key)?.should_stop() {
621                            break;
622                        }
623                    }
624                } else {
625                    for (key, _row) in map.range(bounds) {
626                        if visitor(key)?.should_stop() {
627                            break;
628                        }
629                    }
630                }
631            }
632            DataStoreBackend::Journaled { .. } => {
633                Self::visit_journaled_keys_in_bounds(&self.backend, bounds, reverse, visitor)?;
634            }
635        }
636
637        Ok(())
638    }
639
640    fn visit_journaled_keys_in_bounds<E>(
641        backend: &DataStoreBackend,
642        bounds: (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>),
643        reverse: bool,
644        mut visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
645    ) -> Result<(), E> {
646        let DataStoreBackend::Journaled {
647            canonical,
648            live,
649            tombstones,
650        } = backend
651        else {
652            return Ok(());
653        };
654
655        if canonical.is_empty() {
656            if reverse {
657                for (key, _row) in live.range(bounds).rev() {
658                    if visitor(key)?.should_stop() {
659                        return Ok(());
660                    }
661                }
662            } else {
663                for (key, _row) in live.range(bounds) {
664                    if visitor(key)?.should_stop() {
665                        return Ok(());
666                    }
667                }
668            }
669            return Ok(());
670        }
671
672        if live.is_empty() && tombstones.is_empty() {
673            if reverse {
674                for entry in canonical.range(bounds).rev() {
675                    if visitor(entry.key())?.should_stop() {
676                        return Ok(());
677                    }
678                }
679            } else {
680                for entry in canonical.range(bounds) {
681                    if visitor(entry.key())?.should_stop() {
682                        return Ok(());
683                    }
684                }
685            }
686            return Ok(());
687        }
688
689        let direction = if reverse {
690            Direction::Desc
691        } else {
692            Direction::Asc
693        };
694        match direction {
695            Direction::Asc => visit_ordered_overlay(
696                canonical.range((bounds.0.clone(), bounds.1.clone())),
697                live.range((bounds.0, bounds.1)),
698                direction,
699                |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
700                |canonical_entry| !tombstones.contains(canonical_entry.key()),
701                |live_entry| !tombstones.contains(live_entry.0),
702                |entry| {
703                    let visit = match entry {
704                        OrderedOverlayEntry::Canonical(canonical_entry) => {
705                            visitor(canonical_entry.key())?
706                        }
707                        OrderedOverlayEntry::Live((key, _row)) => visitor(key)?,
708                    };
709                    Ok(if visit.should_stop() {
710                        OrderedOverlayVisit::Stop
711                    } else {
712                        OrderedOverlayVisit::Continue
713                    })
714                },
715            ),
716            Direction::Desc => visit_ordered_overlay(
717                canonical.range((bounds.0.clone(), bounds.1.clone())).rev(),
718                live.range((bounds.0, bounds.1)).rev(),
719                direction,
720                |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
721                |canonical_entry| !tombstones.contains(canonical_entry.key()),
722                |live_entry| !tombstones.contains(live_entry.0),
723                |entry| {
724                    let visit = match entry {
725                        OrderedOverlayEntry::Canonical(canonical_entry) => {
726                            visitor(canonical_entry.key())?
727                        }
728                        OrderedOverlayEntry::Live((key, _row)) => visitor(key)?,
729                    };
730                    Ok(if visit.should_stop() {
731                        OrderedOverlayVisit::Stop
732                    } else {
733                        OrderedOverlayVisit::Continue
734                    })
735                },
736            ),
737        }
738    }
739
740    fn visit_journaled_entries_in_bounds<E>(
741        backend: &DataStoreBackend,
742        bounds: (Bound<RawDataStoreKey>, Bound<RawDataStoreKey>),
743        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
744    ) -> Result<(), E> {
745        let DataStoreBackend::Journaled {
746            canonical,
747            live,
748            tombstones,
749        } = backend
750        else {
751            return Ok(());
752        };
753
754        if canonical.is_empty() {
755            for (key, row) in live.range(bounds) {
756                if visitor(key, row)?.should_stop() {
757                    return Ok(());
758                }
759            }
760            return Ok(());
761        }
762
763        if live.is_empty() && tombstones.is_empty() {
764            for entry in canonical.range(bounds) {
765                if visitor(entry.key(), &entry.value())?.should_stop() {
766                    return Ok(());
767                }
768            }
769            return Ok(());
770        }
771
772        visit_ordered_overlay(
773            canonical.range((bounds.0.clone(), bounds.1.clone())),
774            live.range((bounds.0, bounds.1)),
775            Direction::Asc,
776            |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
777            |canonical_entry| !tombstones.contains(canonical_entry.key()),
778            |live_entry| !tombstones.contains(live_entry.0),
779            |entry| {
780                let visit = match entry {
781                    OrderedOverlayEntry::Canonical(canonical_entry) => {
782                        visitor(canonical_entry.key(), &canonical_entry.value())?
783                    }
784                    OrderedOverlayEntry::Live((key, row)) => visitor(key, row)?,
785                };
786                Ok(if visit.should_stop() {
787                    OrderedOverlayVisit::Stop
788                } else {
789                    OrderedOverlayVisit::Continue
790                })
791            },
792        )
793    }
794}
795
796#[derive(Clone, Debug)]
797struct EntityCardinality {
798    counts: HeapBTreeMap<EntityTag, u64>,
799    decodable: bool,
800}
801
802impl EntityCardinality {
803    const fn empty() -> Self {
804        Self {
805            counts: HeapBTreeMap::new(),
806            decodable: true,
807        }
808    }
809
810    fn exact_count(&self, entity: EntityTag) -> Option<u64> {
811        self.decodable
812            .then(|| self.counts.get(&entity).copied().unwrap_or(0))
813    }
814
815    fn apply_insert(&mut self, key: &RawDataStoreKey, previous: Option<&RawRow>) {
816        if previous.is_some() {
817            return;
818        }
819        self.apply_present_key(key);
820    }
821
822    fn apply_remove(&mut self, key: &RawDataStoreKey, previous: Option<&RawRow>) {
823        if previous.is_none() {
824            return;
825        }
826        self.apply_removed_key(key);
827    }
828
829    fn apply_present_key(&mut self, key: &RawDataStoreKey) {
830        if !self.decodable {
831            return;
832        }
833        let Some(entity) = key.entity_tag_prefix() else {
834            self.invalidate();
835            return;
836        };
837
838        let count = self.counts.entry(entity).or_insert(0);
839        *count = count.saturating_add(1);
840    }
841
842    fn apply_removed_key(&mut self, key: &RawDataStoreKey) {
843        if !self.decodable {
844            return;
845        }
846        let Some(entity) = key.entity_tag_prefix() else {
847            self.invalidate();
848            return;
849        };
850
851        if let Some(count) = self.counts.get_mut(&entity) {
852            *count = count.saturating_sub(1);
853            if *count == 0 {
854                self.counts.remove(&entity);
855            }
856        }
857    }
858
859    fn invalidate(&mut self) {
860        self.counts.clear();
861        self.decodable = false;
862    }
863}
864
865#[cfg(test)]
866mod tests;
icydb_core/db/data/store.rs

icydb_core/db/data/
store.rs
