icydb_core/db/index/
store.rs

1//! Module: index::store
2//! Responsibility: journaled-or-heap index-entry storage behind the index-store boundary.
3//! Does not own: range-scan resolution, continuation semantics, or predicate execution.
4//! Boundary: scan/executor layers depend on this storage boundary.
5
6use crate::db::index::{IndexId, IndexKeyKind};
7use crate::db::{
8    direction::Direction,
9    index::{IndexEntryValue, cardinality::IndexPrefixCardinality, key::RawIndexStoreKey},
10    ordered_overlay::{OrderedOverlayEntry, OrderedOverlayVisit, visit_ordered_overlay},
11};
12
13use candid::CandidType;
14use ic_stable_structures::{
15    BTreeMap as StableBTreeMap, DefaultMemoryImpl, memory_manager::VirtualMemory,
16};
17use serde::Deserialize;
18#[cfg(any(test, all(feature = "sql", feature = "diagnostics")))]
19use std::cell::Cell;
20use std::collections::{BTreeMap as HeapBTreeMap, BTreeSet};
21use std::ops::Bound;
22
23#[cfg(test)]
24thread_local! {
25    static JOURNALED_SNAPSHOT_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
26}
27
28#[cfg(all(feature = "sql", feature = "diagnostics"))]
29thread_local! {
30    static INDEX_STORE_GET_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
31    static INDEX_STORE_RANGE_SCAN_CALL_COUNT: Cell<u64> = const { Cell::new(0) };
32    static INDEX_STORE_ENTRY_READ_COUNT: Cell<u64> = const { Cell::new(0) };
33}
34
35#[cfg(all(feature = "sql", feature = "diagnostics"))]
36fn record_index_store_get_call() {
37    INDEX_STORE_GET_CALL_COUNT.with(|count| {
38        count.set(count.get().saturating_add(1));
39    });
40}
41
42#[cfg(all(feature = "sql", feature = "diagnostics"))]
43fn record_index_store_range_scan_call() {
44    INDEX_STORE_RANGE_SCAN_CALL_COUNT.with(|count| {
45        count.set(count.get().saturating_add(1));
46    });
47}
48
49#[cfg(all(feature = "sql", feature = "diagnostics"))]
50fn record_index_store_entry_read() {
51    INDEX_STORE_ENTRY_READ_COUNT.with(|count| {
52        count.set(count.get().saturating_add(1));
53    });
54}
55
56fn visit_index_store_entry<E>(
57    key: &RawIndexStoreKey,
58    value: &IndexEntryValue,
59    visit: &mut impl FnMut(&RawIndexStoreKey, &IndexEntryValue) -> Result<bool, E>,
60) -> Result<bool, E> {
61    #[cfg(all(feature = "sql", feature = "diagnostics"))]
62    record_index_store_entry_read();
63
64    visit(key, value)
65}
66
67#[cfg(test)]
68fn record_journaled_snapshot_call() {
69    JOURNALED_SNAPSHOT_CALL_COUNT.with(|count| {
70        count.set(count.get().saturating_add(1));
71    });
72}
73
74#[cfg(test)]
75fn reset_journaled_snapshot_call_count_for_tests() {
76    JOURNALED_SNAPSHOT_CALL_COUNT.with(|count| count.set(0));
77}
78
79#[cfg(test)]
80fn journaled_snapshot_call_count_for_tests() -> u64 {
81    JOURNALED_SNAPSHOT_CALL_COUNT.with(Cell::get)
82}
83
84//
85// IndexState
86//
87// Explicit lifecycle visibility state for one index store.
88// Visibility matters because planner-visible indexes must already be complete:
89// the index contents are fully built and query-visible for reads.
90//
91#[derive(CandidType, Clone, Copy, Debug, Default, Deserialize, Eq, PartialEq)]
92pub enum IndexState {
93    Building,
94    #[default]
95    Ready,
96}
97
98impl IndexState {
99    /// Return the stable lowercase text label for this lifecycle state.
100    #[must_use]
101    pub const fn as_str(self) -> &'static str {
102        match self {
103            Self::Building => "building",
104            Self::Ready => "ready",
105        }
106    }
107}
108
109///
110/// IndexStore
111///
112/// Thin persistence wrapper over one journaled or heap BTreeMap.
113///
114/// Invariant: callers provide already-validated `RawIndexStoreKey`/`IndexEntryValue`.
115///
116
117pub struct IndexStore {
118    pub(super) backend: IndexStoreBackend,
119    generation: u64,
120    state: IndexState,
121    access_state_revision: u64,
122    prefix_cardinality: IndexPrefixCardinality,
123}
124
125pub(super) enum IndexStoreBackend {
126    Heap(HeapBTreeMap<RawIndexStoreKey, IndexEntryValue>),
127    Journaled {
128        canonical:
129            StableBTreeMap<RawIndexStoreKey, IndexEntryValue, VirtualMemory<DefaultMemoryImpl>>,
130        live: HeapBTreeMap<RawIndexStoreKey, IndexEntryValue>,
131        tombstones: BTreeSet<RawIndexStoreKey>,
132    },
133}
134
135/// Control-flow result for index-store traversal visitors.
136#[derive(Clone, Copy, Debug, Eq, PartialEq)]
137pub(in crate::db) enum IndexStoreVisit {
138    Continue,
139    Stop,
140}
141
142impl IndexStoreVisit {
143    const fn should_stop(self) -> bool {
144        matches!(self, Self::Stop)
145    }
146}
147
148impl IndexStore {
149    /// Initialize a volatile heap-backed index store.
150    #[must_use]
151    pub const fn init_heap() -> Self {
152        Self {
153            backend: IndexStoreBackend::Heap(HeapBTreeMap::new()),
154            generation: 0,
155            state: IndexState::Ready,
156            access_state_revision: 1,
157            prefix_cardinality: IndexPrefixCardinality::synchronized_empty(),
158        }
159    }
160
161    /// Initialize a journaled cached-stable index store.
162    ///
163    /// Normal writes update only the live materialized projection. The
164    /// canonical stable index is updated by future fold/rebuild paths.
165    #[must_use]
166    pub fn init_journaled(memory: VirtualMemory<DefaultMemoryImpl>) -> Self {
167        let canonical = StableBTreeMap::init(memory);
168        let prefix_cardinality = if canonical.is_empty() {
169            IndexPrefixCardinality::synchronized_empty()
170        } else {
171            IndexPrefixCardinality::unavailable()
172        };
173        Self {
174            backend: IndexStoreBackend::Journaled {
175                canonical,
176                live: HeapBTreeMap::new(),
177                tombstones: BTreeSet::new(),
178            },
179            generation: 0,
180            state: IndexState::Ready,
181            access_state_revision: 1,
182            // Exact zero cardinality is known for an empty canonical map;
183            // populated maps remain unavailable without a startup scan.
184            prefix_cardinality,
185        }
186    }
187
188    /// Visit all index entries in canonical store order without exposing the
189    /// backing stable-map iterator.
190    pub(in crate::db) fn visit_entries<E>(
191        &self,
192        mut visitor: impl FnMut(&RawIndexStoreKey, &IndexEntryValue) -> Result<IndexStoreVisit, E>,
193    ) -> Result<(), E> {
194        match &self.backend {
195            IndexStoreBackend::Heap(map) => {
196                for (key, value) in map {
197                    #[cfg(all(feature = "sql", feature = "diagnostics"))]
198                    record_index_store_entry_read();
199
200                    if visitor(key, value)?.should_stop() {
201                        return Ok(());
202                    }
203                }
204            }
205            IndexStoreBackend::Journaled {
206                canonical: _,
207                live: _,
208                tombstones: _,
209            } => self.visit_journaled_entries_in_range(
210                (&Bound::Unbounded, &Bound::Unbounded),
211                Direction::Asc,
212                |key, value| visitor(key, value).map(IndexStoreVisit::should_stop),
213            )?,
214        }
215
216        Ok(())
217    }
218
219    pub(in crate::db) fn get(&self, key: &RawIndexStoreKey) -> Option<IndexEntryValue> {
220        #[cfg(all(feature = "sql", feature = "diagnostics"))]
221        record_index_store_get_call();
222
223        match &self.backend {
224            IndexStoreBackend::Heap(map) => map.get(key).cloned(),
225            IndexStoreBackend::Journaled { .. } => Self::journaled_get(&self.backend, key),
226        }
227    }
228
229    pub fn len(&self) -> u64 {
230        match &self.backend {
231            IndexStoreBackend::Heap(map) => u64::try_from(map.len()).unwrap_or(u64::MAX),
232            IndexStoreBackend::Journaled { .. } => {
233                let mut count = 0_u64;
234                let _: Result<(), std::convert::Infallible> = self.visit_entries(|_key, _value| {
235                    count = count.saturating_add(1);
236                    Ok(IndexStoreVisit::Continue)
237                });
238                count
239            }
240        }
241    }
242
243    pub fn is_empty(&self) -> bool {
244        match &self.backend {
245            IndexStoreBackend::Heap(map) => map.is_empty(),
246            IndexStoreBackend::Journaled { .. } => {
247                let mut empty = true;
248                let _: Result<(), std::convert::Infallible> = self.visit_entries(|_key, _value| {
249                    empty = false;
250                    Ok(IndexStoreVisit::Stop)
251                });
252                empty
253            }
254        }
255    }
256
257    #[must_use]
258    pub(in crate::db) const fn generation(&self) -> u64 {
259        self.generation
260    }
261
262    /// Return the explicit lifecycle state for this index store.
263    #[must_use]
264    pub(in crate::db) const fn state(&self) -> IndexState {
265        self.state
266    }
267
268    /// Return the current physical access-readiness revision.
269    #[must_use]
270    pub(in crate::db) const fn access_state_revision(&self) -> u64 {
271        self.access_state_revision
272    }
273
274    /// Return an exact user-index prefix count when the index metadata is
275    /// synchronized with the caller's authoritative row-store generation.
276    #[must_use]
277    pub(in crate::db) fn exact_prefix_cardinality(
278        &self,
279        data_generation: u64,
280        key_kind: IndexKeyKind,
281        index_id: IndexId,
282        components: &[Vec<u8>],
283    ) -> Option<u64> {
284        self.prefix_cardinality
285            .exact_count(data_generation, key_kind, index_id, components)
286    }
287
288    /// Return the sum of exact prefix counts for prefixes on the same index
289    /// when synchronized metadata can prove all requested counts.
290    #[must_use]
291    pub(in crate::db) fn exact_prefix_cardinality_sum<'a>(
292        &self,
293        data_generation: u64,
294        key_kind: IndexKeyKind,
295        index_id: IndexId,
296        component_prefixes: impl IntoIterator<Item = &'a [Vec<u8>]>,
297        stop_after: Option<u64>,
298    ) -> Option<u64> {
299        self.prefix_cardinality.exact_count_sum(
300            data_generation,
301            key_kind,
302            index_id,
303            component_prefixes,
304            stop_after,
305        )
306    }
307
308    /// Return non-empty exact child prefixes under a sparse set of already-encoded
309    /// parent prefixes when synchronized metadata can prove the bounded child set.
310    #[must_use]
311    pub(in crate::db) fn exact_child_prefixes_for_parent_set<'a>(
312        &self,
313        data_generation: u64,
314        key_kind: IndexKeyKind,
315        index_id: IndexId,
316        parent_component_prefixes: impl IntoIterator<Item = &'a [Vec<u8>]>,
317        max_children: usize,
318    ) -> Option<Vec<Vec<Vec<u8>>>> {
319        self.prefix_cardinality.exact_child_prefixes_for_parent_set(
320            data_generation,
321            key_kind,
322            index_id,
323            parent_component_prefixes,
324            max_children,
325        )
326    }
327
328    /// Mark prefix-cardinality metadata synchronized with the authoritative
329    /// row-store generation after a committed row/index transition.
330    pub(in crate::db) const fn mark_prefix_cardinality_data_generation(&mut self, generation: u64) {
331        self.prefix_cardinality.mark_synchronized(generation);
332    }
333
334    /// Mark this index store as in-progress and therefore ineligible for
335    /// planner visibility until a full authoritative rebuild ends.
336    pub(in crate::db) const fn set_access_state(&mut self, state: IndexState, revision: u64) {
337        self.state = state;
338        self.access_state_revision = revision;
339    }
340
341    pub(crate) fn insert(
342        &mut self,
343        key: RawIndexStoreKey,
344        entry: IndexEntryValue,
345    ) -> Option<IndexEntryValue> {
346        let previous_journaled = if matches!(self.backend, IndexStoreBackend::Journaled { .. }) {
347            self.get(&key)
348        } else {
349            None
350        };
351        let cardinality_key = key.clone();
352        let previous = match &mut self.backend {
353            IndexStoreBackend::Heap(map) => map.insert(key, entry.clone()),
354            IndexStoreBackend::Journaled {
355                live, tombstones, ..
356            } => {
357                tombstones.remove(&key);
358                live.insert(key, entry.clone());
359                previous_journaled
360            }
361        };
362        self.prefix_cardinality
363            .apply_insert(&cardinality_key, previous.as_ref(), &entry);
364        self.bump_generation();
365        previous
366    }
367
368    pub(crate) fn remove(&mut self, key: &RawIndexStoreKey) -> Option<IndexEntryValue> {
369        let previous_journaled = if matches!(self.backend, IndexStoreBackend::Journaled { .. }) {
370            self.get(key)
371        } else {
372            None
373        };
374        let previous = match &mut self.backend {
375            IndexStoreBackend::Heap(map) => map.remove(key),
376            IndexStoreBackend::Journaled {
377                live, tombstones, ..
378            } => {
379                live.remove(key);
380                tombstones.insert(key.clone());
381                previous_journaled
382            }
383        };
384        self.prefix_cardinality.apply_remove(key, previous.as_ref());
385        self.bump_generation();
386        previous
387    }
388
389    /// Reset the disposable journaled index overlay without traversing or
390    /// mutating the canonical stable index.
391    pub(in crate::db) fn reset_journaled_live_projection(
392        &mut self,
393        data_generation: u64,
394    ) -> Result<(), crate::error::InternalError> {
395        let IndexStoreBackend::Journaled {
396            canonical,
397            live,
398            tombstones,
399        } = &mut self.backend
400        else {
401            return Err(crate::error::InternalError::store_invariant());
402        };
403
404        live.clear();
405        tombstones.clear();
406        self.prefix_cardinality = if canonical.is_empty() {
407            let mut cardinality = IndexPrefixCardinality::synchronized_empty();
408            cardinality.mark_synchronized(data_generation);
409            cardinality
410        } else {
411            IndexPrefixCardinality::unavailable()
412        };
413        self.bump_generation();
414
415        Ok(())
416    }
417
418    /// Apply one recovered index entry directly to canonical stable storage.
419    pub(in crate::db) fn fold_recovered_journal_entry(
420        &mut self,
421        key: RawIndexStoreKey,
422        value: Option<IndexEntryValue>,
423    ) -> Result<(), crate::error::InternalError> {
424        let IndexStoreBackend::Journaled {
425            canonical,
426            live,
427            tombstones,
428        } = &mut self.backend
429        else {
430            return Err(crate::error::InternalError::store_invariant());
431        };
432
433        let visible = !live.contains_key(&key) && !tombstones.contains(&key);
434        let cardinality_key = key.clone();
435        let previous = if let Some(value) = value.as_ref() {
436            canonical.insert(key, value.clone())
437        } else {
438            canonical.remove(&key)
439        };
440        if visible {
441            if let Some(value) = value.as_ref() {
442                self.prefix_cardinality
443                    .apply_insert(&cardinality_key, previous.as_ref(), value);
444            } else {
445                self.prefix_cardinality
446                    .apply_remove(&cardinality_key, previous.as_ref());
447            }
448        }
449        self.bump_generation();
450
451        Ok(())
452    }
453
454    pub fn clear(&mut self) {
455        match &mut self.backend {
456            IndexStoreBackend::Heap(map) => map.clear(),
457            IndexStoreBackend::Journaled {
458                canonical,
459                live,
460                tombstones,
461            } => {
462                live.clear();
463                tombstones.clear();
464                for entry in canonical.iter() {
465                    tombstones.insert(entry.key().clone());
466                }
467            }
468        }
469        self.prefix_cardinality.clear_unsynchronized();
470        self.bump_generation();
471    }
472
473    /// Fold the current journaled materialized index view into the canonical
474    /// stable base and clear volatile projection state.
475    #[cfg(any(test, feature = "migration"))]
476    pub(in crate::db) fn fold_journaled_materialized_view(
477        &mut self,
478    ) -> Result<(), crate::error::InternalError> {
479        let entries = Self::journaled_entries_snapshot_for_fold(&self.backend);
480        let IndexStoreBackend::Journaled {
481            canonical,
482            live,
483            tombstones,
484        } = &mut self.backend
485        else {
486            return Err(crate::error::InternalError::store_invariant());
487        };
488
489        canonical.clear_new();
490        for (key, value) in entries {
491            canonical.insert(key, value);
492        }
493        live.clear();
494        tombstones.clear();
495        let data_generation = self.prefix_cardinality.synchronized_generation();
496        self.rebuild_prefix_cardinality_from_entries(data_generation);
497        self.bump_generation();
498
499        Ok(())
500    }
501
502    /// Sum of bytes used by all stored index entries.
503    pub fn memory_bytes(&self) -> u64 {
504        let mut bytes = 0u64;
505        let _: Result<(), std::convert::Infallible> = self.visit_entries(|key, value| {
506            bytes = bytes.saturating_add(key.as_bytes().len() as u64 + value.len() as u64);
507            Ok(IndexStoreVisit::Continue)
508        });
509        bytes
510    }
511
512    /// Return the monotonic perf-only count of index-entry fetches seen by this process.
513    #[cfg(all(feature = "sql", feature = "diagnostics"))]
514    pub(in crate::db) fn current_get_call_count() -> u64 {
515        INDEX_STORE_GET_CALL_COUNT.with(Cell::get)
516    }
517
518    /// Return the monotonic perf-only count of index range traversal probes seen by this process.
519    #[cfg(all(feature = "sql", feature = "diagnostics"))]
520    pub(in crate::db) fn current_range_scan_call_count() -> u64 {
521        INDEX_STORE_RANGE_SCAN_CALL_COUNT.with(Cell::get)
522    }
523
524    /// Return the monotonic perf-only count of index entries yielded by traversal.
525    #[cfg(all(feature = "sql", feature = "diagnostics"))]
526    pub(in crate::db) fn current_entry_read_count() -> u64 {
527        INDEX_STORE_ENTRY_READ_COUNT.with(Cell::get)
528    }
529
530    #[cfg(all(feature = "sql", feature = "diagnostics"))]
531    pub(in crate::db::index) fn record_range_scan_call() {
532        record_index_store_range_scan_call();
533    }
534
535    #[cfg(all(feature = "sql", feature = "diagnostics"))]
536    pub(in crate::db::index) fn record_merged_entry_reads(count: u64) {
537        INDEX_STORE_ENTRY_READ_COUNT.with(|total| {
538            total.set(total.get().saturating_add(count));
539        });
540    }
541
542    const fn bump_generation(&mut self) {
543        self.generation = self.generation.saturating_add(1);
544    }
545
546    #[cfg(any(test, feature = "migration"))]
547    fn rebuild_prefix_cardinality_from_entries(&mut self, data_generation: Option<u64>) {
548        self.prefix_cardinality.clear_unsynchronized();
549        let entries = Self::entries_snapshot_for_cardinality(&self.backend);
550        for (key, value) in &entries {
551            self.prefix_cardinality.apply_insert(key, None, value);
552        }
553        if let Some(data_generation) = data_generation {
554            self.prefix_cardinality.mark_synchronized(data_generation);
555        }
556    }
557
558    #[cfg(any(test, feature = "migration"))]
559    fn entries_snapshot_for_cardinality(
560        backend: &IndexStoreBackend,
561    ) -> HeapBTreeMap<RawIndexStoreKey, IndexEntryValue> {
562        match backend {
563            IndexStoreBackend::Heap(map) => map.clone(),
564            IndexStoreBackend::Journaled { .. } => {
565                Self::journaled_entries_snapshot_for_fold(backend)
566            }
567        }
568    }
569
570    fn journaled_get(
571        backend: &IndexStoreBackend,
572        key: &RawIndexStoreKey,
573    ) -> Option<IndexEntryValue> {
574        let IndexStoreBackend::Journaled {
575            canonical,
576            live,
577            tombstones,
578        } = backend
579        else {
580            return None;
581        };
582
583        if tombstones.contains(key) {
584            return None;
585        }
586        live.get(key).cloned().or_else(|| canonical.get(key))
587    }
588
589    #[cfg(any(test, feature = "migration"))]
590    pub(super) fn journaled_entries_snapshot_for_fold(
591        backend: &IndexStoreBackend,
592    ) -> HeapBTreeMap<RawIndexStoreKey, IndexEntryValue> {
593        #[cfg(test)]
594        record_journaled_snapshot_call();
595
596        let IndexStoreBackend::Journaled {
597            canonical,
598            live,
599            tombstones,
600        } = backend
601        else {
602            return HeapBTreeMap::new();
603        };
604
605        let mut entries = HeapBTreeMap::new();
606        for entry in canonical.iter() {
607            let key = entry.key().clone();
608            if !tombstones.contains(&key) {
609                entries.insert(key, entry.value());
610            }
611        }
612        for (key, value) in live {
613            if !tombstones.contains(key) {
614                entries.insert(key.clone(), value.clone());
615            }
616        }
617
618        entries
619    }
620
621    pub(super) fn visit_journaled_entries_in_range<E>(
622        &self,
623        bounds: (&Bound<RawIndexStoreKey>, &Bound<RawIndexStoreKey>),
624        direction: Direction,
625        mut visit: impl FnMut(&RawIndexStoreKey, &IndexEntryValue) -> Result<bool, E>,
626    ) -> Result<(), E> {
627        let IndexStoreBackend::Journaled {
628            canonical,
629            live,
630            tombstones,
631        } = &self.backend
632        else {
633            return Ok(());
634        };
635
636        let lower = bounds.0.clone();
637        let upper = bounds.1.clone();
638        match direction {
639            Direction::Asc if canonical.is_empty() => {
640                for (key, value) in live.range((lower, upper)) {
641                    if visit_index_store_entry(key, value, &mut visit)? {
642                        return Ok(());
643                    }
644                }
645            }
646            Direction::Desc if canonical.is_empty() => {
647                for (key, value) in live.range((lower, upper)).rev() {
648                    if visit_index_store_entry(key, value, &mut visit)? {
649                        return Ok(());
650                    }
651                }
652            }
653            Direction::Asc if live.is_empty() && tombstones.is_empty() => {
654                for entry in canonical.range((lower, upper)) {
655                    if visit_index_store_entry(entry.key(), &entry.value(), &mut visit)? {
656                        return Ok(());
657                    }
658                }
659            }
660            Direction::Desc if live.is_empty() && tombstones.is_empty() => {
661                for entry in canonical.range((lower, upper)).rev() {
662                    if visit_index_store_entry(entry.key(), &entry.value(), &mut visit)? {
663                        return Ok(());
664                    }
665                }
666            }
667            Direction::Asc => {
668                visit_ordered_overlay(
669                    canonical.range((lower.clone(), upper.clone())),
670                    live.range((lower, upper)),
671                    direction,
672                    |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
673                    |canonical_entry| !tombstones.contains(canonical_entry.key()),
674                    |live_entry| !tombstones.contains(live_entry.0),
675                    |entry| {
676                        let should_stop = match entry {
677                            OrderedOverlayEntry::Canonical(canonical_entry) => {
678                                visit_index_store_entry(
679                                    canonical_entry.key(),
680                                    &canonical_entry.value(),
681                                    &mut visit,
682                                )?
683                            }
684                            OrderedOverlayEntry::Live((key, value)) => {
685                                visit_index_store_entry(key, value, &mut visit)?
686                            }
687                        };
688                        Ok(if should_stop {
689                            OrderedOverlayVisit::Stop
690                        } else {
691                            OrderedOverlayVisit::Continue
692                        })
693                    },
694                )?;
695            }
696            Direction::Desc => {
697                visit_ordered_overlay(
698                    canonical.range((lower.clone(), upper.clone())).rev(),
699                    live.range((lower, upper)).rev(),
700                    direction,
701                    |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
702                    |canonical_entry| !tombstones.contains(canonical_entry.key()),
703                    |live_entry| !tombstones.contains(live_entry.0),
704                    |entry| {
705                        let should_stop = match entry {
706                            OrderedOverlayEntry::Canonical(canonical_entry) => {
707                                visit_index_store_entry(
708                                    canonical_entry.key(),
709                                    &canonical_entry.value(),
710                                    &mut visit,
711                                )?
712                            }
713                            OrderedOverlayEntry::Live((key, value)) => {
714                                visit_index_store_entry(key, value, &mut visit)?
715                            }
716                        };
717                        Ok(if should_stop {
718                            OrderedOverlayVisit::Stop
719                        } else {
720                            OrderedOverlayVisit::Continue
721                        })
722                    },
723                )?;
724            }
725        }
726
727        Ok(())
728    }
729}
730
731#[cfg(test)]
732mod tests {
733    use super::*;
734    use crate::{
735        db::{
736            direction::Direction,
737            index::{IndexId, IndexKey, IndexKeyKind},
738            key_taxonomy::{PrimaryKeyComponent, PrimaryKeyValue},
739        },
740        testing::test_memory,
741        types::EntityTag,
742    };
743    use ic_stable_structures::Storable;
744    use std::{borrow::Cow, convert::Infallible};
745
746    fn raw_key(value: u8) -> RawIndexStoreKey {
747        <RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(vec![value]))
748    }
749
750    fn indexed_raw_key(
751        index_id: &IndexId,
752        components: Vec<Vec<u8>>,
753        primary_key: u64,
754    ) -> RawIndexStoreKey {
755        indexed_raw_key_with_kind(index_id, IndexKeyKind::User, components, primary_key)
756    }
757
758    fn indexed_raw_key_with_kind(
759        index_id: &IndexId,
760        key_kind: IndexKeyKind,
761        components: Vec<Vec<u8>>,
762        primary_key: u64,
763    ) -> RawIndexStoreKey {
764        IndexKey::new_from_components_with_primary_key_value(
765            index_id,
766            key_kind,
767            components.as_slice(),
768            &PrimaryKeyValue::from(PrimaryKeyComponent::Nat64(primary_key)),
769        )
770        .expect("test index key should build")
771        .to_raw()
772        .expect("test index key should encode")
773    }
774
775    fn malformed_index_entry_value() -> IndexEntryValue {
776        <IndexEntryValue as Storable>::from_bytes(Cow::Owned(vec![0xFF]))
777    }
778
779    fn missing_index_entry_value() -> IndexEntryValue {
780        <IndexEntryValue as Storable>::from_bytes(Cow::Owned(vec![1]))
781    }
782
783    #[test]
784    fn index_prefix_cardinality_requires_explicit_data_generation_sync() {
785        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
786        let collection = b"collection-a".to_vec();
787        let draft = b"Draft".to_vec();
788        let review = b"Review".to_vec();
789        let mut store = IndexStore::init_heap();
790
791        store.insert(
792            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
793            IndexEntryValue::presence(),
794        );
795        store.insert(
796            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2),
797            IndexEntryValue::presence(),
798        );
799        store.insert(
800            indexed_raw_key(&index_id, vec![collection.clone(), review.clone()], 3),
801            IndexEntryValue::presence(),
802        );
803
804        assert_eq!(
805            store.exact_prefix_cardinality(
806                0,
807                IndexKeyKind::User,
808                index_id,
809                std::slice::from_ref(&collection),
810            ),
811            None,
812            "raw index mutations must not be trusted until row generation sync is stamped",
813        );
814
815        store.mark_prefix_cardinality_data_generation(7);
816
817        assert_eq!(
818            store.exact_prefix_cardinality(
819                7,
820                IndexKeyKind::User,
821                index_id,
822                std::slice::from_ref(&collection),
823            ),
824            Some(3),
825        );
826        assert_eq!(
827            store.exact_prefix_cardinality(
828                7,
829                IndexKeyKind::User,
830                index_id,
831                &[collection.clone(), draft],
832            ),
833            Some(2),
834        );
835        assert_eq!(
836            store.exact_prefix_cardinality(8, IndexKeyKind::User, index_id, &[collection, review],),
837            None,
838            "row generation drift should force the caller to use the existing-row fallback",
839        );
840    }
841
842    #[test]
843    fn index_prefix_cardinality_enumerates_bounded_child_prefixes() {
844        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
845        let collection = b"collection-a".to_vec();
846        let other_collection = b"collection-b".to_vec();
847        let draft = b"Draft".to_vec();
848        let review = b"Review".to_vec();
849        let published = b"Published".to_vec();
850        let mut store = IndexStore::init_heap();
851
852        store.insert(
853            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
854            IndexEntryValue::presence(),
855        );
856        store.insert(
857            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2),
858            IndexEntryValue::presence(),
859        );
860        store.insert(
861            indexed_raw_key(&index_id, vec![collection.clone(), review.clone()], 3),
862            IndexEntryValue::presence(),
863        );
864        store.insert(
865            indexed_raw_key(
866                &index_id,
867                vec![other_collection.clone(), published.clone()],
868                4,
869            ),
870            IndexEntryValue::presence(),
871        );
872        store.mark_prefix_cardinality_data_generation(7);
873
874        assert_eq!(
875            store.exact_child_prefixes_for_parent_set(
876                7,
877                IndexKeyKind::User,
878                index_id,
879                [std::slice::from_ref(&collection)],
880                4,
881            ),
882            Some(vec![
883                vec![collection.clone(), draft],
884                vec![collection.clone(), review],
885            ]),
886            "child-prefix enumeration should return deterministic unique children under the requested parent",
887        );
888        assert_eq!(
889            store.exact_child_prefixes_for_parent_set(
890                7,
891                IndexKeyKind::User,
892                index_id,
893                [std::slice::from_ref(&other_collection)],
894                4,
895            ),
896            Some(vec![vec![other_collection, published]]),
897            "child-prefix enumeration must stay scoped to the requested parent prefix",
898        );
899        assert_eq!(
900            store.exact_child_prefixes_for_parent_set(
901                8,
902                IndexKeyKind::User,
903                index_id,
904                [std::slice::from_ref(&collection)],
905                4,
906            ),
907            None,
908            "row generation drift should keep child-prefix expansion fail-closed",
909        );
910        assert_eq!(
911            store.exact_child_prefixes_for_parent_set(
912                7,
913                IndexKeyKind::User,
914                index_id,
915                [std::slice::from_ref(&collection)],
916                1,
917            ),
918            None,
919            "over-cap child-prefix expansion should fall back to the existing route",
920        );
921    }
922
923    #[test]
924    fn index_prefix_cardinality_batches_sparse_child_prefixes() {
925        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
926        let collection = b"collection-a".to_vec();
927        let other_collection = b"collection-b".to_vec();
928        let missing_a = b"missing-a".to_vec();
929        let missing_b = b"missing-b".to_vec();
930        let draft = b"Draft".to_vec();
931        let review = b"Review".to_vec();
932        let published = b"Published".to_vec();
933        let mut store = IndexStore::init_heap();
934
935        store.insert(
936            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
937            IndexEntryValue::presence(),
938        );
939        store.insert(
940            indexed_raw_key(&index_id, vec![collection.clone(), review.clone()], 2),
941            IndexEntryValue::presence(),
942        );
943        store.insert(
944            indexed_raw_key(
945                &index_id,
946                vec![other_collection.clone(), published.clone()],
947                3,
948            ),
949            IndexEntryValue::presence(),
950        );
951        store.mark_prefix_cardinality_data_generation(7);
952
953        let parents = [
954            std::slice::from_ref(&missing_a),
955            std::slice::from_ref(&collection),
956            std::slice::from_ref(&missing_b),
957            std::slice::from_ref(&other_collection),
958        ];
959        assert_eq!(
960            store.exact_child_prefixes_for_parent_set(7, IndexKeyKind::User, index_id, parents, 4,),
961            Some(vec![
962                vec![collection.clone(), draft],
963                vec![collection.clone(), review],
964                vec![other_collection.clone(), published],
965            ]),
966            "batched child-prefix enumeration should skip missing sparse parents and return deterministic real children",
967        );
968        assert_eq!(
969            store.exact_child_prefixes_for_parent_set(
970                7,
971                IndexKeyKind::User,
972                index_id,
973                [
974                    std::slice::from_ref(&missing_a),
975                    std::slice::from_ref(&missing_b)
976                ],
977                4,
978            ),
979            Some(Vec::new()),
980            "missing-only sparse parent sets should be proven empty when cardinality is synchronized",
981        );
982        assert_eq!(
983            store.exact_child_prefixes_for_parent_set(
984                7,
985                IndexKeyKind::User,
986                index_id,
987                [
988                    std::slice::from_ref(&collection),
989                    std::slice::from_ref(&other_collection)
990                ],
991                2,
992            ),
993            None,
994            "over-cap sparse parent-set expansion should fail closed",
995        );
996        assert_eq!(
997            store.exact_child_prefixes_for_parent_set(
998                8,
999                IndexKeyKind::User,
1000                index_id,
1001                [std::slice::from_ref(&collection)],
1002                4,
1003            ),
1004            None,
1005            "generation drift should keep batched child-prefix expansion fail-closed",
1006        );
1007    }
1008
1009    #[test]
1010    fn index_prefix_cardinality_ignores_system_index_mutations() {
1011        let user_index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1012        let system_index_id = IndexId::new(EntityTag::new(0xCA7D), 2);
1013        let collection = b"collection-a".to_vec();
1014        let draft = b"Draft".to_vec();
1015        let system_component = b"reverse-edge".to_vec();
1016        let mut store = IndexStore::init_heap();
1017
1018        store.insert(
1019            indexed_raw_key(&user_index_id, vec![collection.clone(), draft.clone()], 1),
1020            IndexEntryValue::presence(),
1021        );
1022        store.mark_prefix_cardinality_data_generation(7);
1023
1024        assert_eq!(
1025            store.exact_prefix_cardinality(
1026                7,
1027                IndexKeyKind::User,
1028                user_index_id,
1029                &[collection.clone(), draft.clone()],
1030            ),
1031            Some(1),
1032        );
1033
1034        let system_key = indexed_raw_key_with_kind(
1035            &system_index_id,
1036            IndexKeyKind::System,
1037            vec![system_component],
1038            1,
1039        );
1040        store.insert(system_key.clone(), IndexEntryValue::presence());
1041        assert_eq!(
1042            store.exact_prefix_cardinality(
1043                7,
1044                IndexKeyKind::User,
1045                user_index_id,
1046                &[collection.clone(), draft.clone()],
1047            ),
1048            Some(1),
1049            "system index writes must not invalidate synchronized user-prefix cardinality",
1050        );
1051
1052        store.remove(&system_key);
1053        assert_eq!(
1054            store.exact_prefix_cardinality(
1055                7,
1056                IndexKeyKind::User,
1057                user_index_id,
1058                &[collection.clone(), draft.clone()],
1059            ),
1060            Some(1),
1061            "system index removals must not invalidate synchronized user-prefix cardinality",
1062        );
1063
1064        let malformed_system_key = indexed_raw_key_with_kind(
1065            &system_index_id,
1066            IndexKeyKind::System,
1067            vec![b"malformed-reverse-edge".to_vec()],
1068            2,
1069        );
1070        store.insert(malformed_system_key.clone(), malformed_index_entry_value());
1071        assert_eq!(
1072            store.exact_prefix_cardinality(
1073                7,
1074                IndexKeyKind::User,
1075                user_index_id,
1076                &[collection.clone(), draft.clone()],
1077            ),
1078            Some(1),
1079            "malformed system index payloads must not invalidate user-prefix cardinality",
1080        );
1081
1082        store.remove(&malformed_system_key);
1083        assert_eq!(
1084            store.exact_prefix_cardinality(
1085                7,
1086                IndexKeyKind::User,
1087                user_index_id,
1088                &[collection.clone(), draft],
1089            ),
1090            Some(1),
1091            "malformed system index removals must not invalidate user-prefix cardinality",
1092        );
1093
1094        let review = b"Review".to_vec();
1095        store.insert(
1096            indexed_raw_key(&user_index_id, vec![collection.clone(), review.clone()], 2),
1097            IndexEntryValue::presence(),
1098        );
1099        assert_eq!(
1100            store.exact_prefix_cardinality(
1101                7,
1102                IndexKeyKind::User,
1103                user_index_id,
1104                &[collection, review]
1105            ),
1106            None,
1107            "user-prefix count changes must still require a fresh row-generation stamp",
1108        );
1109    }
1110
1111    #[test]
1112    fn index_prefix_cardinality_ignores_missing_user_index_mutations() {
1113        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1114        let collection = b"collection-a".to_vec();
1115        let draft = b"Draft".to_vec();
1116        let mut store = IndexStore::init_heap();
1117
1118        store.insert(
1119            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
1120            IndexEntryValue::presence(),
1121        );
1122        store.mark_prefix_cardinality_data_generation(7);
1123
1124        let stale_key = indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2);
1125        store.insert(stale_key.clone(), missing_index_entry_value());
1126        assert_eq!(
1127            store.exact_prefix_cardinality(
1128                7,
1129                IndexKeyKind::User,
1130                index_id,
1131                &[collection.clone(), draft.clone()],
1132            ),
1133            Some(1),
1134            "missing user index entries must not affect synchronized prefix cardinality",
1135        );
1136
1137        store.remove(&stale_key);
1138        assert_eq!(
1139            store.exact_prefix_cardinality(7, IndexKeyKind::User, index_id, &[collection, draft],),
1140            Some(1),
1141            "missing user index removals must not affect synchronized prefix cardinality",
1142        );
1143    }
1144
1145    #[cfg(all(feature = "sql", feature = "diagnostics"))]
1146    #[test]
1147    fn index_store_diagnostic_counters_record_gets_range_scans_and_entry_reads() {
1148        let mut store = IndexStore::init_heap();
1149        store.insert(raw_key(7), IndexEntryValue::presence());
1150        store.insert(raw_key(9), IndexEntryValue::presence());
1151
1152        let gets_before = IndexStore::current_get_call_count();
1153        assert_eq!(store.get(&raw_key(7)), Some(IndexEntryValue::presence()));
1154        assert_eq!(store.get(&raw_key(8)), None);
1155
1156        assert_eq!(
1157            IndexStore::current_get_call_count().saturating_sub(gets_before),
1158            2,
1159            "diagnostic index-store get counter should count both hit and miss reads",
1160        );
1161
1162        let range_scans_before = IndexStore::current_range_scan_call_count();
1163        let lower = Bound::Included(raw_key(7));
1164        let upper = Bound::Included(raw_key(9));
1165        store
1166            .visit_raw_entries_in_range((&lower, &upper), Direction::Asc, |_key, _entry| Ok(false))
1167            .expect("raw index range visit should succeed");
1168
1169        assert_eq!(
1170            IndexStore::current_range_scan_call_count().saturating_sub(range_scans_before),
1171            1,
1172            "diagnostic index-store range-scan counter should count one range traversal probe",
1173        );
1174
1175        let entries_before = IndexStore::current_entry_read_count();
1176        store
1177            .visit_entries(|_key, _entry| Ok::<_, Infallible>(IndexStoreVisit::Continue))
1178            .expect("index entry visit should succeed");
1179
1180        assert_eq!(
1181            IndexStore::current_entry_read_count().saturating_sub(entries_before),
1182            2,
1183            "diagnostic index-store entry counter should count yielded traversal entries",
1184        );
1185    }
1186
1187    #[test]
1188    fn journaled_mixed_index_range_traversal_streams_without_snapshot() {
1189        let mut store = IndexStore::init_journaled(test_memory(93));
1190        for value in [1_u8, 3, 5] {
1191            store.insert(raw_key(value), IndexEntryValue::presence());
1192        }
1193        store
1194            .fold_journaled_materialized_view()
1195            .expect("canonical index seed should fold");
1196
1197        store.insert(raw_key(0), IndexEntryValue::presence());
1198        store.insert(raw_key(4), IndexEntryValue::presence());
1199        store.insert(raw_key(5), IndexEntryValue::presence());
1200        store.remove(&raw_key(1));
1201
1202        let lower = Bound::Included(raw_key(0));
1203        let upper = Bound::Included(raw_key(5));
1204
1205        reset_journaled_snapshot_call_count_for_tests();
1206        let mut asc = Vec::new();
1207        store
1208            .visit_journaled_entries_in_range((&lower, &upper), Direction::Asc, |key, _value| {
1209                asc.push(key.as_bytes()[0]);
1210                Ok::<_, Infallible>(asc.len() == 2)
1211            })
1212            .expect("asc journaled index range traversal should succeed");
1213        assert_eq!(asc, vec![0, 3]);
1214        assert_eq!(
1215            journaled_snapshot_call_count_for_tests(),
1216            0,
1217            "mixed journaled index range traversal should preserve early stop without materializing a snapshot",
1218        );
1219
1220        reset_journaled_snapshot_call_count_for_tests();
1221        let mut desc = Vec::new();
1222        store
1223            .visit_journaled_entries_in_range((&lower, &upper), Direction::Desc, |key, _value| {
1224                desc.push(key.as_bytes()[0]);
1225                Ok::<_, Infallible>(desc.len() == 2)
1226            })
1227            .expect("desc journaled index range traversal should succeed");
1228        assert_eq!(desc, vec![5, 4]);
1229        assert_eq!(
1230            journaled_snapshot_call_count_for_tests(),
1231            0,
1232            "mixed reverse journaled index range traversal should preserve early stop without materializing a snapshot",
1233        );
1234    }
1235
1236    #[test]
1237    fn journaled_index_store_reopens_without_materializing_prefix_cardinality() {
1238        let memory = test_memory(94);
1239        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1240        let collection = b"collection-a".to_vec();
1241        let mut store = IndexStore::init_journaled(memory.clone());
1242        let key = indexed_raw_key(&index_id, vec![collection.clone()], 1);
1243        store.insert(key.clone(), IndexEntryValue::presence());
1244        store
1245            .fold_journaled_materialized_view()
1246            .expect("canonical index seed should fold");
1247        drop(store);
1248
1249        let reopened = IndexStore::init_journaled(memory);
1250
1251        assert_eq!(reopened.get(&key), Some(IndexEntryValue::presence()));
1252        assert_eq!(
1253            reopened.exact_prefix_cardinality(
1254                0,
1255                IndexKeyKind::User,
1256                index_id,
1257                std::slice::from_ref(&collection),
1258            ),
1259            None,
1260            "startup must leave optional prefix cardinality unavailable without scanning stable entries",
1261        );
1262    }
1263
1264    #[test]
1265    fn empty_journaled_index_store_retains_exact_prefix_cardinality_without_scanning() {
1266        let memory = test_memory(95);
1267        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1268        let collection = b"collection-a".to_vec();
1269        let mut store = IndexStore::init_journaled(memory.clone());
1270
1271        assert_eq!(
1272            store.exact_prefix_cardinality(
1273                0,
1274                IndexKeyKind::User,
1275                index_id,
1276                std::slice::from_ref(&collection),
1277            ),
1278            Some(0),
1279        );
1280        store
1281            .reset_journaled_live_projection(7)
1282            .expect("empty projection reset should succeed");
1283        assert_eq!(
1284            store.exact_prefix_cardinality(
1285                7,
1286                IndexKeyKind::User,
1287                index_id,
1288                std::slice::from_ref(&collection),
1289            ),
1290            Some(0),
1291        );
1292        drop(store);
1293
1294        let reopened = IndexStore::init_journaled(memory);
1295        assert_eq!(
1296            reopened.exact_prefix_cardinality(
1297                0,
1298                IndexKeyKind::User,
1299                index_id,
1300                std::slice::from_ref(&collection),
1301            ),
1302            Some(0),
1303        );
1304    }
1305
1306    #[test]
1307    fn recovered_index_fold_maintains_available_prefix_cardinality() {
1308        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1309        let collection = b"collection-a".to_vec();
1310        let key = indexed_raw_key(&index_id, vec![collection.clone()], 1);
1311        let mut store = IndexStore::init_journaled(test_memory(96));
1312
1313        store
1314            .fold_recovered_journal_entry(key.clone(), Some(IndexEntryValue::presence()))
1315            .expect("recovered index put should fold");
1316        store.mark_prefix_cardinality_data_generation(1);
1317        assert_eq!(
1318            store.exact_prefix_cardinality(
1319                1,
1320                IndexKeyKind::User,
1321                index_id,
1322                std::slice::from_ref(&collection),
1323            ),
1324            Some(1),
1325        );
1326
1327        store
1328            .fold_recovered_journal_entry(key, None)
1329            .expect("recovered index delete should fold");
1330        store.mark_prefix_cardinality_data_generation(2);
1331        assert_eq!(
1332            store.exact_prefix_cardinality(
1333                2,
1334                IndexKeyKind::User,
1335                index_id,
1336                std::slice::from_ref(&collection),
1337            ),
1338            Some(0),
1339        );
1340    }
1341}
icydb_core/db/index/store.rs

icydb_core/db/index/
store.rs
