use polydigest::{CmtConfig, EpochTree, Hasher};
pub use polydigest::{CmtError as Error, CmtResult as Result, ProofStep, Sealed};
use sha2::{Digest, Sha256};
const ARITY: u64 = 2;
const ALG: u64 = 0;
#[derive(Debug, Clone, Copy)]
struct Sha256Hasher;
impl Hasher for Sha256Hasher {
fn leaf(&self, data: &[u8]) -> Vec<u8> {
Sha256::digest(data).to_vec()
}
fn node(&self, children: &[&[u8]]) -> Vec<u8> {
let mut h = Sha256::new();
for child in children {
h.update(child);
}
h.finalize().to_vec()
}
fn empty(&self) -> Vec<u8> {
Sha256::digest(b"").to_vec()
}
fn hash(&self, data: &[u8]) -> Vec<u8> {
Sha256::digest(data).to_vec()
}
fn clone_box(&self) -> Box<dyn Hasher> {
Box::new(*self)
}
}
#[derive(Debug)]
pub struct Emt {
inner: EpochTree,
}
impl Emt {
#[must_use]
pub fn new() -> Self {
let mut inner = EpochTree::new(CmtConfig { arity: ARITY }).expect("k=2 is in range");
inner
.register_algorithm(ALG, Box::new(Sha256Hasher))
.expect("a fresh tree has no algorithm registered");
Self { inner }
}
pub fn set(&mut self, index: u64, payload: Vec<u8>) -> Result<()> {
self.inner.set(index, payload, Vec::new())
}
#[must_use]
pub fn get(&self, index: u64) -> Option<&[u8]> {
self.inner.get(index)
}
#[must_use]
pub fn len(&self) -> u64 {
self.inner.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
#[must_use]
pub fn root(&self) -> Option<Vec<u8>> {
self.inner.root(ALG)
}
#[must_use]
pub fn inclusion_proof(&self, index: u64) -> Option<(Vec<u8>, Vec<ProofStep>)> {
self.inner.inclusion_proof(ALG, index)
}
pub fn seal(self) -> Result<Sealed> {
self.inner.seal()
}
}
impl Default for Emt {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use sha2::{Digest, Sha256};
use super::*;
#[test]
fn reads_in_a_few_lines_and_verifies() {
let mut tree = Emt::new();
tree.set(0, b"genesis".to_vec()).unwrap();
tree.set(1, b"second".to_vec()).unwrap();
tree.set(2, b"third".to_vec()).unwrap();
let root = tree.root().unwrap();
let (leaf, path) = tree.inclusion_proof(1).unwrap();
let h = Sha256Hasher;
let sk = polydigest::rebalanced_skeleton(3, ARITY, 1).expect("valid position");
assert!(polydigest::verify_inclusion(&h, &leaf, &sk, &path, &root));
assert_eq!(leaf, Sha256::digest(b"second").to_vec());
}
#[test]
fn overwrite_changes_the_root() {
let mut tree = Emt::new();
tree.set(0, b"a".to_vec()).unwrap();
tree.set(1, b"b".to_vec()).unwrap();
let before = tree.root().unwrap();
tree.set(0, b"A".to_vec()).unwrap();
assert_ne!(tree.root().unwrap(), before);
}
#[test]
fn seals_one_way() {
let mut tree = Emt::new();
tree.set(0, b"x".to_vec()).unwrap();
let sealed = tree.seal().unwrap();
assert_eq!(sealed.tree_size(), 1);
}
}