Skip to main content

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, ordered_overlay_entries},
11        positioned_overlay::{
12            JournalOverlayPosition, PositionedOverlayMetadata, PositionedOverlayRetirement,
13        },
14    },
15    types::EntityTag,
16};
17use ic_memory::{DefaultMemoryImpl, RuntimeMemory};
18use ic_storage::adapter::StableBTreeMap;
19#[cfg(test)]
20use std::cell::Cell;
21use std::collections::{BTreeMap as HeapBTreeMap, BTreeSet};
22use std::convert::Infallible;
23use std::ops::{Bound, RangeBounds};
24
25#[cfg(test)]
26thread_local! {
27    static DATA_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
28}
29
30#[cfg(test)]
31fn record_data_store_get_call() {
32    DATA_STORE_GET_CALL_COUNT.with(|count| {
33        count.set(count.get().saturating_add(1));
34    });
35}
36
37///
38/// DataStore
39///
40/// Thin persistence wrapper over one journaled or heap BTreeMap.
41///
42/// Invariant: callers provide already-validated `RawDataStoreKey` and canonical row bytes.
43/// This type intentionally does not enforce commit-phase ordering.
44///
45
46pub struct DataStore {
47    backend: DataStoreBackend,
48    generation: u64,
49    entity_cardinality: EntityCardinality,
50}
51
52enum DataStoreBackend {
53    Heap(HeapBTreeMap<RawDataStoreKey, RawRow>),
54    Journaled {
55        canonical: StableBTreeMap<RawDataStoreKey, RawRow, DefaultMemoryImpl>,
56        live: HeapBTreeMap<RawDataStoreKey, RawRow>,
57        tombstones: BTreeSet<RawDataStoreKey>,
58        positions: PositionedOverlayMetadata<RawDataStoreKey>,
59        entity_cardinality_delta: EntityCardinalityDelta,
60    },
61}
62
63/// Preflighted provenance publication for direct journal record families.
64#[cfg(feature = "migration")]
65pub(in crate::db) struct PreparedDataPositionPublication {
66    keys: Vec<RawDataStoreKey>,
67    position: JournalOverlayPosition,
68}
69
70/// Preflighted exact retirement for one complete journal batch.
71pub(in crate::db) struct PreparedDataPositionRetirement {
72    entries: Vec<(RawDataStoreKey, PositionedOverlayRetirement)>,
73}
74
75/// One visible row read that borrows heap/live state and owns stable state.
76///
77/// Callers that only need selected fields can evaluate a borrowed row while
78/// the store handle is active. Stable-structure reads remain owned because
79/// that backend cannot expose a value reference beyond its storage call.
80pub(in crate::db) enum StoredRowRead<'a> {
81    Missing,
82    Borrowed(&'a RawRow),
83    Owned(RawRow),
84}
85
86impl StoredRowRead<'_> {
87    /// Borrow the visible row regardless of its physical backing.
88    #[must_use]
89    pub(in crate::db) const fn as_row(&self) -> Option<&RawRow> {
90        match self {
91            Self::Missing => None,
92            Self::Borrowed(row) => Some(row),
93            Self::Owned(row) => Some(row),
94        }
95    }
96
97    /// Convert the visible row into the existing owned get contract.
98    #[must_use]
99    fn into_owned(self) -> Option<RawRow> {
100        match self {
101            Self::Missing => None,
102            Self::Borrowed(row) => Some(row.clone()),
103            Self::Owned(row) => Some(row),
104        }
105    }
106}
107
108/// Control-flow result for store traversal visitors.
109#[derive(Clone, Copy, Debug, Eq, PartialEq)]
110pub(in crate::db) enum StoreVisit {
111    Continue,
112    Stop,
113}
114
115impl StoreVisit {
116    const fn should_stop(self) -> bool {
117        matches!(self, Self::Stop)
118    }
119}
120
121impl DataStore {
122    /// Initialize a volatile heap-backed data store.
123    #[must_use]
124    pub const fn init_heap() -> Self {
125        Self {
126            backend: DataStoreBackend::Heap(HeapBTreeMap::new()),
127            generation: 0,
128            entity_cardinality: EntityCardinality::empty(),
129        }
130    }
131
132    /// Initialize a journaled cached-stable data store.
133    ///
134    /// Normal writes update only the live projection. The canonical stable map
135    /// is the future fold target and is not mutated by this wrapper's write
136    /// methods.
137    #[must_use]
138    pub fn init_journaled(memory: RuntimeMemory<DefaultMemoryImpl>) -> Self {
139        let canonical = StableBTreeMap::init(memory);
140        let entity_cardinality = if canonical.is_empty() {
141            EntityCardinality::empty()
142        } else {
143            EntityCardinality::unavailable()
144        };
145        Self {
146            backend: DataStoreBackend::Journaled {
147                canonical,
148                live: HeapBTreeMap::new(),
149                tombstones: BTreeSet::new(),
150                positions: PositionedOverlayMetadata::new(),
151                entity_cardinality_delta: EntityCardinalityDelta::empty(),
152            },
153            generation: 0,
154            // Stable rows remain authoritative after reinitialization. Exact
155            // zero cardinality is still known for an empty canonical map;
156            // populated maps remain unavailable without a startup scan.
157            entity_cardinality,
158        }
159    }
160
161    /// Insert or replace one row by raw key.
162    pub(in crate::db) fn insert(
163        &mut self,
164        key: RawDataStoreKey,
165        row: CanonicalRow,
166    ) -> Option<RawRow> {
167        let row = row.into_raw_row();
168        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
169            self.get(&key)
170        } else {
171            None
172        };
173        let cardinality_key = key.clone();
174        let previous = match &mut self.backend {
175            DataStoreBackend::Heap(map) => map.insert(key, row),
176            DataStoreBackend::Journaled {
177                live, tombstones, ..
178            } => {
179                tombstones.remove(&key);
180                live.insert(key, row);
181                previous_journaled
182            }
183        };
184        self.entity_cardinality
185            .apply_insert(&cardinality_key, previous.is_some());
186        self.apply_entity_overlay_delta(&cardinality_key, previous.is_some(), true);
187        self.bump_generation();
188        previous
189    }
190
191    /// Insert one raw row directly for corruption-focused test setup only.
192    #[cfg(test)]
193    pub(in crate::db) fn insert_raw_for_test(
194        &mut self,
195        key: RawDataStoreKey,
196        row: RawRow,
197    ) -> Option<RawRow> {
198        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
199            self.get(&key)
200        } else {
201            None
202        };
203        let cardinality_key = key.clone();
204        let previous = match &mut self.backend {
205            DataStoreBackend::Heap(map) => map.insert(key, row),
206            DataStoreBackend::Journaled {
207                live, tombstones, ..
208            } => {
209                tombstones.remove(&key);
210                live.insert(key, row);
211                previous_journaled
212            }
213        };
214        self.entity_cardinality
215            .apply_insert(&cardinality_key, previous.is_some());
216        self.apply_entity_overlay_delta(&cardinality_key, previous.is_some(), true);
217        self.bump_generation();
218        previous
219    }
220
221    /// Remove one row by raw key.
222    pub(in crate::db) fn remove(&mut self, key: &RawDataStoreKey) -> Option<RawRow> {
223        let previous_journaled = if matches!(self.backend, DataStoreBackend::Journaled { .. }) {
224            self.get(key)
225        } else {
226            None
227        };
228        let previous = match &mut self.backend {
229            DataStoreBackend::Heap(map) => map.remove(key),
230            DataStoreBackend::Journaled {
231                live, tombstones, ..
232            } => {
233                live.remove(key);
234                tombstones.insert(key.clone());
235                previous_journaled
236            }
237        };
238        self.entity_cardinality
239            .apply_remove(key, previous.is_some());
240        self.apply_entity_overlay_delta(key, previous.is_some(), false);
241        self.bump_generation();
242        previous
243    }
244
245    /// Reset the volatile projection for journaled recovery without mutating
246    /// the canonical stable base.
247    pub(in crate::db) fn reset_journaled_live_projection(
248        &mut self,
249    ) -> Result<(), crate::error::InternalError> {
250        let DataStoreBackend::Journaled {
251            canonical,
252            live,
253            tombstones,
254            positions,
255            entity_cardinality_delta,
256        } = &mut self.backend
257        else {
258            return Err(crate::error::InternalError::store_invariant());
259        };
260
261        live.clear();
262        tombstones.clear();
263        positions.clear();
264        *entity_cardinality_delta = EntityCardinalityDelta::empty();
265        self.entity_cardinality = if canonical.is_empty() {
266            EntityCardinality::empty()
267        } else {
268            EntityCardinality::unavailable()
269        };
270        self.bump_generation();
271
272        Ok(())
273    }
274
275    /// Apply one recovered journal row put into the volatile projection.
276    #[cfg(any(test, feature = "migration"))]
277    pub(in crate::db) fn apply_recovered_journal_put(
278        &mut self,
279        key: RawDataStoreKey,
280        row: RawRow,
281    ) -> Result<Option<RawRow>, crate::error::InternalError> {
282        let DataStoreBackend::Journaled {
283            canonical,
284            live,
285            tombstones,
286            ..
287        } = &mut self.backend
288        else {
289            return Err(crate::error::InternalError::store_invariant());
290        };
291
292        let previous = if tombstones.contains(&key) {
293            None
294        } else {
295            live.get(&key).cloned().or_else(|| canonical.get(&key))
296        };
297        tombstones.remove(&key);
298        let cardinality_key = key.clone();
299        live.insert(key, row);
300        self.entity_cardinality
301            .apply_insert(&cardinality_key, previous.is_some());
302        self.apply_entity_overlay_delta(&cardinality_key, previous.is_some(), true);
303        self.bump_generation();
304
305        Ok(previous)
306    }
307
308    /// Publish one preflighted positioned row value or tombstone.
309    pub(in crate::db) fn publish_preflighted_journal_entry(
310        &mut self,
311        key: RawDataStoreKey,
312        row: Option<RawRow>,
313        position: JournalOverlayPosition,
314    ) -> Result<(), crate::error::InternalError> {
315        // Cardinality needs presence only; never materialize a previous row.
316        let previous_present = self.contains(&key);
317        let DataStoreBackend::Journaled {
318            live,
319            tombstones,
320            positions,
321            entity_cardinality_delta,
322            ..
323        } = &mut self.backend
324        else {
325            return Err(crate::error::InternalError::store_invariant());
326        };
327        let next_present = row.is_some();
328        if let Some(row) = row {
329            tombstones.remove(&key);
330            live.insert(key.clone(), row);
331            self.entity_cardinality.apply_insert(&key, previous_present);
332        } else {
333            live.remove(&key);
334            tombstones.insert(key.clone());
335            self.entity_cardinality.apply_remove(&key, previous_present);
336        }
337        entity_cardinality_delta.apply_presence_transition(&key, previous_present, next_present);
338        positions.publish_preflighted(key, position);
339        self.bump_generation();
340
341        Ok(())
342    }
343
344    /// Validate and publish one positioned row for direct store tests.
345    #[cfg(test)]
346    pub(in crate::db) fn publish_positioned_journal_entry(
347        &mut self,
348        key: RawDataStoreKey,
349        row: Option<RawRow>,
350        position: JournalOverlayPosition,
351    ) -> Result<(), crate::error::InternalError> {
352        self.preflight_positioned_journal_entry(&key, position)?;
353        self.publish_preflighted_journal_entry(key, row, position)
354    }
355
356    /// Preflight row provenance before marker publication.
357    pub(in crate::db) fn preflight_positioned_journal_entry(
358        &self,
359        key: &RawDataStoreKey,
360        position: JournalOverlayPosition,
361    ) -> Result<(), crate::error::InternalError> {
362        let DataStoreBackend::Journaled { positions, .. } = &self.backend else {
363            return Err(crate::error::InternalError::store_invariant());
364        };
365        positions.preflight_publish(key, position)
366    }
367
368    /// Preflight direct row provenance before marker publication.
369    #[cfg(feature = "migration")]
370    pub(in crate::db) fn prepare_position_publication(
371        &self,
372        keys: impl IntoIterator<Item = RawDataStoreKey>,
373        position: JournalOverlayPosition,
374    ) -> Result<PreparedDataPositionPublication, crate::error::InternalError> {
375        let DataStoreBackend::Journaled { positions, .. } = &self.backend else {
376            return Err(crate::error::InternalError::store_invariant());
377        };
378        let keys = keys.into_iter().collect::<BTreeSet<_>>();
379        for key in &keys {
380            positions.preflight_publish(key, position)?;
381        }
382        Ok(PreparedDataPositionPublication {
383            keys: keys.into_iter().collect(),
384            position,
385        })
386    }
387
388    /// Publish direct row provenance after its values have been applied.
389    #[cfg(feature = "migration")]
390    pub(in crate::db) fn publish_prepared_positions(
391        &mut self,
392        prepared: PreparedDataPositionPublication,
393    ) {
394        let DataStoreBackend::Journaled { positions, .. } = &mut self.backend else {
395            debug_assert!(false, "preflighted row positions require a journaled store");
396            return;
397        };
398        for key in prepared.keys {
399            positions.publish_preflighted(key, prepared.position);
400        }
401    }
402
403    /// Preflight exact row-overlay retirement before canonical mutation.
404    pub(in crate::db) fn prepare_position_retirement(
405        &self,
406        keys: impl IntoIterator<Item = RawDataStoreKey>,
407        position: JournalOverlayPosition,
408    ) -> Result<PreparedDataPositionRetirement, crate::error::InternalError> {
409        let DataStoreBackend::Journaled { positions, .. } = &self.backend else {
410            return Err(crate::error::InternalError::store_invariant());
411        };
412        let entries = keys
413            .into_iter()
414            .collect::<BTreeSet<_>>()
415            .into_iter()
416            .map(|key| {
417                positions
418                    .preflight_retirement(&key, position)
419                    .map(|retirement| (key, retirement))
420            })
421            .collect::<Result<Vec<_>, _>>()?;
422        Ok(PreparedDataPositionRetirement { entries })
423    }
424
425    /// Retire only exact row overlays after canonical mutation succeeds.
426    pub(in crate::db) fn apply_prepared_position_retirement(
427        &mut self,
428        prepared: PreparedDataPositionRetirement,
429    ) {
430        let DataStoreBackend::Journaled {
431            live,
432            tombstones,
433            positions,
434            ..
435        } = &mut self.backend
436        else {
437            debug_assert!(
438                false,
439                "preflighted row retirement requires a journaled store"
440            );
441            return;
442        };
443        for (key, retirement) in prepared.entries {
444            if retirement == PositionedOverlayRetirement::Exact {
445                live.remove(&key);
446                tombstones.remove(&key);
447                positions.retire_preflighted(&key, retirement);
448            }
449        }
450    }
451
452    /// Apply one folded journal row put into the canonical stable base.
453    pub(in crate::db) fn fold_recovered_journal_put(
454        &mut self,
455        key: RawDataStoreKey,
456        row: RawRow,
457    ) -> Result<Option<RawRow>, crate::error::InternalError> {
458        let DataStoreBackend::Journaled {
459            canonical,
460            live,
461            tombstones,
462            ..
463        } = &mut self.backend
464        else {
465            return Err(crate::error::InternalError::store_invariant());
466        };
467
468        let visible = !live.contains_key(&key) && !tombstones.contains(&key);
469        let cardinality_key = key.clone();
470        let previous = canonical.insert(key, row);
471        if visible {
472            self.entity_cardinality
473                .apply_insert(&cardinality_key, previous.is_some());
474        } else if previous.is_none() {
475            self.apply_entity_overlay_delta(&cardinality_key, true, false);
476        }
477        self.bump_generation();
478
479        Ok(previous)
480    }
481
482    /// Apply one folded journal row delete into the canonical stable base.
483    pub(in crate::db) fn fold_recovered_journal_delete(
484        &mut self,
485        key: &RawDataStoreKey,
486    ) -> Result<Option<RawRow>, crate::error::InternalError> {
487        let DataStoreBackend::Journaled {
488            canonical,
489            live,
490            tombstones,
491            ..
492        } = &mut self.backend
493        else {
494            return Err(crate::error::InternalError::store_invariant());
495        };
496
497        let visible = !live.contains_key(key) && !tombstones.contains(key);
498        let previous = canonical.remove(key);
499        if visible {
500            self.entity_cardinality
501                .apply_remove(key, previous.is_some());
502        } else if previous.is_some() {
503            self.apply_entity_overlay_delta(key, false, true);
504        }
505        self.bump_generation();
506
507        Ok(previous)
508    }
509
510    /// Prove that recovered journal rows can be folded into canonical storage.
511    pub(in crate::db) fn preflight_fold_recovered_journal(
512        &self,
513    ) -> Result<(), crate::error::InternalError> {
514        match self.backend {
515            DataStoreBackend::Journaled { .. } => Ok(()),
516            DataStoreBackend::Heap(_) => Err(crate::error::InternalError::store_invariant()),
517        }
518    }
519
520    /// Load one row by raw key.
521    pub(in crate::db) fn get(&self, key: &RawDataStoreKey) -> Option<RawRow> {
522        self.read(key).into_owned()
523    }
524
525    /// Load one row from the canonical predecessor view.
526    ///
527    /// Online journal folding uses this view so a newer positioned live effect
528    /// cannot replace the predecessor evidence for the older batch being
529    /// canonicalized.
530    pub(in crate::db) fn get_canonical(&self, key: &RawDataStoreKey) -> Option<RawRow> {
531        match &self.backend {
532            DataStoreBackend::Heap(map) => map.get(key).cloned(),
533            DataStoreBackend::Journaled { canonical, .. } => canonical.get(key),
534        }
535    }
536
537    /// Read one visible row without cloning heap/live payloads.
538    pub(in crate::db) fn read<'a>(&'a self, key: &RawDataStoreKey) -> StoredRowRead<'a> {
539        #[cfg(test)]
540        record_data_store_get_call();
541
542        match &self.backend {
543            DataStoreBackend::Heap(map) => map
544                .get(key)
545                .map_or(StoredRowRead::Missing, StoredRowRead::Borrowed),
546            DataStoreBackend::Journaled {
547                canonical,
548                live,
549                tombstones,
550                ..
551            } => {
552                if tombstones.contains(key) {
553                    StoredRowRead::Missing
554                } else if let Some(row) = live.get(key) {
555                    StoredRowRead::Borrowed(row)
556                } else {
557                    canonical
558                        .get(key)
559                        .map_or(StoredRowRead::Missing, StoredRowRead::Owned)
560                }
561            }
562        }
563    }
564
565    /// Return whether one raw key exists without cloning the row payload.
566    #[must_use]
567    pub(in crate::db) fn contains(&self, key: &RawDataStoreKey) -> bool {
568        match &self.backend {
569            DataStoreBackend::Heap(map) => map.contains_key(key),
570            DataStoreBackend::Journaled {
571                canonical,
572                live,
573                tombstones,
574                ..
575            } => {
576                !tombstones.contains(key) && (live.contains_key(key) || canonical.contains_key(key))
577            }
578        }
579    }
580
581    /// Return whether the effective row domain is empty without reading payloads.
582    ///
583    /// Stop at the first visible key. Journal tombstones may hide canonical
584    /// keys, so physical canonical emptiness alone is not sufficient.
585    #[must_use]
586    pub(in crate::db) fn is_empty(&self) -> bool {
587        let mut empty = true;
588        let _: Result<(), Infallible> = self.visit_key_range(.., |_key| {
589            empty = false;
590            Ok(StoreVisit::Stop)
591        });
592        empty
593    }
594
595    /// Return an exact row count for tests without exposing the backing map.
596    #[cfg(test)]
597    #[must_use]
598    pub(in crate::db) fn len(&self) -> u64 {
599        match &self.backend {
600            DataStoreBackend::Heap(map) => u64::try_from(map.len()).unwrap_or(u64::MAX),
601            DataStoreBackend::Journaled { .. } => {
602                let mut count = 0_u64;
603                let _: Result<(), Infallible> = self.visit_entries(|_key, _row| {
604                    count = count.saturating_add(1);
605                    Ok(StoreVisit::Continue)
606                });
607                count
608            }
609        }
610    }
611
612    /// Return the row-store generation used to prove index metadata freshness.
613    #[must_use]
614    pub(in crate::db) const fn generation(&self) -> u64 {
615        self.generation
616    }
617
618    /// Return an exact current row count for one entity when store metadata is valid.
619    #[must_use]
620    pub(in crate::db) fn exact_entity_count(&self, entity: EntityTag) -> Option<u64> {
621        self.entity_cardinality.exact_count(entity)
622    }
623
624    /// Return the exact live-overlay delta from the canonical base for one entity.
625    #[must_use]
626    pub(in crate::db) fn exact_entity_cardinality_delta(&self, entity: EntityTag) -> Option<i64> {
627        match &self.backend {
628            DataStoreBackend::Heap(_) => Some(0),
629            DataStoreBackend::Journaled {
630                entity_cardinality_delta,
631                ..
632            } => entity_cardinality_delta.exact_delta(entity),
633        }
634    }
635
636    /// Visit raw row entries in canonical storage order.
637    pub(in crate::db) fn visit_entries<E>(
638        &self,
639        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
640    ) -> Result<(), E> {
641        match &self.backend {
642            DataStoreBackend::Heap(map) => {
643                for (key, row) in map {
644                    if visitor(key, row)?.should_stop() {
645                        break;
646                    }
647                }
648            }
649            DataStoreBackend::Journaled { .. } => Self::visit_journaled_entries_in_bounds(
650                &self.backend,
651                (Bound::Unbounded, Bound::Unbounded),
652                &mut visitor,
653            )?,
654        }
655
656        Ok(())
657    }
658
659    /// Visit canonical predecessor rows after one exclusive physical checkpoint.
660    ///
661    /// Optional derived-evidence builders use this view so newer live overlays
662    /// cannot contaminate a generation bound to the canonical fold watermark.
663    pub(in crate::db) fn visit_canonical_entries_after(
664        &self,
665        checkpoint: Option<&RawDataStoreKey>,
666        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<bool, crate::error::InternalError>,
667    ) -> Result<(), crate::error::InternalError> {
668        let lower = checkpoint.map_or(Bound::Unbounded, Bound::Excluded);
669        match &self.backend {
670            DataStoreBackend::Heap(map) => {
671                for (key, row) in map.range((lower, Bound::Unbounded)) {
672                    if visitor(key, row)? {
673                        break;
674                    }
675                }
676            }
677            DataStoreBackend::Journaled { canonical, .. } => {
678                for entry in canonical.range((lower, Bound::Unbounded)) {
679                    if visitor(entry.key(), &entry.value())? {
680                        break;
681                    }
682                }
683            }
684        }
685        Ok(())
686    }
687
688    /// Return whether the canonical predecessor domain is physically empty.
689    ///
690    /// This bounded root observation never walks live overlays or materializes
691    /// row payloads.
692    pub(in crate::db) fn canonical_is_empty(&self) -> Result<bool, crate::error::InternalError> {
693        match &self.backend {
694            DataStoreBackend::Journaled { canonical, .. } => Ok(canonical.is_empty()),
695            DataStoreBackend::Heap(_) => Err(crate::error::InternalError::store_invariant()),
696        }
697    }
698
699    /// Visit raw row entries whose keys belong to the provided storage range.
700    pub(in crate::db) fn visit_range<E>(
701        &self,
702        key_range: impl RangeBounds<RawDataStoreKey>,
703        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
704    ) -> Result<(), E> {
705        let bounds = (key_range.start_bound(), key_range.end_bound());
706        match &self.backend {
707            DataStoreBackend::Heap(map) => {
708                for (key, row) in map.range(bounds) {
709                    if visitor(key, row)?.should_stop() {
710                        break;
711                    }
712                }
713            }
714            DataStoreBackend::Journaled { .. } => {
715                Self::visit_journaled_entries_in_bounds(&self.backend, bounds, &mut visitor)?;
716            }
717        }
718
719        Ok(())
720    }
721
722    /// Visit one ascending row range while allowing the caller to stop after
723    /// seeing the key but before a stable row payload is materialized.
724    ///
725    /// Mixed journal overlays retain the ordinary key-then-point-read path;
726    /// one physical backing can keep the range iterator open for the complete
727    /// scan and avoid a second tree lookup per row.
728    pub(in crate::db) fn try_visit_range_with_row_preflight<E>(
729        &self,
730        key_range: impl RangeBounds<RawDataStoreKey>,
731        mut preflight: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
732        mut visitor: impl FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
733    ) -> Result<Option<bool>, E> {
734        let bounds = (key_range.start_bound(), key_range.end_bound());
735        let mut stopped = false;
736        match &self.backend {
737            DataStoreBackend::Heap(map) => {
738                for (key, row) in map.range(bounds) {
739                    if preflight(key)?.should_stop() {
740                        stopped = true;
741                        break;
742                    }
743                    if visitor(key, row)?.should_stop() {
744                        stopped = true;
745                        break;
746                    }
747                }
748            }
749            DataStoreBackend::Journaled {
750                canonical,
751                live,
752                tombstones,
753                ..
754            } if canonical.is_empty() => {
755                for (key, row) in live.range(bounds) {
756                    if tombstones.contains(key) {
757                        continue;
758                    }
759                    if preflight(key)?.should_stop() {
760                        stopped = true;
761                        break;
762                    }
763                    if visitor(key, row)?.should_stop() {
764                        stopped = true;
765                        break;
766                    }
767                }
768            }
769            DataStoreBackend::Journaled {
770                canonical,
771                live,
772                tombstones,
773                ..
774            } if live.is_empty() && tombstones.is_empty() => {
775                for entry in canonical.range(bounds) {
776                    if preflight(entry.key())?.should_stop() {
777                        stopped = true;
778                        break;
779                    }
780                    if visitor(entry.key(), &entry.value())?.should_stop() {
781                        stopped = true;
782                        break;
783                    }
784                }
785            }
786            DataStoreBackend::Journaled { .. } => return Ok(None),
787        }
788
789        Ok(Some(!stopped))
790    }
791
792    /// Visit only raw keys in storage order without fetching row payloads.
793    ///
794    /// Primary-key access streams use this boundary to discover candidate
795    /// identities before the terminal row runtime decides whether the payload
796    /// is needed. Journaled traversal merges canonical and live keys while
797    /// preserving live overrides and tombstones without reading stable values.
798    pub(in crate::db) fn visit_key_range<E>(
799        &self,
800        key_range: impl RangeBounds<RawDataStoreKey>,
801        visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
802    ) -> Result<(), E> {
803        self.visit_keys_in_bounds(
804            (key_range.start_bound(), key_range.end_bound()),
805            false,
806            visitor,
807        )
808    }
809
810    /// Visit only raw keys in reverse storage order without fetching row payloads.
811    pub(in crate::db) fn visit_key_range_rev<E>(
812        &self,
813        key_range: impl RangeBounds<RawDataStoreKey>,
814        visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
815    ) -> Result<(), E> {
816        self.visit_keys_in_bounds(
817            (key_range.start_bound(), key_range.end_bound()),
818            true,
819            visitor,
820        )
821    }
822
823    /// Independent physical byte baseline for storage-report tests.
824    #[cfg(test)]
825    pub(in crate::db) fn memory_bytes(&self) -> u64 {
826        // Report map footprint as key bytes + row bytes per entry.
827        let mut bytes = 0u64;
828        let _: Result<(), Infallible> = self.visit_entries(|key, row| {
829            bytes = bytes.saturating_add(key.as_bytes().len() as u64 + row.len() as u64);
830            Ok(StoreVisit::Continue)
831        });
832        bytes
833    }
834
835    const fn bump_generation(&mut self) {
836        self.generation = self.generation.saturating_add(1);
837    }
838
839    #[cfg(test)]
840    fn rebuild_entity_cardinality_from_entries(&mut self) {
841        let mut cardinality = EntityCardinality::empty();
842        let _: Result<(), Infallible> = self.visit_entries(|key, _row| {
843            cardinality.apply_present_key(key);
844            Ok(StoreVisit::Continue)
845        });
846        self.entity_cardinality = cardinality;
847    }
848
849    fn apply_entity_overlay_delta(
850        &mut self,
851        key: &RawDataStoreKey,
852        previous_present: bool,
853        next_present: bool,
854    ) {
855        let DataStoreBackend::Journaled {
856            entity_cardinality_delta,
857            ..
858        } = &mut self.backend
859        else {
860            return;
861        };
862        entity_cardinality_delta.apply_presence_transition(key, previous_present, next_present);
863    }
864
865    /// Return the monotonic perf-only count of stable row fetches seen by this process.
866    #[cfg(test)]
867    pub(in crate::db) fn current_get_call_count() -> u64 {
868        DATA_STORE_GET_CALL_COUNT.with(Cell::get)
869    }
870
871    fn visit_keys_in_bounds<E>(
872        &self,
873        bounds: (Bound<&RawDataStoreKey>, Bound<&RawDataStoreKey>),
874        reverse: bool,
875        mut visitor: impl FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
876    ) -> Result<(), E> {
877        match &self.backend {
878            DataStoreBackend::Heap(map) => {
879                if reverse {
880                    for (key, _row) in map.range(bounds).rev() {
881                        if visitor(key)?.should_stop() {
882                            break;
883                        }
884                    }
885                } else {
886                    for (key, _row) in map.range(bounds) {
887                        if visitor(key)?.should_stop() {
888                            break;
889                        }
890                    }
891                }
892            }
893            DataStoreBackend::Journaled { .. } => {
894                Self::visit_journaled_keys_in_bounds(&self.backend, bounds, reverse, &mut visitor)?;
895            }
896        }
897
898        Ok(())
899    }
900
901    // Share journaled traversal across callback types while retaining typed
902    // caller errors. Key-only traversal must never materialize row payloads.
903    fn visit_journaled_keys_in_bounds<E>(
904        backend: &DataStoreBackend,
905        bounds: (Bound<&RawDataStoreKey>, Bound<&RawDataStoreKey>),
906        reverse: bool,
907        visitor: &mut dyn FnMut(&RawDataStoreKey) -> Result<StoreVisit, E>,
908    ) -> Result<(), E> {
909        let DataStoreBackend::Journaled {
910            canonical,
911            live,
912            tombstones,
913            ..
914        } = backend
915        else {
916            return Ok(());
917        };
918
919        if canonical.is_empty() {
920            if reverse {
921                for (key, _row) in live.range(bounds).rev() {
922                    if visitor(key)?.should_stop() {
923                        return Ok(());
924                    }
925                }
926            } else {
927                for (key, _row) in live.range(bounds) {
928                    if visitor(key)?.should_stop() {
929                        return Ok(());
930                    }
931                }
932            }
933            return Ok(());
934        }
935
936        if live.is_empty() && tombstones.is_empty() {
937            if reverse {
938                for entry in canonical.range(bounds).rev() {
939                    if visitor(entry.key())?.should_stop() {
940                        return Ok(());
941                    }
942                }
943            } else {
944                for entry in canonical.range(bounds) {
945                    if visitor(entry.key())?.should_stop() {
946                        return Ok(());
947                    }
948                }
949            }
950            return Ok(());
951        }
952
953        let direction = if reverse {
954            Direction::Desc
955        } else {
956            Direction::Asc
957        };
958        match direction {
959            Direction::Asc => {
960                for entry in ordered_overlay_entries(
961                    canonical.range(bounds),
962                    live.range(bounds),
963                    direction,
964                    |entry| entry.key(),
965                    |entry| entry.0,
966                    tombstones,
967                ) {
968                    let visit = match entry {
969                        OrderedOverlayEntry::Canonical(canonical_entry) => {
970                            visitor(canonical_entry.key())?
971                        }
972                        OrderedOverlayEntry::Live((key, _row)) => visitor(key)?,
973                    };
974                    if visit.should_stop() {
975                        return Ok(());
976                    }
977                }
978            }
979            Direction::Desc => {
980                for entry in ordered_overlay_entries(
981                    canonical.range(bounds).rev(),
982                    live.range(bounds).rev(),
983                    direction,
984                    |entry| entry.key(),
985                    |entry| entry.0,
986                    tombstones,
987                ) {
988                    let visit = match entry {
989                        OrderedOverlayEntry::Canonical(canonical_entry) => {
990                            visitor(canonical_entry.key())?
991                        }
992                        OrderedOverlayEntry::Live((key, _row)) => visitor(key)?,
993                    };
994                    if visit.should_stop() {
995                        return Ok(());
996                    }
997                }
998            }
999        }
1000
1001        Ok(())
1002    }
1003
1004    // Keep one journaled row traversal body per error type, independently of
1005    // each caller's closure. The ordered overlay remains the merge authority.
1006    fn visit_journaled_entries_in_bounds<E>(
1007        backend: &DataStoreBackend,
1008        bounds: (Bound<&RawDataStoreKey>, Bound<&RawDataStoreKey>),
1009        visitor: &mut dyn FnMut(&RawDataStoreKey, &RawRow) -> Result<StoreVisit, E>,
1010    ) -> Result<(), E> {
1011        let DataStoreBackend::Journaled {
1012            canonical,
1013            live,
1014            tombstones,
1015            ..
1016        } = backend
1017        else {
1018            return Ok(());
1019        };
1020
1021        if canonical.is_empty() {
1022            for (key, row) in live.range(bounds) {
1023                if visitor(key, row)?.should_stop() {
1024                    return Ok(());
1025                }
1026            }
1027            return Ok(());
1028        }
1029
1030        if live.is_empty() && tombstones.is_empty() {
1031            for entry in canonical.range(bounds) {
1032                if visitor(entry.key(), &entry.value())?.should_stop() {
1033                    return Ok(());
1034                }
1035            }
1036            return Ok(());
1037        }
1038
1039        for entry in ordered_overlay_entries(
1040            canonical.range(bounds),
1041            live.range(bounds),
1042            Direction::Asc,
1043            |entry| entry.key(),
1044            |entry| entry.0,
1045            tombstones,
1046        ) {
1047            let visit = match entry {
1048                OrderedOverlayEntry::Canonical(canonical_entry) => {
1049                    visitor(canonical_entry.key(), &canonical_entry.value())?
1050                }
1051                OrderedOverlayEntry::Live((key, row)) => visitor(key, row)?,
1052            };
1053            if visit.should_stop() {
1054                return Ok(());
1055            }
1056        }
1057
1058        Ok(())
1059    }
1060}
1061
1062#[derive(Clone, Debug)]
1063struct EntityCardinality {
1064    counts: HeapBTreeMap<EntityTag, u64>,
1065    decodable: bool,
1066}
1067
1068#[derive(Clone, Debug)]
1069struct EntityCardinalityDelta {
1070    counts: HeapBTreeMap<EntityTag, i64>,
1071    decodable: bool,
1072}
1073
1074impl EntityCardinalityDelta {
1075    const fn empty() -> Self {
1076        Self {
1077            counts: HeapBTreeMap::new(),
1078            decodable: true,
1079        }
1080    }
1081
1082    fn exact_delta(&self, entity: EntityTag) -> Option<i64> {
1083        self.decodable
1084            .then(|| self.counts.get(&entity).copied().unwrap_or(0))
1085    }
1086
1087    fn apply_presence_transition(
1088        &mut self,
1089        key: &RawDataStoreKey,
1090        previous_present: bool,
1091        next_present: bool,
1092    ) {
1093        if !self.decodable || previous_present == next_present {
1094            return;
1095        }
1096        let Some(entity) = key.entity_tag_prefix() else {
1097            self.counts.clear();
1098            self.decodable = false;
1099            return;
1100        };
1101        let delta = if next_present { 1_i64 } else { -1_i64 };
1102        let count = self.counts.entry(entity).or_insert(0);
1103        let Some(next) = count.checked_add(delta) else {
1104            self.counts.clear();
1105            self.decodable = false;
1106            return;
1107        };
1108        *count = next;
1109        if *count == 0 {
1110            self.counts.remove(&entity);
1111        }
1112    }
1113}
1114
1115impl EntityCardinality {
1116    const fn empty() -> Self {
1117        Self {
1118            counts: HeapBTreeMap::new(),
1119            decodable: true,
1120        }
1121    }
1122
1123    const fn unavailable() -> Self {
1124        Self {
1125            counts: HeapBTreeMap::new(),
1126            decodable: false,
1127        }
1128    }
1129
1130    fn exact_count(&self, entity: EntityTag) -> Option<u64> {
1131        self.decodable
1132            .then(|| self.counts.get(&entity).copied().unwrap_or(0))
1133    }
1134
1135    fn apply_insert(&mut self, key: &RawDataStoreKey, previous_present: bool) {
1136        if previous_present {
1137            return;
1138        }
1139        self.apply_present_key(key);
1140    }
1141
1142    fn apply_remove(&mut self, key: &RawDataStoreKey, previous_present: bool) {
1143        if !previous_present {
1144            return;
1145        }
1146        self.apply_removed_key(key);
1147    }
1148
1149    fn apply_present_key(&mut self, key: &RawDataStoreKey) {
1150        if !self.decodable {
1151            return;
1152        }
1153        let Some(entity) = key.entity_tag_prefix() else {
1154            self.invalidate();
1155            return;
1156        };
1157
1158        let count = self.counts.entry(entity).or_insert(0);
1159        *count = count.saturating_add(1);
1160    }
1161
1162    fn apply_removed_key(&mut self, key: &RawDataStoreKey) {
1163        if !self.decodable {
1164            return;
1165        }
1166        let Some(entity) = key.entity_tag_prefix() else {
1167            self.invalidate();
1168            return;
1169        };
1170
1171        if let Some(count) = self.counts.get_mut(&entity) {
1172            *count = count.saturating_sub(1);
1173            if *count == 0 {
1174                self.counts.remove(&entity);
1175            }
1176        }
1177    }
1178
1179    fn invalidate(&mut self) {
1180        self.counts.clear();
1181        self.decodable = false;
1182    }
1183}
1184
1185#[cfg(test)]
1186mod tests;