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    /// Prove that recovered journal entries can be folded into canonical storage.
455    pub(in crate::db) fn preflight_fold_recovered_journal(
456        &self,
457    ) -> Result<(), crate::error::InternalError> {
458        match self.backend {
459            IndexStoreBackend::Journaled { .. } => Ok(()),
460            IndexStoreBackend::Heap(_) => Err(crate::error::InternalError::store_invariant()),
461        }
462    }
463
464    pub fn clear(&mut self) {
465        match &mut self.backend {
466            IndexStoreBackend::Heap(map) => map.clear(),
467            IndexStoreBackend::Journaled {
468                canonical,
469                live,
470                tombstones,
471            } => {
472                live.clear();
473                tombstones.clear();
474                for entry in canonical.iter() {
475                    tombstones.insert(entry.key().clone());
476                }
477            }
478        }
479        self.prefix_cardinality.clear_unsynchronized();
480        self.bump_generation();
481    }
482
483    /// Fold the current journaled materialized index view into the canonical
484    /// stable base and clear volatile projection state.
485    #[cfg(any(test, feature = "migration"))]
486    pub(in crate::db) fn fold_journaled_materialized_view(
487        &mut self,
488    ) -> Result<(), crate::error::InternalError> {
489        let entries = Self::journaled_entries_snapshot_for_fold(&self.backend);
490        let IndexStoreBackend::Journaled {
491            canonical,
492            live,
493            tombstones,
494        } = &mut self.backend
495        else {
496            return Err(crate::error::InternalError::store_invariant());
497        };
498
499        canonical.clear_new();
500        for (key, value) in entries {
501            canonical.insert(key, value);
502        }
503        live.clear();
504        tombstones.clear();
505        let data_generation = self.prefix_cardinality.synchronized_generation();
506        self.rebuild_prefix_cardinality_from_entries(data_generation);
507        self.bump_generation();
508
509        Ok(())
510    }
511
512    /// Sum of bytes used by all stored index entries.
513    pub fn memory_bytes(&self) -> u64 {
514        let mut bytes = 0u64;
515        let _: Result<(), std::convert::Infallible> = self.visit_entries(|key, value| {
516            bytes = bytes.saturating_add(key.as_bytes().len() as u64 + value.len() as u64);
517            Ok(IndexStoreVisit::Continue)
518        });
519        bytes
520    }
521
522    /// Return the monotonic perf-only count of index-entry fetches seen by this process.
523    #[cfg(all(feature = "sql", feature = "diagnostics"))]
524    pub(in crate::db) fn current_get_call_count() -> u64 {
525        INDEX_STORE_GET_CALL_COUNT.with(Cell::get)
526    }
527
528    /// Return the monotonic perf-only count of index range traversal probes seen by this process.
529    #[cfg(all(feature = "sql", feature = "diagnostics"))]
530    pub(in crate::db) fn current_range_scan_call_count() -> u64 {
531        INDEX_STORE_RANGE_SCAN_CALL_COUNT.with(Cell::get)
532    }
533
534    /// Return the monotonic perf-only count of index entries yielded by traversal.
535    #[cfg(all(feature = "sql", feature = "diagnostics"))]
536    pub(in crate::db) fn current_entry_read_count() -> u64 {
537        INDEX_STORE_ENTRY_READ_COUNT.with(Cell::get)
538    }
539
540    #[cfg(all(feature = "sql", feature = "diagnostics"))]
541    pub(in crate::db::index) fn record_range_scan_call() {
542        record_index_store_range_scan_call();
543    }
544
545    #[cfg(all(feature = "sql", feature = "diagnostics"))]
546    pub(in crate::db::index) fn record_merged_entry_reads(count: u64) {
547        INDEX_STORE_ENTRY_READ_COUNT.with(|total| {
548            total.set(total.get().saturating_add(count));
549        });
550    }
551
552    const fn bump_generation(&mut self) {
553        self.generation = self.generation.saturating_add(1);
554    }
555
556    #[cfg(any(test, feature = "migration"))]
557    fn rebuild_prefix_cardinality_from_entries(&mut self, data_generation: Option<u64>) {
558        self.prefix_cardinality.clear_unsynchronized();
559        let entries = Self::entries_snapshot_for_cardinality(&self.backend);
560        for (key, value) in &entries {
561            self.prefix_cardinality.apply_insert(key, None, value);
562        }
563        if let Some(data_generation) = data_generation {
564            self.prefix_cardinality.mark_synchronized(data_generation);
565        }
566    }
567
568    #[cfg(any(test, feature = "migration"))]
569    fn entries_snapshot_for_cardinality(
570        backend: &IndexStoreBackend,
571    ) -> HeapBTreeMap<RawIndexStoreKey, IndexEntryValue> {
572        match backend {
573            IndexStoreBackend::Heap(map) => map.clone(),
574            IndexStoreBackend::Journaled { .. } => {
575                Self::journaled_entries_snapshot_for_fold(backend)
576            }
577        }
578    }
579
580    fn journaled_get(
581        backend: &IndexStoreBackend,
582        key: &RawIndexStoreKey,
583    ) -> Option<IndexEntryValue> {
584        let IndexStoreBackend::Journaled {
585            canonical,
586            live,
587            tombstones,
588        } = backend
589        else {
590            return None;
591        };
592
593        if tombstones.contains(key) {
594            return None;
595        }
596        live.get(key).cloned().or_else(|| canonical.get(key))
597    }
598
599    #[cfg(any(test, feature = "migration"))]
600    pub(super) fn journaled_entries_snapshot_for_fold(
601        backend: &IndexStoreBackend,
602    ) -> HeapBTreeMap<RawIndexStoreKey, IndexEntryValue> {
603        #[cfg(test)]
604        record_journaled_snapshot_call();
605
606        let IndexStoreBackend::Journaled {
607            canonical,
608            live,
609            tombstones,
610        } = backend
611        else {
612            return HeapBTreeMap::new();
613        };
614
615        let mut entries = HeapBTreeMap::new();
616        for entry in canonical.iter() {
617            let key = entry.key().clone();
618            if !tombstones.contains(&key) {
619                entries.insert(key, entry.value());
620            }
621        }
622        for (key, value) in live {
623            if !tombstones.contains(key) {
624                entries.insert(key.clone(), value.clone());
625            }
626        }
627
628        entries
629    }
630
631    pub(super) fn visit_journaled_entries_in_range<E>(
632        &self,
633        bounds: (&Bound<RawIndexStoreKey>, &Bound<RawIndexStoreKey>),
634        direction: Direction,
635        mut visit: impl FnMut(&RawIndexStoreKey, &IndexEntryValue) -> Result<bool, E>,
636    ) -> Result<(), E> {
637        let IndexStoreBackend::Journaled {
638            canonical,
639            live,
640            tombstones,
641        } = &self.backend
642        else {
643            return Ok(());
644        };
645
646        let lower = bounds.0.clone();
647        let upper = bounds.1.clone();
648        match direction {
649            Direction::Asc if canonical.is_empty() => {
650                for (key, value) in live.range((lower, upper)) {
651                    if visit_index_store_entry(key, value, &mut visit)? {
652                        return Ok(());
653                    }
654                }
655            }
656            Direction::Desc if canonical.is_empty() => {
657                for (key, value) in live.range((lower, upper)).rev() {
658                    if visit_index_store_entry(key, value, &mut visit)? {
659                        return Ok(());
660                    }
661                }
662            }
663            Direction::Asc if live.is_empty() && tombstones.is_empty() => {
664                for entry in canonical.range((lower, upper)) {
665                    if visit_index_store_entry(entry.key(), &entry.value(), &mut visit)? {
666                        return Ok(());
667                    }
668                }
669            }
670            Direction::Desc if live.is_empty() && tombstones.is_empty() => {
671                for entry in canonical.range((lower, upper)).rev() {
672                    if visit_index_store_entry(entry.key(), &entry.value(), &mut visit)? {
673                        return Ok(());
674                    }
675                }
676            }
677            Direction::Asc => {
678                visit_ordered_overlay(
679                    canonical.range((lower.clone(), upper.clone())),
680                    live.range((lower, upper)),
681                    direction,
682                    |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
683                    |canonical_entry| !tombstones.contains(canonical_entry.key()),
684                    |live_entry| !tombstones.contains(live_entry.0),
685                    |entry| {
686                        let should_stop = match entry {
687                            OrderedOverlayEntry::Canonical(canonical_entry) => {
688                                visit_index_store_entry(
689                                    canonical_entry.key(),
690                                    &canonical_entry.value(),
691                                    &mut visit,
692                                )?
693                            }
694                            OrderedOverlayEntry::Live((key, value)) => {
695                                visit_index_store_entry(key, value, &mut visit)?
696                            }
697                        };
698                        Ok(if should_stop {
699                            OrderedOverlayVisit::Stop
700                        } else {
701                            OrderedOverlayVisit::Continue
702                        })
703                    },
704                )?;
705            }
706            Direction::Desc => {
707                visit_ordered_overlay(
708                    canonical.range((lower.clone(), upper.clone())).rev(),
709                    live.range((lower, upper)).rev(),
710                    direction,
711                    |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
712                    |canonical_entry| !tombstones.contains(canonical_entry.key()),
713                    |live_entry| !tombstones.contains(live_entry.0),
714                    |entry| {
715                        let should_stop = match entry {
716                            OrderedOverlayEntry::Canonical(canonical_entry) => {
717                                visit_index_store_entry(
718                                    canonical_entry.key(),
719                                    &canonical_entry.value(),
720                                    &mut visit,
721                                )?
722                            }
723                            OrderedOverlayEntry::Live((key, value)) => {
724                                visit_index_store_entry(key, value, &mut visit)?
725                            }
726                        };
727                        Ok(if should_stop {
728                            OrderedOverlayVisit::Stop
729                        } else {
730                            OrderedOverlayVisit::Continue
731                        })
732                    },
733                )?;
734            }
735        }
736
737        Ok(())
738    }
739}
740
741#[cfg(test)]
742mod tests {
743    use super::*;
744    use crate::{
745        db::{
746            direction::Direction,
747            index::{IndexId, IndexKey, IndexKeyKind},
748            key_taxonomy::{PrimaryKeyComponent, PrimaryKeyValue},
749        },
750        testing::test_memory,
751        types::EntityTag,
752    };
753    use ic_stable_structures::Storable;
754    use std::{borrow::Cow, convert::Infallible};
755
756    fn raw_key(value: u8) -> RawIndexStoreKey {
757        <RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(vec![value]))
758    }
759
760    fn indexed_raw_key(
761        index_id: &IndexId,
762        components: Vec<Vec<u8>>,
763        primary_key: u64,
764    ) -> RawIndexStoreKey {
765        indexed_raw_key_with_kind(index_id, IndexKeyKind::User, components, primary_key)
766    }
767
768    fn indexed_raw_key_with_kind(
769        index_id: &IndexId,
770        key_kind: IndexKeyKind,
771        components: Vec<Vec<u8>>,
772        primary_key: u64,
773    ) -> RawIndexStoreKey {
774        IndexKey::new_from_components_with_primary_key_value(
775            index_id,
776            key_kind,
777            components.as_slice(),
778            &PrimaryKeyValue::from(PrimaryKeyComponent::Nat64(primary_key)),
779        )
780        .expect("test index key should build")
781        .to_raw()
782        .expect("test index key should encode")
783    }
784
785    fn malformed_index_entry_value() -> IndexEntryValue {
786        <IndexEntryValue as Storable>::from_bytes(Cow::Owned(vec![0xFF]))
787    }
788
789    fn missing_index_entry_value() -> IndexEntryValue {
790        <IndexEntryValue as Storable>::from_bytes(Cow::Owned(vec![1]))
791    }
792
793    #[test]
794    fn index_prefix_cardinality_requires_explicit_data_generation_sync() {
795        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
796        let collection = b"collection-a".to_vec();
797        let draft = b"Draft".to_vec();
798        let review = b"Review".to_vec();
799        let mut store = IndexStore::init_heap();
800
801        store.insert(
802            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
803            IndexEntryValue::presence(),
804        );
805        store.insert(
806            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2),
807            IndexEntryValue::presence(),
808        );
809        store.insert(
810            indexed_raw_key(&index_id, vec![collection.clone(), review.clone()], 3),
811            IndexEntryValue::presence(),
812        );
813
814        assert_eq!(
815            store.exact_prefix_cardinality(
816                0,
817                IndexKeyKind::User,
818                index_id,
819                std::slice::from_ref(&collection),
820            ),
821            None,
822            "raw index mutations must not be trusted until row generation sync is stamped",
823        );
824
825        store.mark_prefix_cardinality_data_generation(7);
826
827        assert_eq!(
828            store.exact_prefix_cardinality(
829                7,
830                IndexKeyKind::User,
831                index_id,
832                std::slice::from_ref(&collection),
833            ),
834            Some(3),
835        );
836        assert_eq!(
837            store.exact_prefix_cardinality(
838                7,
839                IndexKeyKind::User,
840                index_id,
841                &[collection.clone(), draft],
842            ),
843            Some(2),
844        );
845        assert_eq!(
846            store.exact_prefix_cardinality(8, IndexKeyKind::User, index_id, &[collection, review],),
847            None,
848            "row generation drift should force the caller to use the existing-row fallback",
849        );
850    }
851
852    #[test]
853    fn index_prefix_cardinality_enumerates_bounded_child_prefixes() {
854        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
855        let collection = b"collection-a".to_vec();
856        let other_collection = b"collection-b".to_vec();
857        let draft = b"Draft".to_vec();
858        let review = b"Review".to_vec();
859        let published = b"Published".to_vec();
860        let mut store = IndexStore::init_heap();
861
862        store.insert(
863            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
864            IndexEntryValue::presence(),
865        );
866        store.insert(
867            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2),
868            IndexEntryValue::presence(),
869        );
870        store.insert(
871            indexed_raw_key(&index_id, vec![collection.clone(), review.clone()], 3),
872            IndexEntryValue::presence(),
873        );
874        store.insert(
875            indexed_raw_key(
876                &index_id,
877                vec![other_collection.clone(), published.clone()],
878                4,
879            ),
880            IndexEntryValue::presence(),
881        );
882        store.mark_prefix_cardinality_data_generation(7);
883
884        assert_eq!(
885            store.exact_child_prefixes_for_parent_set(
886                7,
887                IndexKeyKind::User,
888                index_id,
889                [std::slice::from_ref(&collection)],
890                4,
891            ),
892            Some(vec![
893                vec![collection.clone(), draft],
894                vec![collection.clone(), review],
895            ]),
896            "child-prefix enumeration should return deterministic unique children under the requested parent",
897        );
898        assert_eq!(
899            store.exact_child_prefixes_for_parent_set(
900                7,
901                IndexKeyKind::User,
902                index_id,
903                [std::slice::from_ref(&other_collection)],
904                4,
905            ),
906            Some(vec![vec![other_collection, published]]),
907            "child-prefix enumeration must stay scoped to the requested parent prefix",
908        );
909        assert_eq!(
910            store.exact_child_prefixes_for_parent_set(
911                8,
912                IndexKeyKind::User,
913                index_id,
914                [std::slice::from_ref(&collection)],
915                4,
916            ),
917            None,
918            "row generation drift should keep child-prefix expansion fail-closed",
919        );
920        assert_eq!(
921            store.exact_child_prefixes_for_parent_set(
922                7,
923                IndexKeyKind::User,
924                index_id,
925                [std::slice::from_ref(&collection)],
926                1,
927            ),
928            None,
929            "over-cap child-prefix expansion should fall back to the existing route",
930        );
931    }
932
933    #[test]
934    fn index_prefix_cardinality_batches_sparse_child_prefixes() {
935        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
936        let collection = b"collection-a".to_vec();
937        let other_collection = b"collection-b".to_vec();
938        let missing_a = b"missing-a".to_vec();
939        let missing_b = b"missing-b".to_vec();
940        let draft = b"Draft".to_vec();
941        let review = b"Review".to_vec();
942        let published = b"Published".to_vec();
943        let mut store = IndexStore::init_heap();
944
945        store.insert(
946            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
947            IndexEntryValue::presence(),
948        );
949        store.insert(
950            indexed_raw_key(&index_id, vec![collection.clone(), review.clone()], 2),
951            IndexEntryValue::presence(),
952        );
953        store.insert(
954            indexed_raw_key(
955                &index_id,
956                vec![other_collection.clone(), published.clone()],
957                3,
958            ),
959            IndexEntryValue::presence(),
960        );
961        store.mark_prefix_cardinality_data_generation(7);
962
963        let parents = [
964            std::slice::from_ref(&missing_a),
965            std::slice::from_ref(&collection),
966            std::slice::from_ref(&missing_b),
967            std::slice::from_ref(&other_collection),
968        ];
969        assert_eq!(
970            store.exact_child_prefixes_for_parent_set(7, IndexKeyKind::User, index_id, parents, 4,),
971            Some(vec![
972                vec![collection.clone(), draft],
973                vec![collection.clone(), review],
974                vec![other_collection.clone(), published],
975            ]),
976            "batched child-prefix enumeration should skip missing sparse parents and return deterministic real children",
977        );
978        assert_eq!(
979            store.exact_child_prefixes_for_parent_set(
980                7,
981                IndexKeyKind::User,
982                index_id,
983                [
984                    std::slice::from_ref(&missing_a),
985                    std::slice::from_ref(&missing_b)
986                ],
987                4,
988            ),
989            Some(Vec::new()),
990            "missing-only sparse parent sets should be proven empty when cardinality is synchronized",
991        );
992        assert_eq!(
993            store.exact_child_prefixes_for_parent_set(
994                7,
995                IndexKeyKind::User,
996                index_id,
997                [
998                    std::slice::from_ref(&collection),
999                    std::slice::from_ref(&other_collection)
1000                ],
1001                2,
1002            ),
1003            None,
1004            "over-cap sparse parent-set expansion should fail closed",
1005        );
1006        assert_eq!(
1007            store.exact_child_prefixes_for_parent_set(
1008                8,
1009                IndexKeyKind::User,
1010                index_id,
1011                [std::slice::from_ref(&collection)],
1012                4,
1013            ),
1014            None,
1015            "generation drift should keep batched child-prefix expansion fail-closed",
1016        );
1017    }
1018
1019    #[test]
1020    fn index_prefix_cardinality_ignores_system_index_mutations() {
1021        let user_index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1022        let system_index_id = IndexId::new(EntityTag::new(0xCA7D), 2);
1023        let collection = b"collection-a".to_vec();
1024        let draft = b"Draft".to_vec();
1025        let system_component = b"reverse-edge".to_vec();
1026        let mut store = IndexStore::init_heap();
1027
1028        store.insert(
1029            indexed_raw_key(&user_index_id, vec![collection.clone(), draft.clone()], 1),
1030            IndexEntryValue::presence(),
1031        );
1032        store.mark_prefix_cardinality_data_generation(7);
1033
1034        assert_eq!(
1035            store.exact_prefix_cardinality(
1036                7,
1037                IndexKeyKind::User,
1038                user_index_id,
1039                &[collection.clone(), draft.clone()],
1040            ),
1041            Some(1),
1042        );
1043
1044        let system_key = indexed_raw_key_with_kind(
1045            &system_index_id,
1046            IndexKeyKind::System,
1047            vec![system_component],
1048            1,
1049        );
1050        store.insert(system_key.clone(), IndexEntryValue::presence());
1051        assert_eq!(
1052            store.exact_prefix_cardinality(
1053                7,
1054                IndexKeyKind::User,
1055                user_index_id,
1056                &[collection.clone(), draft.clone()],
1057            ),
1058            Some(1),
1059            "system index writes must not invalidate synchronized user-prefix cardinality",
1060        );
1061
1062        store.remove(&system_key);
1063        assert_eq!(
1064            store.exact_prefix_cardinality(
1065                7,
1066                IndexKeyKind::User,
1067                user_index_id,
1068                &[collection.clone(), draft.clone()],
1069            ),
1070            Some(1),
1071            "system index removals must not invalidate synchronized user-prefix cardinality",
1072        );
1073
1074        let malformed_system_key = indexed_raw_key_with_kind(
1075            &system_index_id,
1076            IndexKeyKind::System,
1077            vec![b"malformed-reverse-edge".to_vec()],
1078            2,
1079        );
1080        store.insert(malformed_system_key.clone(), malformed_index_entry_value());
1081        assert_eq!(
1082            store.exact_prefix_cardinality(
1083                7,
1084                IndexKeyKind::User,
1085                user_index_id,
1086                &[collection.clone(), draft.clone()],
1087            ),
1088            Some(1),
1089            "malformed system index payloads must not invalidate user-prefix cardinality",
1090        );
1091
1092        store.remove(&malformed_system_key);
1093        assert_eq!(
1094            store.exact_prefix_cardinality(
1095                7,
1096                IndexKeyKind::User,
1097                user_index_id,
1098                &[collection.clone(), draft],
1099            ),
1100            Some(1),
1101            "malformed system index removals must not invalidate user-prefix cardinality",
1102        );
1103
1104        let review = b"Review".to_vec();
1105        store.insert(
1106            indexed_raw_key(&user_index_id, vec![collection.clone(), review.clone()], 2),
1107            IndexEntryValue::presence(),
1108        );
1109        assert_eq!(
1110            store.exact_prefix_cardinality(
1111                7,
1112                IndexKeyKind::User,
1113                user_index_id,
1114                &[collection, review]
1115            ),
1116            None,
1117            "user-prefix count changes must still require a fresh row-generation stamp",
1118        );
1119    }
1120
1121    #[test]
1122    fn index_prefix_cardinality_ignores_missing_user_index_mutations() {
1123        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1124        let collection = b"collection-a".to_vec();
1125        let draft = b"Draft".to_vec();
1126        let mut store = IndexStore::init_heap();
1127
1128        store.insert(
1129            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
1130            IndexEntryValue::presence(),
1131        );
1132        store.mark_prefix_cardinality_data_generation(7);
1133
1134        let stale_key = indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2);
1135        store.insert(stale_key.clone(), missing_index_entry_value());
1136        assert_eq!(
1137            store.exact_prefix_cardinality(
1138                7,
1139                IndexKeyKind::User,
1140                index_id,
1141                &[collection.clone(), draft.clone()],
1142            ),
1143            Some(1),
1144            "missing user index entries must not affect synchronized prefix cardinality",
1145        );
1146
1147        store.remove(&stale_key);
1148        assert_eq!(
1149            store.exact_prefix_cardinality(7, IndexKeyKind::User, index_id, &[collection, draft],),
1150            Some(1),
1151            "missing user index removals must not affect synchronized prefix cardinality",
1152        );
1153    }
1154
1155    #[cfg(all(feature = "sql", feature = "diagnostics"))]
1156    #[test]
1157    fn index_store_diagnostic_counters_record_gets_range_scans_and_entry_reads() {
1158        let mut store = IndexStore::init_heap();
1159        store.insert(raw_key(7), IndexEntryValue::presence());
1160        store.insert(raw_key(9), IndexEntryValue::presence());
1161
1162        let gets_before = IndexStore::current_get_call_count();
1163        assert_eq!(store.get(&raw_key(7)), Some(IndexEntryValue::presence()));
1164        assert_eq!(store.get(&raw_key(8)), None);
1165
1166        assert_eq!(
1167            IndexStore::current_get_call_count().saturating_sub(gets_before),
1168            2,
1169            "diagnostic index-store get counter should count both hit and miss reads",
1170        );
1171
1172        let range_scans_before = IndexStore::current_range_scan_call_count();
1173        let lower = Bound::Included(raw_key(7));
1174        let upper = Bound::Included(raw_key(9));
1175        store
1176            .visit_raw_entries_in_range((&lower, &upper), Direction::Asc, |_key, _entry| Ok(false))
1177            .expect("raw index range visit should succeed");
1178
1179        assert_eq!(
1180            IndexStore::current_range_scan_call_count().saturating_sub(range_scans_before),
1181            1,
1182            "diagnostic index-store range-scan counter should count one range traversal probe",
1183        );
1184
1185        let entries_before = IndexStore::current_entry_read_count();
1186        store
1187            .visit_entries(|_key, _entry| Ok::<_, Infallible>(IndexStoreVisit::Continue))
1188            .expect("index entry visit should succeed");
1189
1190        assert_eq!(
1191            IndexStore::current_entry_read_count().saturating_sub(entries_before),
1192            2,
1193            "diagnostic index-store entry counter should count yielded traversal entries",
1194        );
1195    }
1196
1197    #[test]
1198    fn journaled_mixed_index_range_traversal_streams_without_snapshot() {
1199        let mut store = IndexStore::init_journaled(test_memory(93));
1200        for value in [1_u8, 3, 5] {
1201            store.insert(raw_key(value), IndexEntryValue::presence());
1202        }
1203        store
1204            .fold_journaled_materialized_view()
1205            .expect("canonical index seed should fold");
1206
1207        store.insert(raw_key(0), IndexEntryValue::presence());
1208        store.insert(raw_key(4), IndexEntryValue::presence());
1209        store.insert(raw_key(5), IndexEntryValue::presence());
1210        store.remove(&raw_key(1));
1211
1212        let lower = Bound::Included(raw_key(0));
1213        let upper = Bound::Included(raw_key(5));
1214
1215        reset_journaled_snapshot_call_count_for_tests();
1216        let mut asc = Vec::new();
1217        store
1218            .visit_journaled_entries_in_range((&lower, &upper), Direction::Asc, |key, _value| {
1219                asc.push(key.as_bytes()[0]);
1220                Ok::<_, Infallible>(asc.len() == 2)
1221            })
1222            .expect("asc journaled index range traversal should succeed");
1223        assert_eq!(asc, vec![0, 3]);
1224        assert_eq!(
1225            journaled_snapshot_call_count_for_tests(),
1226            0,
1227            "mixed journaled index range traversal should preserve early stop without materializing a snapshot",
1228        );
1229
1230        reset_journaled_snapshot_call_count_for_tests();
1231        let mut desc = Vec::new();
1232        store
1233            .visit_journaled_entries_in_range((&lower, &upper), Direction::Desc, |key, _value| {
1234                desc.push(key.as_bytes()[0]);
1235                Ok::<_, Infallible>(desc.len() == 2)
1236            })
1237            .expect("desc journaled index range traversal should succeed");
1238        assert_eq!(desc, vec![5, 4]);
1239        assert_eq!(
1240            journaled_snapshot_call_count_for_tests(),
1241            0,
1242            "mixed reverse journaled index range traversal should preserve early stop without materializing a snapshot",
1243        );
1244    }
1245
1246    #[test]
1247    fn journaled_index_store_reopens_without_materializing_prefix_cardinality() {
1248        let memory = test_memory(94);
1249        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1250        let collection = b"collection-a".to_vec();
1251        let mut store = IndexStore::init_journaled(memory.clone());
1252        let key = indexed_raw_key(&index_id, vec![collection.clone()], 1);
1253        store.insert(key.clone(), IndexEntryValue::presence());
1254        store
1255            .fold_journaled_materialized_view()
1256            .expect("canonical index seed should fold");
1257        drop(store);
1258
1259        let reopened = IndexStore::init_journaled(memory);
1260
1261        assert_eq!(reopened.get(&key), Some(IndexEntryValue::presence()));
1262        assert_eq!(
1263            reopened.exact_prefix_cardinality(
1264                0,
1265                IndexKeyKind::User,
1266                index_id,
1267                std::slice::from_ref(&collection),
1268            ),
1269            None,
1270            "startup must leave optional prefix cardinality unavailable without scanning stable entries",
1271        );
1272    }
1273
1274    #[test]
1275    fn empty_journaled_index_store_retains_exact_prefix_cardinality_without_scanning() {
1276        let memory = test_memory(95);
1277        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1278        let collection = b"collection-a".to_vec();
1279        let mut store = IndexStore::init_journaled(memory.clone());
1280
1281        assert_eq!(
1282            store.exact_prefix_cardinality(
1283                0,
1284                IndexKeyKind::User,
1285                index_id,
1286                std::slice::from_ref(&collection),
1287            ),
1288            Some(0),
1289        );
1290        store
1291            .reset_journaled_live_projection(7)
1292            .expect("empty projection reset should succeed");
1293        assert_eq!(
1294            store.exact_prefix_cardinality(
1295                7,
1296                IndexKeyKind::User,
1297                index_id,
1298                std::slice::from_ref(&collection),
1299            ),
1300            Some(0),
1301        );
1302        drop(store);
1303
1304        let reopened = IndexStore::init_journaled(memory);
1305        assert_eq!(
1306            reopened.exact_prefix_cardinality(
1307                0,
1308                IndexKeyKind::User,
1309                index_id,
1310                std::slice::from_ref(&collection),
1311            ),
1312            Some(0),
1313        );
1314    }
1315
1316    #[test]
1317    fn recovered_index_fold_maintains_available_prefix_cardinality() {
1318        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
1319        let collection = b"collection-a".to_vec();
1320        let key = indexed_raw_key(&index_id, vec![collection.clone()], 1);
1321        let mut store = IndexStore::init_journaled(test_memory(96));
1322
1323        store
1324            .fold_recovered_journal_entry(key.clone(), Some(IndexEntryValue::presence()))
1325            .expect("recovered index put should fold");
1326        store.mark_prefix_cardinality_data_generation(1);
1327        assert_eq!(
1328            store.exact_prefix_cardinality(
1329                1,
1330                IndexKeyKind::User,
1331                index_id,
1332                std::slice::from_ref(&collection),
1333            ),
1334            Some(1),
1335        );
1336
1337        store
1338            .fold_recovered_journal_entry(key, None)
1339            .expect("recovered index delete should fold");
1340        store.mark_prefix_cardinality_data_generation(2);
1341        assert_eq!(
1342            store.exact_prefix_cardinality(
1343                2,
1344                IndexKeyKind::User,
1345                index_id,
1346                std::slice::from_ref(&collection),
1347            ),
1348            Some(0),
1349        );
1350    }
1351}
icydb_core/db/index/store.rs

icydb_core/db/index/
store.rs
