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