icydb_core/db/index/
store.rs

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

icydb_core/db/index/
store.rs
