refactor: optimize code, fix bugs

This commit is contained in:
Andrey Khmuro
2023-06-05 18:02:16 +03:00
parent 43f1a4cb64
commit 2708a23649
3 changed files with 93 additions and 131 deletions

View File

@@ -1,18 +1,29 @@
use crate::hash::rpo::RpoDigest;
use super::{
super::{int_to_digest, int_to_node, NodeIndex},
PartialMerkleTree, Rpo256,
super::{int_to_node, MerkleStore, MerkleTree, NodeIndex, PartialMerkleTree},
Word,
};
// TEST DATA
// ================================================================================================
const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0);
const NODE22: NodeIndex = NodeIndex::new_unchecked(2, 2);
const NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2);
const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3);
const VALUES8: [Word; 8] = [
int_to_node(1),
int_to_node(2),
int_to_node(3),
int_to_node(4),
int_to_node(5),
int_to_node(6),
int_to_node(7),
int_to_node(8),
];
// TESTS
// ================================================================================================
@@ -21,107 +32,92 @@ const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3);
#[test]
fn get_root() {
let leaf0 = int_to_digest(0);
let leaf1 = int_to_digest(1);
let leaf2 = int_to_digest(2);
let leaf3 = int_to_digest(3);
let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
let expected_root = mt.root();
let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
let mut store = MerkleStore::new();
let ms = MerkleStore::extend(&mut store, mt.inner_nodes());
let root_exp = calculate_parent_hash(parent0, 0, parent1);
let path33 = ms.get_path(expected_root, NODE33).unwrap();
let set =
super::PartialMerkleTree::with_paths([(0, leaf0, vec![*leaf1, *parent1].into())]).unwrap();
let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
assert_eq!(set.root(), root_exp);
assert_eq!(pmt.root(), expected_root.into());
}
#[test]
fn add_and_get_paths() {
let value32 = int_to_digest(32);
let value33 = int_to_digest(33);
let value20 = int_to_digest(20);
let value22 = int_to_digest(22);
let value23 = int_to_digest(23);
let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
let expected_root = mt.root();
let value21 = Rpo256::merge(&[value32, value33]);
let value10 = Rpo256::merge(&[value20, value21]);
let value11 = Rpo256::merge(&[value22, value23]);
let mut store = MerkleStore::new();
let ms = MerkleStore::extend(&mut store, mt.inner_nodes());
let path_33 = vec![*value32, *value20, *value11];
let path_22 = vec![*value23, *value10];
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, value33, path_33.clone().into()),
(2, value22, path_22.clone().into()),
(3_u64, expected_path33.value.into(), expected_path33.path.clone()),
(2, expected_path22.value.into(), expected_path22.path.clone()),
])
.unwrap();
let stored_path_33 = pmt.get_path(NODE33).unwrap();
let stored_path_22 = pmt.get_path(NODE22).unwrap();
assert_eq!(path_33, *stored_path_33);
assert_eq!(path_22, *stored_path_22);
let path33 = pmt.get_path(NODE33).unwrap();
let path22 = pmt.get_path(NODE22).unwrap();
assert_eq!(expected_path33.path, path33);
assert_eq!(expected_path22.path, path22);
}
#[test]
fn get_node() {
let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
let hash_6 = int_to_digest(6);
let index = NodeIndex::make(3, 6);
let pmt = PartialMerkleTree::with_paths([(index.value(), hash_6, path_6.into())]).unwrap();
let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
let expected_root = mt.root();
assert_eq!(int_to_digest(6u64), pmt.get_node(index).unwrap());
let mut store = MerkleStore::new();
let ms = MerkleStore::extend(&mut store, mt.inner_nodes());
let path33 = ms.get_path(expected_root, NODE33).unwrap();
let pmt = PartialMerkleTree::with_paths([(3_u64, 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());
}
#[test]
fn update_leaf() {
let value32 = int_to_digest(32);
let value33 = int_to_digest(33);
let value20 = int_to_digest(20);
let value22 = int_to_digest(22);
let value23 = int_to_digest(23);
let mut mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
let root = mt.root();
let value21 = Rpo256::merge(&[value32, value33]);
let value10 = Rpo256::merge(&[value20, value21]);
let value11 = Rpo256::merge(&[value22, value23]);
let path_33 = vec![*value32, *value20, *value11];
let path_22 = vec![*value23, *value10];
let mut store = MerkleStore::new();
let ms = MerkleStore::extend(&mut store, mt.inner_nodes());
let path33 = ms.get_path(root, NODE33).unwrap();
let mut pmt =
PartialMerkleTree::with_paths([(3, value33, path_33.into()), (2, value22, path_22.into())])
.unwrap();
PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
let new_value32 = int_to_digest(132);
let new_value21 = Rpo256::merge(&[new_value32, value33]);
let new_value10 = Rpo256::merge(&[value20, new_value21]);
let expected_root = Rpo256::merge(&[new_value10, value11]);
let new_value32 = int_to_node(132);
mt.update_leaf(2_u64, new_value32).unwrap();
let expected_root = mt.root();
let old_leaf = pmt.update_leaf(NODE32, new_value32).unwrap();
pmt.update_leaf(NODE32, new_value32.into()).unwrap();
let actual_root = pmt.root();
assert_eq!(value32, old_leaf);
let new_root = pmt.root();
assert_eq!(new_root, expected_root);
assert_eq!(expected_root, *actual_root);
}
#[test]
fn check_leaf_depth() {
let value32 = int_to_digest(32);
let value33 = int_to_digest(33);
let value20 = int_to_digest(20);
let value22 = int_to_digest(22);
let value23 = int_to_digest(23);
let mt = MerkleTree::new(VALUES8.to_vec()).unwrap();
let expected_root = mt.root();
let value11 = Rpo256::merge(&[value22, value23]);
let mut store = MerkleStore::new();
let ms = MerkleStore::extend(&mut store, mt.inner_nodes());
let path_33 = vec![*value32, *value20, *value11];
let path33 = ms.get_path(expected_root, NODE33).unwrap();
let pmt = PartialMerkleTree::with_paths([(3, value33, path_33.into())]).unwrap();
let pmt = PartialMerkleTree::with_paths([(3_u64, path33.value.into(), path33.path)]).unwrap();
assert_eq!(pmt.get_leaf_depth(0).unwrap(), 2);
assert_eq!(pmt.get_leaf_depth(1).unwrap(), 2);
@@ -131,23 +127,8 @@ fn check_leaf_depth() {
assert_eq!(pmt.get_leaf_depth(5).unwrap(), 1);
assert_eq!(pmt.get_leaf_depth(6).unwrap(), 1);
assert_eq!(pmt.get_leaf_depth(7).unwrap(), 1);
assert!(pmt.get_leaf_depth(8).is_err());
}
// TODO: add test for add_path function and check correctness of leaf determination (requires
// inner_nodes iter)
// HELPER FUNCTIONS
// --------------------------------------------------------------------------------------------
/// Calculates the hash of the parent node by two sibling ones
/// - node — current node
/// - node_pos — position of the current node
/// - sibling — neighboring vertex in the tree
fn calculate_parent_hash(node: RpoDigest, node_pos: u64, sibling: RpoDigest) -> RpoDigest {
let parity = node_pos & 1;
if parity == 0 {
Rpo256::merge(&[node, sibling])
} else {
Rpo256::merge(&[sibling, node])
}
}