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

icydb_core/db/index/
store.rs
