refactor: improve tests, small fixes
This commit is contained in:
parent
218a64b5c7
commit
766702e37a
4 changed files with 158 additions and 30 deletions
|
@ -40,7 +40,7 @@ mod node;
|
|||
pub use node::InnerNodeInfo;
|
||||
|
||||
mod partial_mt;
|
||||
pub use partial_mt::PartialMerkleTree;
|
||||
pub use partial_mt::{pmt_to_text, PartialMerkleTree};
|
||||
|
||||
// ERRORS
|
||||
// ================================================================================================
|
||||
|
|
|
@ -2,9 +2,8 @@ use super::{
|
|||
BTreeMap, BTreeSet, MerkleError, MerklePath, NodeIndex, Rpo256, RpoDigest, ValuePath, Vec,
|
||||
Word, EMPTY_WORD,
|
||||
};
|
||||
|
||||
extern crate alloc;
|
||||
use alloc::format;
|
||||
use crate::utils::{format, string::String, word_to_hex};
|
||||
use core::fmt;
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests;
|
||||
|
@ -102,10 +101,8 @@ impl PartialMerkleTree {
|
|||
self.leaves.contains(&index)
|
||||
}
|
||||
|
||||
pub fn get_leaf_depth(&self, index: u64) -> u8 {
|
||||
let trunc_value = index.min(2_u64.pow(self.max_depth() as u32));
|
||||
let mut node_index =
|
||||
NodeIndex::new(self.max_depth(), trunc_value).expect("Truncated value is not valid.");
|
||||
pub fn get_leaf_depth(&self, node_index: NodeIndex) -> u8 {
|
||||
let mut node_index = node_index;
|
||||
for _ in 0..node_index.depth() {
|
||||
if self.leaves.contains(&node_index) {
|
||||
return node_index.depth();
|
||||
|
@ -304,3 +301,32 @@ impl PartialMerkleTree {
|
|||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Utility to visualize a [PartialMerkleTree] in text.
|
||||
pub fn pmt_to_text(tree: &PartialMerkleTree) -> Result<String, fmt::Error> {
|
||||
let indent = " ";
|
||||
let mut s = String::new();
|
||||
s.push_str("root: ");
|
||||
s.push_str(&word_to_hex(&tree.root())?);
|
||||
s.push('\n');
|
||||
for d in 1..=tree.max_depth() {
|
||||
let entries = 2u64.pow(d.into());
|
||||
for i in 0..entries {
|
||||
let index = NodeIndex::new(d, i).expect("The index must always be valid");
|
||||
let node = tree.get_node(index);
|
||||
let node = match node {
|
||||
Err(_) => continue,
|
||||
Ok(node) => node,
|
||||
};
|
||||
|
||||
for _ in 0..d {
|
||||
s.push_str(indent);
|
||||
}
|
||||
s.push_str(&format!("({}, {}): ", index.depth(), index.value()));
|
||||
s.push_str(&word_to_hex(&node)?);
|
||||
s.push('\n');
|
||||
}
|
||||
}
|
||||
|
||||
Ok(s)
|
||||
}
|
||||
|
|
|
@ -1,14 +1,17 @@
|
|||
use super::{
|
||||
super::{int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree},
|
||||
Word,
|
||||
Rpo256, ValuePath, Vec, Word,
|
||||
};
|
||||
|
||||
// TEST DATA
|
||||
// ================================================================================================
|
||||
|
||||
const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0);
|
||||
const NODE11: NodeIndex = NodeIndex::new_unchecked(1, 1);
|
||||
|
||||
const NODE20: NodeIndex = NodeIndex::new_unchecked(2, 0);
|
||||
const NODE22: NodeIndex = NodeIndex::new_unchecked(2, 2);
|
||||
const NODE23: NodeIndex = NodeIndex::new_unchecked(2, 3);
|
||||
|
||||
const NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2);
|
||||
const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3);
|
||||
|
@ -39,7 +42,7 @@ fn get_root() {
|
|||
|
||||
let path33 = ms.get_path(expected_root, NODE33).unwrap();
|
||||
|
||||
let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
|
||||
let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
|
||||
|
||||
assert_eq!(pmt.root(), expected_root.into());
|
||||
}
|
||||
|
@ -54,17 +57,19 @@ fn add_and_get_paths() {
|
|||
let expected_path33 = ms.get_path(expected_root, NODE33).unwrap();
|
||||
let expected_path22 = ms.get_path(expected_root, NODE22).unwrap();
|
||||
|
||||
let pmt = PartialMerkleTree::with_paths([
|
||||
(3_u64, expected_path33.value.into(), expected_path33.path.clone()),
|
||||
(2, expected_path22.value.into(), expected_path22.path.clone()),
|
||||
])
|
||||
.unwrap();
|
||||
let mut pmt = PartialMerkleTree::new();
|
||||
pmt.add_path(3, expected_path33.value.into(), expected_path33.path.clone())
|
||||
.unwrap();
|
||||
pmt.add_path(2, expected_path22.value.into(), expected_path22.path.clone())
|
||||
.unwrap();
|
||||
|
||||
let path33 = pmt.get_path(NODE33).unwrap();
|
||||
let path22 = pmt.get_path(NODE22).unwrap();
|
||||
let actual_root = pmt.root();
|
||||
|
||||
assert_eq!(expected_path33.path, path33);
|
||||
assert_eq!(expected_path22.path, path22);
|
||||
assert_eq!(expected_root, *actual_root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
@ -76,7 +81,7 @@ fn get_node() {
|
|||
|
||||
let path33 = ms.get_path(expected_root, NODE33).unwrap();
|
||||
|
||||
let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
|
||||
let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
|
||||
|
||||
assert_eq!(ms.get_node(expected_root, NODE32).unwrap(), *pmt.get_node(NODE32).unwrap());
|
||||
assert_eq!(ms.get_node(expected_root, NODE10).unwrap(), *pmt.get_node(NODE10).unwrap());
|
||||
|
@ -90,17 +95,28 @@ fn update_leaf() {
|
|||
let ms = MerkleStore::from(&mt);
|
||||
let path33 = ms.get_path(root, NODE33).unwrap();
|
||||
|
||||
let mut pmt =
|
||||
PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
|
||||
let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
|
||||
|
||||
let new_value32 = int_to_node(132);
|
||||
mt.update_leaf(2_u64, new_value32).unwrap();
|
||||
mt.update_leaf(2, new_value32).unwrap();
|
||||
let expected_root = mt.root();
|
||||
|
||||
pmt.update_leaf(NODE32, new_value32.into()).unwrap();
|
||||
let actual_root = pmt.root();
|
||||
|
||||
assert_eq!(expected_root, *actual_root);
|
||||
|
||||
let mut new_vals = VALUES8.clone();
|
||||
new_vals[1] = int_to_node(131);
|
||||
new_vals[2] = int_to_node(132);
|
||||
let new_value20 = Rpo256::merge(&[new_vals[0].into(), new_vals[1].into()]);
|
||||
let mt = MerkleTree::new(new_vals.to_vec()).unwrap();
|
||||
let expected_root = mt.root();
|
||||
|
||||
pmt.update_leaf(NODE20, new_value20).unwrap();
|
||||
let actual_root = pmt.root();
|
||||
|
||||
assert_eq!(expected_root, *actual_root);
|
||||
}
|
||||
|
||||
#[test]
|
||||
|
@ -112,17 +128,97 @@ fn check_leaf_depth() {
|
|||
|
||||
let path33 = ms.get_path(expected_root, NODE33).unwrap();
|
||||
|
||||
let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
|
||||
let pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
|
||||
|
||||
assert_eq!(pmt.get_leaf_depth(0), 2);
|
||||
assert_eq!(pmt.get_leaf_depth(1), 2);
|
||||
assert_eq!(pmt.get_leaf_depth(2), 3);
|
||||
assert_eq!(pmt.get_leaf_depth(3), 3);
|
||||
assert_eq!(pmt.get_leaf_depth(4), 1);
|
||||
assert_eq!(pmt.get_leaf_depth(5), 1);
|
||||
assert_eq!(pmt.get_leaf_depth(6), 1);
|
||||
assert_eq!(pmt.get_leaf_depth(7), 1);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(4, 1)), 2);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(4, 6)), 3);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(4, 10)), 1);
|
||||
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(3, 1)), 2);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(3, 2)), 3);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(3, 5)), 1);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(3, 7)), 1);
|
||||
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(2, 0)), 2);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(2, 1)), 0);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(2, 2)), 1);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(2, 3)), 1);
|
||||
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(1, 0)), 0);
|
||||
assert_eq!(pmt.get_leaf_depth(NodeIndex::make(1, 1)), 1);
|
||||
}
|
||||
|
||||
// TODO: add test for add_path function and check correctness of leaf determination (requires
|
||||
// inner_nodes iter)
|
||||
#[test]
|
||||
fn get_paths() {
|
||||
let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
|
||||
let expected_root = mt.root();
|
||||
|
||||
let ms = MerkleStore::from(&mt);
|
||||
|
||||
let path33 = ms.get_path(expected_root, NODE33).unwrap();
|
||||
let path22 = ms.get_path(expected_root, NODE22).unwrap();
|
||||
|
||||
let mut pmt = PartialMerkleTree::new();
|
||||
pmt.add_path(3, path33.value.into(), path33.path.clone()).unwrap();
|
||||
pmt.add_path(2, path22.value.into(), path22.path.clone()).unwrap();
|
||||
|
||||
let leaves = vec![NODE20, NODE22, NODE23, NODE32, NODE33];
|
||||
let expected_paths: Vec<(NodeIndex, ValuePath)> = leaves
|
||||
.iter()
|
||||
.map(|&leaf| {
|
||||
(
|
||||
leaf,
|
||||
ValuePath {
|
||||
value: mt.get_node(leaf).unwrap().into(),
|
||||
path: mt.get_path(leaf).unwrap(),
|
||||
},
|
||||
)
|
||||
})
|
||||
.collect();
|
||||
|
||||
let actual_paths = pmt.paths();
|
||||
|
||||
assert_eq!(expected_paths, actual_paths);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn leaves() {
|
||||
let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
|
||||
let expected_root = mt.root();
|
||||
|
||||
let ms = MerkleStore::from(&mt);
|
||||
|
||||
let path33 = ms.get_path(expected_root, NODE33).unwrap();
|
||||
let path22 = ms.get_path(expected_root, NODE22).unwrap();
|
||||
|
||||
let mut pmt = PartialMerkleTree::with_paths([(3, path33.value.into(), path33.path)]).unwrap();
|
||||
|
||||
let value11 = mt.get_node(NODE11).unwrap().into();
|
||||
let value20 = mt.get_node(NODE20).unwrap().into();
|
||||
let value32 = mt.get_node(NODE32).unwrap().into();
|
||||
let value33 = mt.get_node(NODE33).unwrap().into();
|
||||
|
||||
let leaves = vec![(NODE11, value11), (NODE20, value20), (NODE32, value32), (NODE33, value33)];
|
||||
|
||||
let expected_leaves = leaves.iter().map(|&tuple| tuple);
|
||||
assert!(expected_leaves.eq(pmt.leaves()));
|
||||
|
||||
pmt.add_path(2, path22.value.into(), path22.path).unwrap();
|
||||
|
||||
let value20 = mt.get_node(NODE20).unwrap().into();
|
||||
let value22 = mt.get_node(NODE22).unwrap().into();
|
||||
let value23 = mt.get_node(NODE23).unwrap().into();
|
||||
let value32 = mt.get_node(NODE32).unwrap().into();
|
||||
let value33 = mt.get_node(NODE33).unwrap().into();
|
||||
|
||||
let leaves = vec![
|
||||
(NODE20, value20),
|
||||
(NODE22, value22),
|
||||
(NODE23, value23),
|
||||
(NODE32, value32),
|
||||
(NODE33, value33),
|
||||
];
|
||||
|
||||
let expected_leaves = leaves.iter().map(|&tuple| tuple);
|
||||
assert!(expected_leaves.eq(pmt.leaves()));
|
||||
}
|
||||
|
|
|
@ -2,6 +2,12 @@ use super::Word;
|
|||
use crate::utils::string::String;
|
||||
use core::fmt::{self, Write};
|
||||
|
||||
#[cfg(not(feature = "std"))]
|
||||
pub use alloc::format;
|
||||
|
||||
#[cfg(feature = "std")]
|
||||
pub use std::format;
|
||||
|
||||
// RE-EXPORTS
|
||||
// ================================================================================================
|
||||
pub use winter_utils::{
|
||||
|
|
Loading…
Add table
Reference in a new issue