pub type Mmb<D> = crate::merkle::mem::Mem<super::Family, D>;
pub type Config<D> = crate::merkle::mem::Config<super::Family, D>;
#[cfg(test)]
mod tests {
use super::*;
use crate::merkle::{
hasher::{Hasher as _, Standard},
mmb::{Error, Location, Position},
};
use commonware_cryptography::Sha256;
type D = <Sha256 as commonware_cryptography::Hasher>::Digest;
type H = Standard<Sha256>;
fn build_mmb(n: u64) -> (H, Mmb<D>) {
let hasher = H::new();
let mut mmb = Mmb::new(&hasher);
let batch = {
let mut batch = mmb.new_batch();
for i in 0..n {
batch = batch.add(&hasher, &i.to_be_bytes());
}
batch.merkleize(&mmb, &hasher)
};
mmb.apply_batch(&batch).unwrap();
(hasher, mmb)
}
#[test]
fn test_append_and_size() {
let hasher = H::new();
let mut mmb = Mmb::new(&hasher);
for i in 0u64..8 {
let batch = {
let mut batch = mmb.new_batch();
let loc = batch.leaves();
batch = batch.add(&hasher, &i.to_be_bytes());
assert_eq!(*loc, i);
batch.merkleize(&mmb, &hasher)
};
mmb.apply_batch(&batch).unwrap();
}
assert_eq!(*mmb.leaves(), 8);
assert_eq!(*mmb.size(), 13);
}
#[test]
fn test_add_eight_values_structure() {
let (hasher, mmb) = build_mmb(8);
assert_eq!(mmb.bounds().start, Location::new(0));
assert_eq!(mmb.size(), Position::new(13));
assert_eq!(mmb.leaves(), Location::new(8));
let peaks: Vec<(Position, u32)> = mmb.peak_iterator().collect();
assert_eq!(
peaks,
vec![
(Position::new(7), 2),
(Position::new(9), 1),
(Position::new(12), 1)
],
"MMB peaks not as expected"
);
let leaf_positions = [0u64, 1, 3, 4, 6, 8, 10, 11];
for (i, pos) in leaf_positions.into_iter().enumerate() {
let expected = hasher.leaf_digest(Position::new(pos), &(i as u64).to_be_bytes());
assert_eq!(
mmb.get_node(Position::new(pos)).unwrap(),
expected,
"leaf digest mismatch at location {i}"
);
}
let digest2 = hasher.node_digest(
Position::new(2),
&mmb.get_node(Position::new(0)).unwrap(),
&mmb.get_node(Position::new(1)).unwrap(),
);
assert_eq!(mmb.get_node(Position::new(2)).unwrap(), digest2);
let digest5 = hasher.node_digest(
Position::new(5),
&mmb.get_node(Position::new(3)).unwrap(),
&mmb.get_node(Position::new(4)).unwrap(),
);
assert_eq!(mmb.get_node(Position::new(5)).unwrap(), digest5);
let digest7 = hasher.node_digest(Position::new(7), &digest2, &digest5);
assert_eq!(mmb.get_node(Position::new(7)).unwrap(), digest7);
let digest9 = hasher.node_digest(
Position::new(9),
&mmb.get_node(Position::new(6)).unwrap(),
&mmb.get_node(Position::new(8)).unwrap(),
);
assert_eq!(mmb.get_node(Position::new(9)).unwrap(), digest9);
let digest12 = hasher.node_digest(
Position::new(12),
&mmb.get_node(Position::new(10)).unwrap(),
&mmb.get_node(Position::new(11)).unwrap(),
);
assert_eq!(mmb.get_node(Position::new(12)).unwrap(), digest12);
let expected_root = hasher.root(Location::new(8), [digest7, digest9, digest12].iter());
assert_eq!(*mmb.root(), expected_root, "incorrect root");
}
#[test]
fn test_prune_and_reinit() {
let (hasher, mut mmb) = build_mmb(24);
let root = *mmb.root();
let prune_loc = Location::new(9);
mmb.prune(prune_loc).unwrap();
assert_eq!(mmb.bounds().start, prune_loc);
assert_eq!(*mmb.root(), root);
assert!(matches!(
mmb.proof(&hasher, Location::new(0)),
Err(Error::ElementPruned(_))
));
for loc in *prune_loc..*mmb.leaves() {
assert!(
mmb.proof(&hasher, Location::new(loc)).is_ok(),
"loc={loc} should remain provable after pruning"
);
}
let mmb_copy = Mmb::init(
Config {
nodes: (*Position::try_from(prune_loc).unwrap()..*mmb.size())
.map(|i| mmb.get_node(Position::new(i)).unwrap())
.collect(),
pruning_boundary: prune_loc,
pinned_nodes: mmb.node_digests_to_pin(prune_loc),
},
&hasher,
)
.unwrap();
assert_eq!(mmb_copy.size(), mmb.size());
assert_eq!(mmb_copy.leaves(), mmb.leaves());
assert_eq!(mmb_copy.bounds(), mmb.bounds());
assert_eq!(*mmb_copy.root(), root);
assert!(mmb_copy.proof(&hasher, Location::new(17)).is_ok());
}
#[test]
fn test_init_size_validation() {
let hasher = H::new();
assert!(Mmb::<D>::init(
Config {
nodes: vec![],
pruning_boundary: Location::new(0),
pinned_nodes: vec![],
},
&hasher,
)
.is_ok());
assert!(matches!(
Mmb::init(
Config {
nodes: vec![hasher.digest(b"node1"), hasher.digest(b"node2")],
pruning_boundary: Location::new(0),
pinned_nodes: vec![],
},
&hasher,
),
Err(Error::InvalidSize(_))
));
assert!(Mmb::init(
Config {
nodes: vec![
hasher.digest(b"leaf1"),
hasher.digest(b"leaf2"),
hasher.digest(b"parent"),
],
pruning_boundary: Location::new(0),
pinned_nodes: vec![],
},
&hasher,
)
.is_ok());
let (_, mmb) = build_mmb(64);
let nodes: Vec<_> = (0..*mmb.size())
.map(|i| *mmb.get_node_unchecked(Position::new(i)))
.collect();
assert!(Mmb::init(
Config {
nodes,
pruning_boundary: Location::new(0),
pinned_nodes: vec![],
},
&hasher,
)
.is_ok());
let (_, mut mmb) = build_mmb(11);
mmb.prune(Location::new(4)).unwrap();
let nodes: Vec<_> = (6..*mmb.size())
.map(|i| *mmb.get_node_unchecked(Position::new(i)))
.collect();
let pinned_nodes = mmb.node_digests_to_pin(Location::new(4));
assert!(Mmb::init(
Config {
nodes: nodes.clone(),
pruning_boundary: Location::new(4),
pinned_nodes: pinned_nodes.clone(),
},
&hasher,
)
.is_ok());
assert!(matches!(
Mmb::init(
Config {
nodes,
pruning_boundary: Location::new(2),
pinned_nodes,
},
&hasher,
),
Err(Error::InvalidSize(_))
));
}
}