pmtree/tests/memory_keccak.rs

132 lines
3.5 KiB
Rust
Raw Normal View History

2022-11-01 09:08:27 +00:00
use hex_literal::hex;
use pmtree::*;
use std::collections::HashMap;
use tiny_keccak::{Hasher as _, Keccak};
2022-11-01 09:08:27 +00:00
2022-11-02 10:45:49 +00:00
struct MemoryDB(HashMap<DBKey, Value>);
struct MyKeccak(Keccak);
2022-11-01 09:08:27 +00:00
#[derive(Default)]
struct MemoryDBConfig;
2022-11-01 09:08:27 +00:00
impl Database for MemoryDB {
type Config = MemoryDBConfig;
2023-04-18 08:52:38 +00:00
fn new(_db_config: MemoryDBConfig) -> PmtreeResult<Self> {
2022-11-06 12:25:19 +00:00
Ok(MemoryDB(HashMap::new()))
2022-11-01 09:08:27 +00:00
}
2023-04-18 08:52:38 +00:00
fn load(_db_config: MemoryDBConfig) -> PmtreeResult<Self> {
Err(DatabaseError(DatabaseErrorKind::CannotLoadDatabase))
2022-11-01 09:08:27 +00:00
}
2023-04-18 08:52:38 +00:00
fn get(&self, key: DBKey) -> PmtreeResult<Option<Value>> {
2022-11-06 12:25:19 +00:00
Ok(self.0.get(&key).cloned())
2022-11-01 09:08:27 +00:00
}
2023-04-18 08:52:38 +00:00
fn put(&mut self, key: DBKey, value: Value) -> PmtreeResult<()> {
self.0.insert(key, value);
2022-11-06 12:25:19 +00:00
Ok(())
2022-11-01 09:08:27 +00:00
}
2023-02-02 10:31:18 +00:00
2023-04-18 08:52:38 +00:00
fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> PmtreeResult<()> {
self.0.extend(subtree.into_iter());
2023-02-02 10:31:18 +00:00
Ok(())
}
2022-11-01 09:08:27 +00:00
}
2022-11-30 18:50:10 +00:00
impl Hasher for MyKeccak {
type Fr = [u8; 32];
2022-11-01 09:08:27 +00:00
2022-11-30 18:50:10 +00:00
fn default_leaf() -> Self::Fr {
[0; 32]
2022-11-01 09:08:27 +00:00
}
2022-11-30 18:50:10 +00:00
fn serialize(value: Self::Fr) -> Value {
value.to_vec()
}
2022-11-01 09:08:27 +00:00
2022-11-30 18:50:10 +00:00
fn deserialize(value: Value) -> Self::Fr {
value.try_into().unwrap()
2022-11-01 09:08:27 +00:00
}
fn hash(input: &[Self::Fr]) -> Self::Fr {
let mut output = [0; 32];
let mut hasher = Keccak::v256();
for element in input {
2022-11-30 18:50:10 +00:00
hasher.update(element);
}
hasher.finalize(&mut output);
2022-11-30 18:50:10 +00:00
output
2022-11-01 09:08:27 +00:00
}
}
#[test]
2023-04-18 08:52:38 +00:00
fn insert_delete() -> PmtreeResult<()> {
let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, MemoryDBConfig)?;
assert_eq!(mt.capacity(), 4);
assert_eq!(mt.depth(), 2);
let leaves = [
hex!("0000000000000000000000000000000000000000000000000000000000000001"),
hex!("0000000000000000000000000000000000000000000000000000000000000002"),
hex!("0000000000000000000000000000000000000000000000000000000000000003"),
hex!("0000000000000000000000000000000000000000000000000000000000000004"),
];
let default_tree_root =
hex!("b4c11951957c6f8f642c4af61cd6b24640fec6dc7fc607ee8206a99e92410d30");
2022-11-30 18:50:10 +00:00
assert_eq!(mt.root(), default_tree_root);
let roots = [
hex!("c1ba1812ff680ce84c1d5b4f1087eeb08147a4d510f3496b2849df3a73f5af95"),
hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b"),
hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c"),
hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36"),
];
for i in 0..leaves.len() {
2022-11-30 18:50:10 +00:00
mt.update_next(leaves[i])?;
assert_eq!(mt.root(), roots[i]);
}
2022-11-08 12:36:17 +00:00
for (i, &leaf) in leaves.iter().enumerate() {
2022-11-30 18:50:10 +00:00
assert!(mt.verify(&leaf, &mt.proof(i)?));
2022-11-08 12:36:17 +00:00
}
for i in (0..leaves.len()).rev() {
2022-11-06 12:25:19 +00:00
mt.delete(i)?;
}
2022-11-30 18:50:10 +00:00
assert_eq!(mt.root(), default_tree_root);
2022-11-06 12:25:19 +00:00
2022-11-30 18:50:10 +00:00
assert!(mt.update_next(leaves[0]).is_err());
2022-11-06 12:25:19 +00:00
Ok(())
}
#[test]
2023-04-18 08:52:38 +00:00
fn batch_insertions() -> PmtreeResult<()> {
let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, MemoryDBConfig)?;
let leaves = [
hex!("0000000000000000000000000000000000000000000000000000000000000001"),
hex!("0000000000000000000000000000000000000000000000000000000000000002"),
hex!("0000000000000000000000000000000000000000000000000000000000000003"),
hex!("0000000000000000000000000000000000000000000000000000000000000004"),
];
mt.batch_insert(&leaves)?;
2023-02-02 16:24:58 +00:00
assert_eq!(
mt.root(),
hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36")
);
Ok(())
}