From 766702e37af3ff65a147d0786ec19b0cf760e6c5 Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Fri, 9 Jun 2023 13:33:56 +0300 Subject: [PATCH] refactor: improve tests, small fixes --- src/merkle/mod.rs | 2 +- src/merkle/partial_mt/mod.rs | 40 ++++++++-- src/merkle/partial_mt/tests.rs | 140 +++++++++++++++++++++++++++------ src/utils.rs | 6 ++ 4 files changed, 158 insertions(+), 30 deletions(-) diff --git a/src/merkle/mod.rs b/src/merkle/mod.rs index 6f7d2ab..73f149e 100644 --- a/src/merkle/mod.rs +++ b/src/merkle/mod.rs @@ -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 // ================================================================================================ diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index ab7b644..fd9f1c6 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -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 { + 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) +} diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index 6151421..85d3036 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/src/merkle/partial_mt/tests.rs @@ -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(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); +} + +#[test] +fn get_paths() { + let mt = MerkleTree::new(VALUES8.to_vec()).unwrap(); + let expected_root = mt.root(); + + let ms = MerkleStore::from(&mt); - 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); + 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); } -// TODO: add test for add_path function and check correctness of leaf determination (requires -// inner_nodes iter) +#[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())); +} diff --git a/src/utils.rs b/src/utils.rs index 522d9ea..b9b7849 100644 --- a/src/utils.rs +++ b/src/utils.rs @@ -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::{