#![allow(clippy::unwrap_used)]
#![allow(clippy::panic)]
use std::collections::BTreeMap;
use proptest::prelude::*;
use super::{HandoffMergeError, merge_committed_union};
use crate::store::MemoryStore;
use crate::sync::ballot::{Ballot, Stamp};
use crate::sync::topology::SyncNodeId;
use crate::tree::{Cursor, Hash, LeafNode, Node, TreePolicy, batch_mutate};
use crate::ttl::entry::{StampedEntry, encode_stamped, encode_stamped_tombstone};
#[derive(Clone, Debug)]
struct Write {
key: Vec<u8>,
stamp: Stamp,
value: Option<Vec<u8>>,
}
fn empty_root(store: &mut MemoryStore) -> Hash {
let leaf = LeafNode::new(Vec::new()).unwrap();
store.put(&Node::Leaf(leaf))
}
fn stored_bytes(write: &Write) -> Vec<u8> {
write.value.as_ref().map_or_else(
|| encode_stamped_tombstone(write.stamp.clone()),
|value| encode_stamped(value.clone(), write.stamp.clone(), None),
)
}
fn build_tree(store: &mut MemoryStore, writes: &[Write]) -> Hash {
let mut by_key: BTreeMap<Vec<u8>, Vec<u8>> = BTreeMap::new();
for write in writes {
by_key.insert(write.key.clone(), stored_bytes(write));
}
let root = empty_root(store);
let mutations: Vec<(Vec<u8>, Option<Vec<u8>>)> =
by_key.into_iter().map(|(k, v)| (k, Some(v))).collect();
batch_mutate(store, root, mutations.as_slice(), TreePolicy::V1_DEFAULT).unwrap()
}
fn logical_get(store: &MemoryStore, root: Hash, key: &[u8]) -> Option<Vec<u8>> {
let cursor = Cursor::new(store, root);
let raw = cursor.get(key).unwrap()?;
StampedEntry::decode(&raw).unwrap().unwrap().into_value()
}
fn reference_winners(all: &[Write]) -> BTreeMap<Vec<u8>, Option<Vec<u8>>> {
let mut winners: BTreeMap<Vec<u8>, (Stamp, Option<Vec<u8>>)> = BTreeMap::new();
for write in all {
match winners.get(&write.key) {
Some((existing_stamp, _)) if *existing_stamp >= write.stamp => {}
_ => {
winners.insert(
write.key.clone(),
(write.stamp.clone(), write.value.clone()),
);
}
}
}
winners
.into_iter()
.map(|(key, (_stamp, value))| (key, value))
.collect()
}
fn permutations(items: &[Hash]) -> Vec<Vec<Hash>> {
if items.len() <= 1 {
return vec![items.to_vec()];
}
let mut out = Vec::new();
for index in 0..items.len() {
let mut rest = items.to_vec();
let head = rest.remove(index);
for mut perm in permutations(&rest) {
perm.insert(0, head);
out.push(perm);
}
}
out
}
fn fold_merge(store: &mut MemoryStore, roots: &[Hash]) -> Result<Option<Hash>, HandoffMergeError> {
let mut acc: Option<Hash> = None;
for &root in roots {
acc = merge_committed_union(acc, Some(root), store, TreePolicy::V1_DEFAULT)?;
}
Ok(acc)
}
fn unique_stamp(index: u64) -> Stamp {
let epoch = Ballot::new(index / 4 + 1, SyncNodeId::new("p"));
Stamp::new(epoch, index % 4)
}
prop_compose! {
fn arb_write(index: u64)(
key_index in 0_u8..6,
is_tombstone in any::<bool>(),
value in proptest::collection::vec(any::<u8>(), 0..8),
) -> Write {
Write {
key: vec![b'k', key_index],
stamp: unique_stamp(index),
value: if is_tombstone { None } else { Some(value) },
}
}
}
fn arb_writes() -> impl Strategy<Value = Vec<Write>> {
(1_usize..=12).prop_flat_map(|count| {
let per_write: Vec<_> = (0..count).map(|i| arb_write(i as u64)).collect();
per_write
})
}
fn split(writes: &[Write], parts: usize) -> Vec<Vec<Write>> {
let mut buckets = vec![Vec::new(); parts];
for (i, write) in writes.iter().enumerate() {
buckets[i % parts].push(write.clone());
}
buckets
}
fn assert_converges(writes: &[Write], parts: usize) -> Result<(), TestCaseError> {
let mut store = MemoryStore::new();
let buckets = split(writes, parts);
let roots: Vec<Hash> = buckets.iter().map(|b| build_tree(&mut store, b)).collect();
let baseline = fold_merge(&mut store, &roots)
.map_err(|e| TestCaseError::fail(format!("baseline merge failed: {e}")))?
.ok_or_else(|| TestCaseError::fail("baseline merge produced an empty root"))?;
for perm in permutations(&roots) {
let root = fold_merge(&mut store, &perm)
.map_err(|e| TestCaseError::fail(format!("merge failed: {e}")))?
.ok_or_else(|| TestCaseError::fail("merge produced an empty root"))?;
prop_assert_eq!(
root,
baseline,
"every merge order must yield the identical root"
);
}
let reference = reference_winners(writes);
for (key, expected) in &reference {
prop_assert_eq!(
logical_get(&store, baseline, key),
expected.clone(),
"merged read must equal the per-key max-stamp chain tip"
);
}
Ok(())
}
proptest! {
#[test]
fn two_way_merge_is_order_independent(writes in arb_writes()) {
assert_converges(&writes, 2)?;
}
#[test]
fn three_way_merge_is_order_independent(writes in arb_writes()) {
assert_converges(&writes, 3)?;
}
}