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

icydb_core/db/data/
store.rs
