#![cfg(feature = "typed-tree")]
use armdb::codec::Codec;
use armdb::{Config, DbError, MigrateAction, TypedTree};
use tempfile::tempdir;
struct U64Codec;
impl Codec<u64> for U64Codec {
fn encode_to(&self, value: &u64, buf: &mut Vec<u8>) {
buf.clear();
buf.extend_from_slice(&value.to_be_bytes());
}
fn decode_from(&self, bytes: &[u8]) -> armdb::DbResult<u64> {
let arr: [u8; 8] = bytes
.try_into()
.map_err(|_| DbError::CorruptedEntry { offset: 0 })?;
Ok(u64::from_be_bytes(arr))
}
}
type Tree = TypedTree<[u8; 8], u64, U64Codec>;
fn open(dir: &std::path::Path) -> Tree {
TypedTree::<[u8; 8], u64, U64Codec>::open(dir, Config::test(), U64Codec).unwrap()
}
fn populate(tree: &Tree, n: u64) {
for i in 0..n {
tree.put(&i.to_be_bytes(), i * 10).unwrap();
}
}
#[test]
fn test_put_get() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
assert!(tree.get(&1u64.to_be_bytes()).is_none());
tree.put(&1u64.to_be_bytes(), 100).unwrap();
let r = tree.get(&1u64.to_be_bytes()).unwrap();
assert_eq!(*r, 100);
}
#[test]
fn test_put_returns_old() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
let old = tree.put(&1u64.to_be_bytes(), 100).unwrap();
assert!(old.is_none());
let old = tree.put(&1u64.to_be_bytes(), 200).unwrap();
assert_eq!(*old.unwrap(), 100);
assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 200);
}
#[test]
fn test_insert() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
tree.insert(&1u64.to_be_bytes(), 100).unwrap();
assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 100);
let err = tree.insert(&1u64.to_be_bytes(), 200).unwrap_err();
assert!(matches!(err, DbError::KeyExists));
assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 100);
}
#[test]
fn test_delete() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
let old = tree.delete(&1u64.to_be_bytes()).unwrap();
assert!(old.is_none());
tree.put(&1u64.to_be_bytes(), 100).unwrap();
let old = tree.delete(&1u64.to_be_bytes()).unwrap();
assert_eq!(*old.unwrap(), 100);
assert!(tree.get(&1u64.to_be_bytes()).is_none());
}
#[test]
fn test_contains() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
assert!(!tree.contains(&1u64.to_be_bytes()));
tree.put(&1u64.to_be_bytes(), 42).unwrap();
assert!(tree.contains(&1u64.to_be_bytes()));
}
#[test]
fn test_len() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
assert_eq!(tree.len(), 0);
assert!(tree.is_empty());
populate(&tree, 10);
assert_eq!(tree.len(), 10);
assert!(!tree.is_empty());
}
#[test]
fn test_cas() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
tree.put(&1u64.to_be_bytes(), 100).unwrap();
let err = tree.cas(&1u64.to_be_bytes(), &999, 200).unwrap_err();
assert!(matches!(err, DbError::CasMismatch));
assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 100);
tree.cas(&1u64.to_be_bytes(), &100, 200).unwrap();
assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 200);
let err = tree.cas(&99u64.to_be_bytes(), &0, 1).unwrap_err();
assert!(matches!(err, DbError::KeyNotFound));
}
#[test]
fn test_update() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
let r = tree.update(&1u64.to_be_bytes(), |v| v + 1).unwrap();
assert!(r.is_none());
tree.put(&1u64.to_be_bytes(), 100).unwrap();
let r = tree.update(&1u64.to_be_bytes(), |v| v * 2).unwrap();
assert_eq!(*r.unwrap(), 200);
assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 200);
}
#[test]
fn test_iter_all() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
populate(&tree, 20);
let entries: Vec<_> = tree.iter().collect();
assert_eq!(entries.len(), 20);
for i in 0..20u64 {
let expected = 19 - i;
assert_eq!(entries[i as usize].0, expected.to_be_bytes());
assert_eq!(*entries[i as usize].1, expected * 10);
}
}
#[test]
fn test_iter_empty() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
assert_eq!(tree.iter().count(), 0);
assert!(tree.iter().next().is_none());
assert!(tree.iter().next_back().is_none());
}
#[test]
fn test_range() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
populate(&tree, 20);
let start = 5u64.to_be_bytes();
let end = 10u64.to_be_bytes();
let entries: Vec<_> = tree.range(&start, &end).collect();
assert_eq!(entries.len(), 5); for (i, (key, val)) in entries.iter().enumerate() {
let expected = 9 - i as u64;
assert_eq!(*key, expected.to_be_bytes());
assert_eq!(**val, expected * 10);
}
}
#[test]
fn test_prefix_iter() {
let dir = tempdir().unwrap();
let tree: TypedTree<[u8; 16], u64, U64Codec> =
TypedTree::open(dir.path(), Config::test(), U64Codec).unwrap();
for i in 0u64..5 {
let mut key = [0u8; 16];
key[3] = 1;
key[8..].copy_from_slice(&i.to_be_bytes());
tree.put(&key, i).unwrap();
}
for i in 0u64..3 {
let mut key = [0u8; 16];
key[3] = 2;
key[8..].copy_from_slice(&i.to_be_bytes());
tree.put(&key, 100 + i).unwrap();
}
let entries: Vec<_> = tree.prefix_iter(&[0, 0, 0, 1]).collect();
assert_eq!(entries.len(), 5);
let entries: Vec<_> = tree.prefix_iter(&[0, 0, 0, 2]).collect();
assert_eq!(entries.len(), 3);
}
#[test]
fn test_iter_rev() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
populate(&tree, 10);
let entries: Vec<_> = tree.iter().rev().collect();
assert_eq!(entries.len(), 10);
for i in 0..10u64 {
assert_eq!(entries[i as usize].0, i.to_be_bytes());
}
}
#[test]
fn test_next_back_mixed() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
populate(&tree, 5);
let mut iter = tree.iter();
let (k, _) = iter.next().unwrap();
assert_eq!(k, 4u64.to_be_bytes());
let (k, _) = iter.next_back().unwrap();
assert_eq!(k, 0u64.to_be_bytes());
let (k, _) = iter.next().unwrap();
assert_eq!(k, 3u64.to_be_bytes());
let (k, _) = iter.next_back().unwrap();
assert_eq!(k, 1u64.to_be_bytes());
let (k, _) = iter.next().unwrap();
assert_eq!(k, 2u64.to_be_bytes());
assert!(iter.next().is_none());
assert!(iter.next_back().is_none());
}
#[test]
fn test_range_rev() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
populate(&tree, 20);
let start = 5u64.to_be_bytes();
let end = 10u64.to_be_bytes();
let entries: Vec<_> = tree.range(&start, &end).rev().collect();
assert_eq!(entries.len(), 5);
for (i, (key, _)) in entries.iter().enumerate() {
let expected = i as u64 + 5;
assert_eq!(*key, expected.to_be_bytes());
}
}
#[test]
fn test_reversed_prefix_iter() {
let dir = tempdir().unwrap();
let mut config = Config::test();
config.reversed = true;
let tree: TypedTree<[u8; 16], u64, U64Codec> =
TypedTree::open(dir.path(), config, U64Codec).unwrap();
for i in 0u64..5 {
let mut key = [0u8; 16];
key[3] = 1;
key[8..].copy_from_slice(&i.to_be_bytes());
tree.put(&key, i).unwrap();
}
let entries: Vec<_> = tree.prefix_iter(&[0, 0, 0, 1]).collect();
assert_eq!(entries.len(), 5);
assert_eq!(*entries[0].1, 4);
assert_eq!(*entries[4].1, 0);
}
#[test]
fn test_reversed_take() {
let dir = tempdir().unwrap();
let mut config = Config::test();
config.reversed = true;
let tree: TypedTree<[u8; 16], u64, U64Codec> =
TypedTree::open(dir.path(), config, U64Codec).unwrap();
for i in 0u64..100 {
let mut key = [0u8; 16];
key[3] = 1;
key[8..].copy_from_slice(&i.to_be_bytes());
tree.put(&key, i).unwrap();
}
let top3: Vec<_> = tree.prefix_iter(&[0, 0, 0, 1]).take(3).collect();
assert_eq!(top3.len(), 3);
assert_eq!(*top3[0].1, 99);
assert_eq!(*top3[1].1, 98);
assert_eq!(*top3[2].1, 97);
}
#[test]
fn test_single_element() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
tree.put(&1u64.to_be_bytes(), 42).unwrap();
let entries: Vec<_> = tree.iter().collect();
assert_eq!(entries.len(), 1);
assert_eq!(*entries[0].1, 42);
let rev: Vec<_> = tree.iter().rev().collect();
assert_eq!(rev.len(), 1);
assert_eq!(*rev[0].1, 42);
}
#[test]
fn test_recovery() {
let dir = tempdir().unwrap();
{
let tree = open(dir.path());
populate(&tree, 50);
tree.close().unwrap();
}
let tree = open(dir.path());
assert_eq!(tree.len(), 50);
for i in 0..50u64 {
let r = tree.get(&i.to_be_bytes()).unwrap();
assert_eq!(*r, i * 10);
}
}
#[test]
fn test_recovery_with_deletes() {
let dir = tempdir().unwrap();
{
let tree = open(dir.path());
populate(&tree, 20);
for i in (0..20u64).step_by(2) {
tree.delete(&i.to_be_bytes()).unwrap();
}
tree.close().unwrap();
}
let tree = open(dir.path());
assert_eq!(tree.len(), 10);
for i in 0..20u64 {
if i % 2 == 1 {
assert!(tree.contains(&i.to_be_bytes()));
} else {
assert!(!tree.contains(&i.to_be_bytes()));
}
}
}
#[test]
fn test_migrate() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
populate(&tree, 10);
let mutations = tree
.migrate(|_key, value| {
if *value < 50 {
MigrateAction::Update(*value + 1000)
} else {
MigrateAction::Keep
}
})
.unwrap();
assert_eq!(mutations, 5);
assert_eq!(*tree.get(&0u64.to_be_bytes()).unwrap(), 1000);
assert_eq!(*tree.get(&3u64.to_be_bytes()).unwrap(), 1030);
assert_eq!(*tree.get(&5u64.to_be_bytes()).unwrap(), 50); }
#[test]
fn test_compact() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
populate(&tree, 100);
for i in 0..100u64 {
tree.put(&i.to_be_bytes(), i * 100).unwrap();
}
tree.flush_buffers().unwrap();
let compacted = tree.compact().unwrap();
let _ = compacted;
for i in 0..100u64 {
assert_eq!(*tree.get(&i.to_be_bytes()).unwrap(), i * 100);
}
}
#[test]
fn test_typed_ref_deref() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
tree.put(&1u64.to_be_bytes(), 42).unwrap();
let r = tree.get(&1u64.to_be_bytes()).unwrap();
let val: &u64 = &r;
assert_eq!(*val, 42);
assert_eq!(*r + 1, 43);
}
#[test]
fn test_put_old_ref_valid() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
tree.put(&1u64.to_be_bytes(), 100).unwrap();
let old = tree.put(&1u64.to_be_bytes(), 200).unwrap().unwrap();
assert_eq!(*old, 100);
assert_eq!(*tree.get(&1u64.to_be_bytes()).unwrap(), 200);
assert_eq!(*old, 100);
}
#[test]
fn test_delete_old_ref_valid() {
let dir = tempdir().unwrap();
let tree = open(dir.path());
tree.put(&1u64.to_be_bytes(), 42).unwrap();
let old = tree.delete(&1u64.to_be_bytes()).unwrap().unwrap();
assert_eq!(*old, 42);
assert!(tree.get(&1u64.to_be_bytes()).is_none());
}