use crate::{
db::{
direction::Direction,
index::{
IndexEntryValue, IndexId, IndexKey, IndexKeyKind, IndexStore, IndexStoreVisit,
RawIndexStoreKey, key::EncodedValue, resume_bounds_for_continuation,
},
journal::FoldWatermark,
key_taxonomy::{PrimaryKeyComponent, PrimaryKeyValue},
},
testing::test_memory,
types::{EntityTag, IntBig},
value::Value,
};
use ic_memory::ic_stable_structures::Storable;
use std::{borrow::Cow, cell::Cell, ops::Bound};
fn raw_key(value: u8) -> RawIndexStoreKey {
<RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(vec![value]))
}
fn seed_entry(store: &mut IndexStore, key: RawIndexStoreKey, canonical: bool) {
if canonical {
store
.apply_canonical_entry(key, Some(IndexEntryValue::presence()))
.unwrap();
} else {
store.insert(key, IndexEntryValue::presence());
}
}
#[test]
fn empty_envelopes_skip_populated_store_traversal() {
for (mut store, journaled) in [
(IndexStore::init_heap(), false),
(IndexStore::init_journaled(test_memory(97)), true),
] {
for value in 1..=3 {
seed_entry(&mut store, bigint_scan_key(value, 1), journaled);
}
store.insert(bigint_scan_key(4, 1), IndexEntryValue::presence());
let low = bigint_scan_key(1, 1);
let high = bigint_scan_key(3, 1);
let lower = Bound::Included(low.clone());
let upper = Bound::Included(high.clone());
let asc_empty =
resume_bounds_for_continuation(Direction::Asc, Some(&high), &lower, &upper).unwrap();
let desc_empty =
resume_bounds_for_continuation(Direction::Desc, Some(&low), &lower, &upper).unwrap();
let cases = [
("ASC endpoint resume", asc_empty, 0, 0),
("DESC endpoint resume", desc_empty, 0, 0),
("inverted", (upper.clone(), lower.clone()), 0, 0),
(
"equal lower excluded",
(Bound::Excluded(low.clone()), lower.clone()),
0,
0,
),
(
"equal upper excluded",
(lower.clone(), Bound::Excluded(low.clone())),
0,
0,
),
(
"equal both excluded",
(Bound::Excluded(low.clone()), Bound::Excluded(low)),
0,
0,
),
("singleton", (lower.clone(), lower), 1, 1),
(
"unbounded control",
(Bound::Unbounded, Bound::Unbounded),
4,
if journaled { 3 } else { 4 },
),
];
for (case, (lower, upper), expected_live, expected_canonical) in cases {
for direction in [Direction::Asc, Direction::Desc] {
let mut visited = 0;
store
.visit_raw_entries_in_range((&lower, &upper), direction, |_, _| {
visited += 1;
Ok(false)
})
.unwrap();
assert_eq!(
visited, expected_live,
"{case}: {journaled:?}, {direction:?}"
);
}
let mut visited = 0;
store
.visit_canonical_raw_entries_in_range((&lower, &upper), |_, _| {
visited += 1;
Ok(false)
})
.unwrap();
assert_eq!(visited, expected_canonical, "{case}: {journaled:?}");
}
}
}
#[test]
fn other_key_probe_preserves_replay_exclusion_and_effective_overlay() {
let key = |value, row_id| {
IndexKey::new_from_components_with_primary_key_value(
&IndexId::new(EntityTag::new(17), 0),
IndexKeyKind::User,
&[EncodedValue::try_new(&Value::Nat64(value))
.unwrap()
.into_bytes()],
&PrimaryKeyValue::from(PrimaryKeyComponent::Nat64(row_id)),
)
.unwrap()
.to_raw()
.unwrap()
};
let candidate = key(7, 1);
let conflicting = key(7, 2);
let outside = key(8, 2);
let (lower, upper) = IndexKey::try_from_raw(&candidate)
.unwrap()
.raw_bounds_for_all_components()
.unwrap();
let bounds = (Bound::Included(lower), Bound::Included(upper));
for (mut store, journaled) in [
(IndexStore::init_heap(), false),
(IndexStore::init_journaled(test_memory(98)), true),
] {
let probe = |store: &IndexStore| {
store
.contains_other_raw_key_in_range((&bounds.0, &bounds.1), &candidate)
.unwrap()
};
assert!(!probe(&store));
seed_entry(&mut store, outside.clone(), journaled);
assert!(!probe(&store));
seed_entry(&mut store, candidate.clone(), journaled);
assert!(!probe(&store));
store.insert(conflicting.clone(), IndexEntryValue::presence());
assert!(probe(&store));
if journaled {
store
.apply_canonical_entry(conflicting.clone(), Some(IndexEntryValue::presence()))
.unwrap();
store
.reset_journaled_live_projection(0, FoldWatermark::initial())
.unwrap();
}
assert!(probe(&store));
store.remove(&conflicting);
assert!(!probe(&store));
store.insert(conflicting.clone(), IndexEntryValue::presence());
store.remove(&candidate);
assert!(probe(&store));
}
}
#[test]
fn visit_raw_entries_in_range_preserves_directional_store_order() {
let mut index_store = IndexStore::init_journaled(test_memory(91));
for value in [1_u8, 2, 3] {
let raw_key = <RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(vec![value]));
let raw_entry = IndexEntryValue::presence();
index_store.insert(raw_key, raw_entry);
}
let lower = Bound::Included(<RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(
vec![1],
)));
let upper = Bound::Included(<RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(
vec![3],
)));
let mut asc = Vec::new();
index_store
.visit_raw_entries_in_range((&lower, &upper), Direction::Asc, |raw_key, _| {
asc.push(raw_key.as_bytes()[0]);
Ok(false)
})
.expect("asc scan should succeed");
assert_eq!(asc, vec![1, 2, 3], "asc scan should follow raw key order");
let mut desc = Vec::new();
index_store
.visit_raw_entries_in_range((&lower, &upper), Direction::Desc, |raw_key, _| {
desc.push(raw_key.as_bytes()[0]);
Ok(false)
})
.expect("desc scan should succeed");
assert_eq!(
desc,
vec![3, 2, 1],
"desc scan should reverse raw key order"
);
}
#[test]
fn visit_entries_preserves_store_order_and_supports_early_stop() {
let mut index_store = IndexStore::init_journaled(test_memory(92));
for value in [3_u8, 1, 2] {
let raw_key = <RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(vec![value]));
let raw_entry = IndexEntryValue::presence();
index_store.insert(raw_key, raw_entry);
}
let mut visited = Vec::new();
let _: Result<(), std::convert::Infallible> = index_store.visit_entries(|raw_key, _| {
visited.push(raw_key.as_bytes()[0]);
Ok(if visited.len() == 2 {
IndexStoreVisit::Stop
} else {
IndexStoreVisit::Continue
})
});
assert_eq!(
visited,
vec![1, 2],
"index entry traversal should preserve raw store order and stop without allocation"
);
}
#[test]
fn heap_index_store_preserves_range_order_and_early_stop() {
let mut index_store = IndexStore::init_heap();
for value in [3_u8, 1, 2] {
let raw_key = <RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(vec![value]));
let raw_entry = IndexEntryValue::presence();
index_store.insert(raw_key, raw_entry);
}
let lower = Bound::Included(<RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(
vec![1],
)));
let upper = Bound::Included(<RawIndexStoreKey as Storable>::from_bytes(Cow::Owned(
vec![3],
)));
let mut asc = Vec::new();
index_store
.visit_raw_entries_in_range((&lower, &upper), Direction::Asc, |raw_key, _| {
asc.push(raw_key.as_bytes()[0]);
Ok(false)
})
.expect("heap asc scan should succeed");
assert_eq!(asc, vec![1, 2, 3]);
let mut desc = Vec::new();
index_store
.visit_raw_entries_in_range((&lower, &upper), Direction::Desc, |raw_key, _| {
desc.push(raw_key.as_bytes()[0]);
Ok(false)
})
.expect("heap desc scan should succeed");
assert_eq!(desc, vec![3, 2, 1]);
let mut stopped = Vec::new();
let _: Result<(), std::convert::Infallible> = index_store.visit_entries(|raw_key, _| {
stopped.push(raw_key.as_bytes()[0]);
Ok(if stopped.len() == 2 {
IndexStoreVisit::Stop
} else {
IndexStoreVisit::Continue
})
});
assert_eq!(stopped, vec![1, 2]);
}
#[test]
fn merged_ranges_preserve_logical_order_direction_and_early_stop() {
for (mut index_store, fold) in [
(IndexStore::init_heap(), false),
(IndexStore::init_journaled(test_memory(94)), false),
(IndexStore::init_journaled(test_memory(96)), true),
] {
for value in [10_u8, 11, 12, 20, 21, 22] {
seed_entry(&mut index_store, raw_key(value), fold);
}
let bounds = [
(Bound::Excluded(raw_key(9)), Bound::Included(raw_key(12))),
(Bound::Included(raw_key(20)), Bound::Unbounded),
];
let decode_order = |key: &RawIndexStoreKey| {
let raw = key.as_bytes()[0];
Ok::<u8, crate::error::InternalError>(match raw {
10..=12 => raw.saturating_sub(10).saturating_mul(2).saturating_add(1),
20..=22 => raw.saturating_sub(20).saturating_mul(2).saturating_add(2),
_ => return Err(crate::error::InternalError::executor_invariant()),
})
};
let mut asc = Vec::new();
assert!(
index_store
.visit_raw_entries_in_merged_ranges(
bounds.as_slice(),
Direction::Asc,
|_bytes| Ok(()),
decode_order,
|order, _key, _value| {
asc.push(order);
Ok(asc.len() == 4)
},
)
.expect("merged ASC ranges should execute"),
);
assert_eq!(asc, vec![1, 2, 3, 4]);
let mut desc = Vec::new();
assert!(
index_store
.visit_raw_entries_in_merged_ranges(
bounds.as_slice(),
Direction::Desc,
|_bytes| Ok(()),
decode_order,
|order, _key, _value| {
desc.push(order);
Ok(false)
},
)
.expect("merged DESC ranges should execute"),
);
assert_eq!(desc, vec![6, 5, 4, 3, 2, 1]);
}
}
#[test]
fn merged_ranges_admit_complete_structural_state_before_reading() {
let mut index_store = IndexStore::init_journaled(test_memory(95));
for value in [10_u8, 11, 20, 21] {
seed_entry(&mut index_store, raw_key(value), true);
}
let bounds = [
(Bound::Included(raw_key(10)), Bound::Included(raw_key(11))),
(Bound::Included(raw_key(20)), Bound::Included(raw_key(21))),
];
let retained_bound_bytes = bounds
.iter()
.flat_map(|(lower, upper)| [lower, upper])
.map(RawIndexStoreKey::bound_backing_bytes)
.sum::<usize>();
let admitted = Cell::new(0usize);
let decode_calls = Cell::new(0usize);
let visit_calls = Cell::new(0usize);
let _error = index_store
.visit_raw_entries_in_merged_ranges(
bounds.as_slice(),
Direction::Asc,
|bytes| {
admitted.set(bytes);
Err(crate::error::InternalError::executor_internal())
},
|_key| {
decode_calls.set(decode_calls.get().saturating_add(1));
Ok::<u8, crate::error::InternalError>(0)
},
|_order, _key, _value| {
visit_calls.set(visit_calls.get().saturating_add(1));
Ok(false)
},
)
.expect_err("structural admission rejection should stop the merged scan");
assert!(admitted.get() > retained_bound_bytes);
assert_eq!(decode_calls.get(), 0);
assert_eq!(visit_calls.get(), 0);
}
fn bigint_scan_key(value: i64, suffix: u64) -> RawIndexStoreKey {
let components = [
Value::Text("tenant".into()),
Value::IntBig(IntBig::from(value)),
Value::Nat64(suffix),
]
.map(|value| EncodedValue::try_new(&value).unwrap().into_bytes());
IndexKey::new_from_components_with_primary_key_value(
&IndexId::new(EntityTag::new(17), 0),
IndexKeyKind::User,
&components,
&PrimaryKeyValue::from(PrimaryKeyComponent::Nat64(suffix)),
)
.unwrap()
.to_raw()
.unwrap()
}
fn bigint_scan_page(
store: &IndexStore,
lower: &Bound<RawIndexStoreKey>,
upper: &Bound<RawIndexStoreKey>,
direction: Direction,
limit: usize,
) -> Vec<RawIndexStoreKey> {
let mut keys = Vec::new();
store
.visit_raw_entries_in_range((lower, upper), direction, |key, _| {
keys.push(key.clone());
Ok(keys.len() == limit)
})
.unwrap();
keys
}
#[test]
fn binary_bigint_index_ranges_resume_from_canonical_after_reopen() {
let memory = test_memory(96);
let mut store = IndexStore::init_journaled(memory.clone());
let values = [-65536, -257, -256, -255, -1, 0, 1, 255, 256, 257, 65536];
for value in values.into_iter().rev() {
for suffix in [2, 1] {
seed_entry(&mut store, bigint_scan_key(value, suffix), true);
}
}
drop(store);
let store = IndexStore::init_journaled(memory);
let prefix = [EncodedValue::try_new(&Value::Text("tenant".into()))
.unwrap()
.into_bytes()];
for inclusive in [false, true] {
let low = EncodedValue::try_new(&Value::IntBig(IntBig::from(-256)))
.unwrap()
.into_bytes();
let high = EncodedValue::try_new(&Value::IntBig(IntBig::from(256)))
.unwrap()
.into_bytes();
let low = if inclusive {
Bound::Included(low)
} else {
Bound::Excluded(low)
};
let high = if inclusive {
Bound::Included(high)
} else {
Bound::Excluded(high)
};
let (lower, upper) = IndexKey::raw_bounds_for_prefix_component_range_with_kind(
&IndexId::new(EntityTag::new(17), 0),
IndexKeyKind::User,
3,
&prefix,
&low,
&high,
)
.unwrap();
let expected: Vec<_> = values
.into_iter()
.filter(|&value| {
if inclusive {
(-256..=256).contains(&value)
} else {
(-255..256).contains(&value)
}
})
.flat_map(|value| [bigint_scan_key(value, 1), bigint_scan_key(value, 2)])
.collect();
for direction in [Direction::Asc, Direction::Desc] {
let mut expected = expected.clone();
if direction == Direction::Desc {
expected.reverse();
}
let mut actual = bigint_scan_page(&store, &lower, &upper, direction, 3);
let boundary = Bound::Excluded(actual.last().unwrap().clone());
let rest = if direction == Direction::Asc {
bigint_scan_page(&store, &boundary, &upper, direction, usize::MAX)
} else {
bigint_scan_page(&store, &lower, &boundary, direction, usize::MAX)
};
actual.extend(rest);
assert_eq!(actual, expected);
}
}
}