Skip to main content

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