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

icydb_core/db/index/
store.rs
