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    #[allow(dead_code)]
486    pub(in crate::db) fn current_prefix_cardinality_lookup_count() -> u64 {
487        INDEX_STORE_PREFIX_CARDINALITY_LOOKUP_COUNT.with(Cell::get)
488    }
489
490    #[cfg(any(test, feature = "diagnostics"))]
491    pub(in crate::db::index) fn record_range_scan_call() {
492        record_index_store_range_scan_call();
493    }
494
495    const fn bump_generation(&mut self) {
496        self.generation = self.generation.saturating_add(1);
497    }
498
499    fn rebuild_prefix_cardinality_from_entries(&mut self, data_generation: Option<u64>) {
500        self.prefix_cardinality.clear_unsynchronized();
501        let entries = Self::entries_snapshot_for_cardinality(&self.backend);
502        for (key, value) in &entries {
503            self.prefix_cardinality.apply_insert(key, None, value);
504        }
505        if let Some(data_generation) = data_generation {
506            self.prefix_cardinality.mark_synchronized(data_generation);
507        }
508    }
509
510    fn entries_snapshot_for_cardinality(
511        backend: &IndexStoreBackend,
512    ) -> HeapBTreeMap<RawIndexStoreKey, IndexEntryValue> {
513        match backend {
514            IndexStoreBackend::Heap(map) => map.clone(),
515            IndexStoreBackend::Journaled { .. } => {
516                Self::journaled_entries_snapshot_for_fold(backend)
517            }
518        }
519    }
520
521    #[cfg(test)]
522    #[must_use]
523    pub(in crate::db) fn canonical_len_for_tests(&self) -> u64 {
524        match &self.backend {
525            IndexStoreBackend::Journaled { canonical: map, .. } => map.len(),
526            IndexStoreBackend::Heap(_) => 0,
527        }
528    }
529
530    fn journaled_get(
531        backend: &IndexStoreBackend,
532        key: &RawIndexStoreKey,
533    ) -> Option<IndexEntryValue> {
534        let IndexStoreBackend::Journaled {
535            canonical,
536            live,
537            tombstones,
538        } = backend
539        else {
540            return None;
541        };
542
543        if tombstones.contains(key) {
544            return None;
545        }
546        live.get(key).cloned().or_else(|| canonical.get(key))
547    }
548
549    pub(super) fn journaled_entries_snapshot_for_fold(
550        backend: &IndexStoreBackend,
551    ) -> HeapBTreeMap<RawIndexStoreKey, IndexEntryValue> {
552        #[cfg(test)]
553        record_journaled_snapshot_call();
554
555        let IndexStoreBackend::Journaled {
556            canonical,
557            live,
558            tombstones,
559        } = backend
560        else {
561            return HeapBTreeMap::new();
562        };
563
564        let mut entries = HeapBTreeMap::new();
565        for entry in canonical.iter() {
566            let key = entry.key().clone();
567            if !tombstones.contains(&key) {
568                entries.insert(key, entry.value());
569            }
570        }
571        for (key, value) in live {
572            if !tombstones.contains(key) {
573                entries.insert(key.clone(), value.clone());
574            }
575        }
576
577        entries
578    }
579
580    pub(super) fn visit_journaled_entries_in_range<E>(
581        &self,
582        bounds: (&Bound<RawIndexStoreKey>, &Bound<RawIndexStoreKey>),
583        direction: Direction,
584        mut visit: impl FnMut(&RawIndexStoreKey, &IndexEntryValue) -> Result<bool, E>,
585    ) -> Result<(), E> {
586        let IndexStoreBackend::Journaled {
587            canonical,
588            live,
589            tombstones,
590        } = &self.backend
591        else {
592            return Ok(());
593        };
594
595        let lower = bounds.0.clone();
596        let upper = bounds.1.clone();
597        match direction {
598            Direction::Asc if canonical.is_empty() => {
599                for (key, value) in live.range((lower, upper)) {
600                    if visit_index_store_entry(key, value, &mut visit)? {
601                        return Ok(());
602                    }
603                }
604            }
605            Direction::Desc if canonical.is_empty() => {
606                for (key, value) in live.range((lower, upper)).rev() {
607                    if visit_index_store_entry(key, value, &mut visit)? {
608                        return Ok(());
609                    }
610                }
611            }
612            Direction::Asc if live.is_empty() && tombstones.is_empty() => {
613                for entry in canonical.range((lower, upper)) {
614                    if visit_index_store_entry(entry.key(), &entry.value(), &mut visit)? {
615                        return Ok(());
616                    }
617                }
618            }
619            Direction::Desc if live.is_empty() && tombstones.is_empty() => {
620                for entry in canonical.range((lower, upper)).rev() {
621                    if visit_index_store_entry(entry.key(), &entry.value(), &mut visit)? {
622                        return Ok(());
623                    }
624                }
625            }
626            Direction::Asc => {
627                visit_ordered_overlay(
628                    canonical.range((lower.clone(), upper.clone())),
629                    live.range((lower, upper)),
630                    direction,
631                    |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
632                    |canonical_entry| !tombstones.contains(canonical_entry.key()),
633                    |live_entry| !tombstones.contains(live_entry.0),
634                    |entry| {
635                        let should_stop = match entry {
636                            OrderedOverlayEntry::Canonical(canonical_entry) => {
637                                visit_index_store_entry(
638                                    canonical_entry.key(),
639                                    &canonical_entry.value(),
640                                    &mut visit,
641                                )?
642                            }
643                            OrderedOverlayEntry::Live((key, value)) => {
644                                visit_index_store_entry(key, value, &mut visit)?
645                            }
646                        };
647                        Ok(if should_stop {
648                            OrderedOverlayVisit::Stop
649                        } else {
650                            OrderedOverlayVisit::Continue
651                        })
652                    },
653                )?;
654            }
655            Direction::Desc => {
656                visit_ordered_overlay(
657                    canonical.range((lower.clone(), upper.clone())).rev(),
658                    live.range((lower, upper)).rev(),
659                    direction,
660                    |canonical_entry, live_entry| canonical_entry.key().cmp(live_entry.0),
661                    |canonical_entry| !tombstones.contains(canonical_entry.key()),
662                    |live_entry| !tombstones.contains(live_entry.0),
663                    |entry| {
664                        let should_stop = match entry {
665                            OrderedOverlayEntry::Canonical(canonical_entry) => {
666                                visit_index_store_entry(
667                                    canonical_entry.key(),
668                                    &canonical_entry.value(),
669                                    &mut visit,
670                                )?
671                            }
672                            OrderedOverlayEntry::Live((key, value)) => {
673                                visit_index_store_entry(key, value, &mut visit)?
674                            }
675                        };
676                        Ok(if should_stop {
677                            OrderedOverlayVisit::Stop
678                        } else {
679                            OrderedOverlayVisit::Continue
680                        })
681                    },
682                )?;
683            }
684        }
685
686        Ok(())
687    }
688}
689
690#[cfg(test)]
691mod tests {
692    use super::*;
693    use crate::{
694        db::{
695            direction::Direction,
696            index::{IndexId, IndexKey, IndexKeyKind},
697            key_taxonomy::{PrimaryKeyComponent, PrimaryKeyValue},
698        },
699        testing::test_memory,
700        traits::Storable,
701        types::EntityTag,
702    };
703    use std::{borrow::Cow, convert::Infallible};
704
705    fn raw_key(value: u8) -> RawIndexStoreKey {
706        <RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(vec![value]))
707    }
708
709    fn indexed_raw_key(
710        index_id: &IndexId,
711        components: Vec<Vec<u8>>,
712        primary_key: u64,
713    ) -> RawIndexStoreKey {
714        indexed_raw_key_with_kind(index_id, IndexKeyKind::User, components, primary_key)
715    }
716
717    fn indexed_raw_key_with_kind(
718        index_id: &IndexId,
719        key_kind: IndexKeyKind,
720        components: Vec<Vec<u8>>,
721        primary_key: u64,
722    ) -> RawIndexStoreKey {
723        IndexKey::new_from_components_with_primary_key_value(
724            index_id,
725            key_kind,
726            components.as_slice(),
727            &PrimaryKeyValue::from(PrimaryKeyComponent::Nat64(primary_key)),
728        )
729        .to_raw()
730    }
731
732    fn malformed_index_entry_value() -> IndexEntryValue {
733        <IndexEntryValue as Storable>::from_bytes(Cow::Owned(vec![0xFF]))
734    }
735
736    fn missing_index_entry_value() -> IndexEntryValue {
737        <IndexEntryValue as Storable>::from_bytes(Cow::Owned(vec![1]))
738    }
739
740    #[test]
741    fn index_prefix_cardinality_requires_explicit_data_generation_sync() {
742        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
743        let collection = b"collection-a".to_vec();
744        let draft = b"Draft".to_vec();
745        let review = b"Review".to_vec();
746        let mut store = IndexStore::init_heap();
747
748        store.insert(
749            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
750            IndexEntryValue::presence(),
751        );
752        store.insert(
753            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2),
754            IndexEntryValue::presence(),
755        );
756        store.insert(
757            indexed_raw_key(&index_id, vec![collection.clone(), review.clone()], 3),
758            IndexEntryValue::presence(),
759        );
760
761        assert_eq!(
762            store.exact_prefix_cardinality(
763                0,
764                IndexKeyKind::User,
765                index_id,
766                std::slice::from_ref(&collection),
767            ),
768            None,
769            "raw index mutations must not be trusted until row generation sync is stamped",
770        );
771
772        store.mark_prefix_cardinality_data_generation(7);
773
774        assert_eq!(
775            store.exact_prefix_cardinality(
776                7,
777                IndexKeyKind::User,
778                index_id,
779                std::slice::from_ref(&collection),
780            ),
781            Some(3),
782        );
783        assert_eq!(
784            store.exact_prefix_cardinality(
785                7,
786                IndexKeyKind::User,
787                index_id,
788                &[collection.clone(), draft],
789            ),
790            Some(2),
791        );
792        assert_eq!(
793            store.exact_prefix_cardinality(8, IndexKeyKind::User, index_id, &[collection, review],),
794            None,
795            "row generation drift should force the caller to use the existing-row fallback",
796        );
797    }
798
799    #[test]
800    fn index_prefix_cardinality_enumerates_bounded_child_prefixes() {
801        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
802        let collection = b"collection-a".to_vec();
803        let other_collection = b"collection-b".to_vec();
804        let draft = b"Draft".to_vec();
805        let review = b"Review".to_vec();
806        let published = b"Published".to_vec();
807        let mut store = IndexStore::init_heap();
808
809        store.insert(
810            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
811            IndexEntryValue::presence(),
812        );
813        store.insert(
814            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2),
815            IndexEntryValue::presence(),
816        );
817        store.insert(
818            indexed_raw_key(&index_id, vec![collection.clone(), review.clone()], 3),
819            IndexEntryValue::presence(),
820        );
821        store.insert(
822            indexed_raw_key(
823                &index_id,
824                vec![other_collection.clone(), published.clone()],
825                4,
826            ),
827            IndexEntryValue::presence(),
828        );
829        store.mark_prefix_cardinality_data_generation(7);
830
831        assert_eq!(
832            store.exact_child_prefixes(
833                7,
834                IndexKeyKind::User,
835                index_id,
836                std::slice::from_ref(&collection),
837                4,
838            ),
839            Some(vec![
840                vec![collection.clone(), draft],
841                vec![collection.clone(), review],
842            ]),
843            "child-prefix enumeration should return deterministic unique children under the requested parent",
844        );
845        assert_eq!(
846            store.exact_child_prefixes(
847                7,
848                IndexKeyKind::User,
849                index_id,
850                std::slice::from_ref(&other_collection),
851                4,
852            ),
853            Some(vec![vec![other_collection, published]]),
854            "child-prefix enumeration must stay scoped to the requested parent prefix",
855        );
856        assert_eq!(
857            store.exact_child_prefixes(
858                8,
859                IndexKeyKind::User,
860                index_id,
861                std::slice::from_ref(&collection),
862                4,
863            ),
864            None,
865            "row generation drift should keep child-prefix expansion fail-closed",
866        );
867        assert_eq!(
868            store.exact_child_prefixes(
869                7,
870                IndexKeyKind::User,
871                index_id,
872                std::slice::from_ref(&collection),
873                1,
874            ),
875            None,
876            "over-cap child-prefix expansion should fall back to the existing route",
877        );
878    }
879
880    #[test]
881    fn index_prefix_cardinality_ignores_system_index_mutations() {
882        let user_index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
883        let system_index_id = IndexId::new(EntityTag::new(0xCA7D), 2);
884        let collection = b"collection-a".to_vec();
885        let draft = b"Draft".to_vec();
886        let system_component = b"reverse-edge".to_vec();
887        let mut store = IndexStore::init_heap();
888
889        store.insert(
890            indexed_raw_key(&user_index_id, vec![collection.clone(), draft.clone()], 1),
891            IndexEntryValue::presence(),
892        );
893        store.mark_prefix_cardinality_data_generation(7);
894
895        assert_eq!(
896            store.exact_prefix_cardinality(
897                7,
898                IndexKeyKind::User,
899                user_index_id,
900                &[collection.clone(), draft.clone()],
901            ),
902            Some(1),
903        );
904
905        let system_key = indexed_raw_key_with_kind(
906            &system_index_id,
907            IndexKeyKind::System,
908            vec![system_component],
909            1,
910        );
911        store.insert(system_key.clone(), IndexEntryValue::presence());
912        assert_eq!(
913            store.exact_prefix_cardinality(
914                7,
915                IndexKeyKind::User,
916                user_index_id,
917                &[collection.clone(), draft.clone()],
918            ),
919            Some(1),
920            "system index writes must not invalidate synchronized user-prefix cardinality",
921        );
922
923        store.remove(&system_key);
924        assert_eq!(
925            store.exact_prefix_cardinality(
926                7,
927                IndexKeyKind::User,
928                user_index_id,
929                &[collection.clone(), draft.clone()],
930            ),
931            Some(1),
932            "system index removals must not invalidate synchronized user-prefix cardinality",
933        );
934
935        let malformed_system_key = indexed_raw_key_with_kind(
936            &system_index_id,
937            IndexKeyKind::System,
938            vec![b"malformed-reverse-edge".to_vec()],
939            2,
940        );
941        store.insert(malformed_system_key.clone(), malformed_index_entry_value());
942        assert_eq!(
943            store.exact_prefix_cardinality(
944                7,
945                IndexKeyKind::User,
946                user_index_id,
947                &[collection.clone(), draft.clone()],
948            ),
949            Some(1),
950            "malformed system index payloads must not invalidate user-prefix cardinality",
951        );
952
953        store.remove(&malformed_system_key);
954        assert_eq!(
955            store.exact_prefix_cardinality(
956                7,
957                IndexKeyKind::User,
958                user_index_id,
959                &[collection.clone(), draft],
960            ),
961            Some(1),
962            "malformed system index removals must not invalidate user-prefix cardinality",
963        );
964
965        let review = b"Review".to_vec();
966        store.insert(
967            indexed_raw_key(&user_index_id, vec![collection.clone(), review.clone()], 2),
968            IndexEntryValue::presence(),
969        );
970        assert_eq!(
971            store.exact_prefix_cardinality(
972                7,
973                IndexKeyKind::User,
974                user_index_id,
975                &[collection, review]
976            ),
977            None,
978            "user-prefix count changes must still require a fresh row-generation stamp",
979        );
980    }
981
982    #[test]
983    fn index_prefix_cardinality_ignores_missing_user_index_mutations() {
984        let index_id = IndexId::new(EntityTag::new(0xCA7D), 1);
985        let collection = b"collection-a".to_vec();
986        let draft = b"Draft".to_vec();
987        let mut store = IndexStore::init_heap();
988
989        store.insert(
990            indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 1),
991            IndexEntryValue::presence(),
992        );
993        store.mark_prefix_cardinality_data_generation(7);
994
995        let stale_key = indexed_raw_key(&index_id, vec![collection.clone(), draft.clone()], 2);
996        store.insert(stale_key.clone(), missing_index_entry_value());
997        assert_eq!(
998            store.exact_prefix_cardinality(
999                7,
1000                IndexKeyKind::User,
1001                index_id,
1002                &[collection.clone(), draft.clone()],
1003            ),
1004            Some(1),
1005            "missing user index entries must not affect synchronized prefix cardinality",
1006        );
1007
1008        store.remove(&stale_key);
1009        assert_eq!(
1010            store.exact_prefix_cardinality(7, IndexKeyKind::User, index_id, &[collection, draft],),
1011            Some(1),
1012            "missing user index removals must not affect synchronized prefix cardinality",
1013        );
1014    }
1015
1016    #[cfg(feature = "diagnostics")]
1017    #[test]
1018    fn index_store_diagnostic_counters_record_gets_range_scans_and_entry_reads() {
1019        let mut store = IndexStore::init_heap();
1020        store.insert(raw_key(7), IndexEntryValue::presence());
1021        store.insert(raw_key(9), IndexEntryValue::presence());
1022
1023        let gets_before = IndexStore::current_get_call_count();
1024        assert_eq!(store.get(&raw_key(7)), Some(IndexEntryValue::presence()));
1025        assert_eq!(store.get(&raw_key(8)), None);
1026
1027        assert_eq!(
1028            IndexStore::current_get_call_count().saturating_sub(gets_before),
1029            2,
1030            "diagnostic index-store get counter should count both hit and miss reads",
1031        );
1032
1033        let range_scans_before = IndexStore::current_range_scan_call_count();
1034        let lower = Bound::Included(raw_key(7));
1035        let upper = Bound::Included(raw_key(9));
1036        store
1037            .visit_raw_entries_in_range((&lower, &upper), Direction::Asc, |_key, _entry| Ok(false))
1038            .expect("raw index range visit should succeed");
1039
1040        assert_eq!(
1041            IndexStore::current_range_scan_call_count().saturating_sub(range_scans_before),
1042            1,
1043            "diagnostic index-store range-scan counter should count one range traversal probe",
1044        );
1045
1046        let entries_before = IndexStore::current_entry_read_count();
1047        store
1048            .visit_entries(|_key, _entry| Ok::<_, Infallible>(IndexStoreVisit::Continue))
1049            .expect("index entry visit should succeed");
1050
1051        assert_eq!(
1052            IndexStore::current_entry_read_count().saturating_sub(entries_before),
1053            2,
1054            "diagnostic index-store entry counter should count yielded traversal entries",
1055        );
1056    }
1057
1058    #[test]
1059    fn journaled_mixed_index_range_traversal_streams_without_snapshot() {
1060        let mut store = IndexStore::init_journaled(test_memory(93));
1061        for value in [1_u8, 3, 5] {
1062            store.insert(raw_key(value), IndexEntryValue::presence());
1063        }
1064        store
1065            .fold_journaled_materialized_view()
1066            .expect("canonical index seed should fold");
1067
1068        store.insert(raw_key(0), IndexEntryValue::presence());
1069        store.insert(raw_key(4), IndexEntryValue::presence());
1070        store.insert(raw_key(5), IndexEntryValue::presence());
1071        store.remove(&raw_key(1));
1072
1073        let lower = Bound::Included(raw_key(0));
1074        let upper = Bound::Included(raw_key(5));
1075
1076        reset_journaled_snapshot_call_count_for_tests();
1077        let mut asc = Vec::new();
1078        store
1079            .visit_journaled_entries_in_range((&lower, &upper), Direction::Asc, |key, _value| {
1080                asc.push(key.as_bytes()[0]);
1081                Ok::<_, Infallible>(asc.len() == 2)
1082            })
1083            .expect("asc journaled index range traversal should succeed");
1084        assert_eq!(asc, vec![0, 3]);
1085        assert_eq!(
1086            journaled_snapshot_call_count_for_tests(),
1087            0,
1088            "mixed journaled index range traversal should preserve early stop without materializing a snapshot",
1089        );
1090
1091        reset_journaled_snapshot_call_count_for_tests();
1092        let mut desc = Vec::new();
1093        store
1094            .visit_journaled_entries_in_range((&lower, &upper), Direction::Desc, |key, _value| {
1095                desc.push(key.as_bytes()[0]);
1096                Ok::<_, Infallible>(desc.len() == 2)
1097            })
1098            .expect("desc journaled index range traversal should succeed");
1099        assert_eq!(desc, vec![5, 4]);
1100        assert_eq!(
1101            journaled_snapshot_call_count_for_tests(),
1102            0,
1103            "mixed reverse journaled index range traversal should preserve early stop without materializing a snapshot",
1104        );
1105    }
1106}
icydb_core/db/index/store.rs

icydb_core/db/index/
store.rs
