#[cfg(feature = "var-collections")]
use armdb::VarTree;
use armdb::{Config, ConstTree, DbError};
use proptest::prelude::*;
use std::collections::BTreeMap;
use tempfile::TempDir;
fn make_key(group: u8, id: u16) -> [u8; 8] {
let mut k = [0u8; 8];
k[0] = group;
k[6..8].copy_from_slice(&id.to_be_bytes());
k
}
fn make_config(reversed: bool) -> Config {
let mut config = Config::test();
config.shard_count = 2;
config.max_file_size = 4096;
config.compaction_threshold = 0.1;
config.reversed = reversed;
config
}
#[derive(Debug, Clone)]
enum ConstOp {
Put {
group: u8,
id: u16,
value: u64,
},
Get {
group: u8,
id: u16,
},
Delete {
group: u8,
id: u16,
},
Insert {
group: u8,
id: u16,
value: u64,
},
Cas {
group: u8,
id: u16,
expected: u64,
new_value: u64,
},
Compact,
CloseReopen,
IterAll,
PrefixIter {
group: u8,
},
}
fn const_op_strategy() -> impl Strategy<Value = ConstOp> {
prop_oneof![
40 => (0..4u8, 0..16u16, any::<u64>())
.prop_map(|(g, id, v)| ConstOp::Put { group: g, id, value: v }),
20 => (0..4u8, 0..16u16)
.prop_map(|(g, id)| ConstOp::Get { group: g, id }),
15 => (0..4u8, 0..16u16)
.prop_map(|(g, id)| ConstOp::Delete { group: g, id }),
5 => (0..4u8, 0..16u16, any::<u64>())
.prop_map(|(g, id, v)| ConstOp::Insert { group: g, id, value: v }),
5 => (0..4u8, 0..16u16, any::<u64>(), any::<u64>())
.prop_map(|(g, id, e, n)| ConstOp::Cas { group: g, id, expected: e, new_value: n }),
5 => Just(ConstOp::Compact),
3 => Just(ConstOp::CloseReopen),
5 => Just(ConstOp::IterAll),
2 => (0..4u8).prop_map(|g| ConstOp::PrefixIter { group: g }),
]
}
fn run_const_tree_model(ops: Vec<ConstOp>, reversed: bool) {
let dir = TempDir::new().unwrap();
let mut model: BTreeMap<[u8; 8], [u8; 8]> = BTreeMap::new();
let mut tree = ConstTree::<[u8; 8], 8>::open(dir.path(), make_config(reversed)).unwrap();
for (i, op) in ops.iter().enumerate() {
match op {
ConstOp::Put { group, id, value } => {
let k = make_key(*group, *id);
let v = value.to_be_bytes();
let tree_old = tree.put(&k, &v).unwrap();
let model_old = model.insert(k, v);
assert_eq!(tree_old, model_old, "put old mismatch at op {i}");
}
ConstOp::Get { group, id } => {
let k = make_key(*group, *id);
let tree_val = tree.get(&k);
let model_val = model.get(&k).copied();
assert_eq!(tree_val, model_val, "get mismatch at op {i}");
}
ConstOp::Delete { group, id } => {
let k = make_key(*group, *id);
let tree_old = tree.delete(&k).unwrap();
let model_old = model.remove(&k);
assert_eq!(tree_old, model_old, "delete old mismatch at op {i}");
}
ConstOp::Insert { group, id, value } => {
let k = make_key(*group, *id);
let v = value.to_be_bytes();
let result = tree.insert(&k, &v);
#[allow(clippy::map_entry)]
if model.contains_key(&k) {
assert!(
matches!(result, Err(DbError::KeyExists)),
"insert on existing key should be KeyExists at op {i}, got {result:?}"
);
} else {
result.unwrap();
model.insert(k, v);
}
}
ConstOp::Cas {
group,
id,
expected,
new_value,
} => {
let k = make_key(*group, *id);
let e = expected.to_be_bytes();
let n = new_value.to_be_bytes();
let result = tree.cas(&k, &e, &n);
match model.get(&k) {
None => assert!(
matches!(result, Err(DbError::KeyNotFound)),
"cas on missing key should be KeyNotFound at op {i}, got {result:?}"
),
Some(v) if *v == e => {
result.unwrap();
model.insert(k, n);
}
Some(_) => assert!(
matches!(result, Err(DbError::CasMismatch)),
"cas mismatch should be CasMismatch at op {i}, got {result:?}"
),
}
}
ConstOp::Compact => {
tree.compact().unwrap();
assert_eq!(
tree.len(),
model.len(),
"len mismatch after compact at op {i}"
);
}
ConstOp::CloseReopen => {
tree.close().unwrap();
tree = ConstTree::<[u8; 8], 8>::open(dir.path(), make_config(reversed)).unwrap();
assert_eq!(
tree.len(),
model.len(),
"len mismatch after reopen at op {i}"
);
for (k, v) in &model {
assert_eq!(
tree.get(k),
Some(*v),
"value mismatch after reopen at op {i}, key {k:?}"
);
}
}
ConstOp::IterAll => {
let tree_entries: Vec<_> = tree.iter().collect();
let model_entries: Vec<_> = if reversed {
model.iter().rev().map(|(k, v)| (*k, *v)).collect()
} else {
model.iter().map(|(k, v)| (*k, *v)).collect()
};
assert_eq!(tree_entries, model_entries, "iter mismatch at op {i}");
}
ConstOp::PrefixIter { group } => {
let tree_entries: Vec<_> = tree.prefix_iter(&[*group]).collect();
let model_entries: Vec<_> = if reversed {
model
.iter()
.rev()
.filter(|(k, _)| k[0] == *group)
.map(|(k, v)| (*k, *v))
.collect()
} else {
model
.iter()
.filter(|(k, _)| k[0] == *group)
.map(|(k, v)| (*k, *v))
.collect()
};
assert_eq!(
tree_entries, model_entries,
"prefix_iter mismatch at op {i}"
);
}
}
}
assert_eq!(tree.len(), model.len(), "final len mismatch");
for (k, v) in &model {
assert_eq!(tree.get(k).unwrap(), *v, "final value mismatch for {k:?}");
}
}
#[cfg(feature = "var-collections")]
#[derive(Debug, Clone)]
enum VarOp {
Put {
group: u8,
id: u16,
value: Vec<u8>,
},
Get {
group: u8,
id: u16,
},
Delete {
group: u8,
id: u16,
},
Insert {
group: u8,
id: u16,
value: Vec<u8>,
},
Cas {
group: u8,
id: u16,
expected: Vec<u8>,
new_value: Vec<u8>,
},
Compact,
CloseReopen,
IterAll,
PrefixIter {
group: u8,
},
}
#[cfg(feature = "var-collections")]
fn var_op_strategy() -> impl Strategy<Value = VarOp> {
prop_oneof![
40 => (0..4u8, 0..16u16, prop::collection::vec(any::<u8>(), 0..128))
.prop_map(|(g, id, v)| VarOp::Put { group: g, id, value: v }),
20 => (0..4u8, 0..16u16)
.prop_map(|(g, id)| VarOp::Get { group: g, id }),
15 => (0..4u8, 0..16u16)
.prop_map(|(g, id)| VarOp::Delete { group: g, id }),
5 => (0..4u8, 0..16u16, prop::collection::vec(any::<u8>(), 0..128))
.prop_map(|(g, id, v)| VarOp::Insert { group: g, id, value: v }),
5 => (0..4u8, 0..16u16, prop::collection::vec(any::<u8>(), 0..128), prop::collection::vec(any::<u8>(), 0..128))
.prop_map(|(g, id, e, n)| VarOp::Cas { group: g, id, expected: e, new_value: n }),
5 => Just(VarOp::Compact),
3 => Just(VarOp::CloseReopen),
5 => Just(VarOp::IterAll),
2 => (0..4u8).prop_map(|g| VarOp::PrefixIter { group: g }),
]
}
#[cfg(feature = "var-collections")]
fn run_var_tree_model(ops: Vec<VarOp>, reversed: bool) {
let dir = TempDir::new().unwrap();
let mut model: BTreeMap<[u8; 8], Vec<u8>> = BTreeMap::new();
let mut tree = VarTree::<[u8; 8]>::open(dir.path(), make_config(reversed)).unwrap();
for (i, op) in ops.iter().enumerate() {
match op {
VarOp::Put { group, id, value } => {
let k = make_key(*group, *id);
tree.put(&k, value).unwrap();
model.insert(k, value.clone());
}
VarOp::Get { group, id } => {
let k = make_key(*group, *id);
let tree_val = tree.get(&k).map(|bv| bv.as_ref().to_vec());
let model_val = model.get(&k).cloned();
assert_eq!(tree_val, model_val, "get mismatch at op {i}");
}
VarOp::Delete { group, id } => {
let k = make_key(*group, *id);
let tree_existed = tree.delete(&k).unwrap();
let model_existed = model.remove(&k).is_some();
assert_eq!(
tree_existed, model_existed,
"delete existed mismatch at op {i}"
);
}
VarOp::Insert { group, id, value } => {
let k = make_key(*group, *id);
let result = tree.insert(&k, value);
#[allow(clippy::map_entry)]
if model.contains_key(&k) {
assert!(
matches!(result, Err(DbError::KeyExists)),
"insert on existing key should be KeyExists at op {i}, got {result:?}"
);
} else {
result.unwrap();
model.insert(k, value.clone());
}
}
VarOp::Cas {
group,
id,
expected,
new_value,
} => {
let k = make_key(*group, *id);
let result = tree.cas(&k, expected, new_value);
match model.get(&k) {
None => assert!(
matches!(result, Err(DbError::KeyNotFound)),
"cas on missing key should be KeyNotFound at op {i}, got {result:?}"
),
Some(v) if v.as_slice() == expected.as_slice() => {
result.unwrap();
model.insert(k, new_value.clone());
}
Some(_) => assert!(
matches!(result, Err(DbError::CasMismatch)),
"cas mismatch should be CasMismatch at op {i}, got {result:?}"
),
}
}
VarOp::Compact => {
tree.compact().unwrap();
assert_eq!(
tree.len(),
model.len(),
"len mismatch after compact at op {i}"
);
}
VarOp::CloseReopen => {
tree.close().unwrap();
tree = VarTree::open(dir.path(), make_config(reversed)).unwrap();
assert_eq!(
tree.len(),
model.len(),
"len mismatch after reopen at op {i}"
);
for (k, v) in &model {
assert_eq!(
tree.get(k).unwrap().as_ref(),
v.as_slice(),
"value mismatch after reopen at op {i}, key {k:?}"
);
}
}
VarOp::IterAll => {
let tree_entries: Vec<_> = tree
.iter()
.map(|(k, bv)| (k, bv.as_ref().to_vec()))
.collect();
let model_entries: Vec<_> = if reversed {
model.iter().rev().map(|(k, v)| (*k, v.clone())).collect()
} else {
model.iter().map(|(k, v)| (*k, v.clone())).collect()
};
assert_eq!(tree_entries, model_entries, "iter mismatch at op {i}");
}
VarOp::PrefixIter { group } => {
let tree_entries: Vec<_> = tree
.prefix_iter(&[*group])
.map(|(k, bv)| (k, bv.as_ref().to_vec()))
.collect();
let model_entries: Vec<_> = if reversed {
model
.iter()
.rev()
.filter(|(k, _)| k[0] == *group)
.map(|(k, v)| (*k, v.clone()))
.collect()
} else {
model
.iter()
.filter(|(k, _)| k[0] == *group)
.map(|(k, v)| (*k, v.clone()))
.collect()
};
assert_eq!(
tree_entries, model_entries,
"prefix_iter mismatch at op {i}"
);
}
}
}
assert_eq!(tree.len(), model.len(), "final len mismatch");
for (k, v) in &model {
assert_eq!(
tree.get(k).unwrap().as_ref(),
v.as_slice(),
"final value mismatch for {k:?}"
);
}
}
proptest! {
#![proptest_config(ProptestConfig::with_cases(256))]
#[test]
fn const_tree_model_reversed(ops in prop::collection::vec(const_op_strategy(), 1..200)) {
run_const_tree_model(ops, true);
}
#[test]
fn const_tree_model_ascending(ops in prop::collection::vec(const_op_strategy(), 1..200)) {
run_const_tree_model(ops, false);
}
}
#[cfg(feature = "var-collections")]
proptest! {
#![proptest_config(ProptestConfig::with_cases(256))]
#[test]
fn var_tree_model_reversed(ops in prop::collection::vec(var_op_strategy(), 1..200)) {
run_var_tree_model(ops, true);
}
#[test]
fn var_tree_model_ascending(ops in prop::collection::vec(var_op_strategy(), 1..200)) {
run_var_tree_model(ops, false);
}
}