From 55cc71dadfc5ce801b083e4f274196523cced61c Mon Sep 17 00:00:00 2001 From: Andrey Khmuro Date: Sat, 3 Jun 2023 15:38:24 +0300 Subject: [PATCH] fix: fix add_path func leaf determination --- src/merkle/partial_mt/mod.rs | 12 +++++------- src/merkle/partial_mt/tests.rs | 3 +++ 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/src/merkle/partial_mt/mod.rs b/src/merkle/partial_mt/mod.rs index be46802..26eb8a3 100644 --- a/src/merkle/partial_mt/mod.rs +++ b/src/merkle/partial_mt/mod.rs @@ -213,19 +213,17 @@ impl PartialMerkleTree { // insert calculated node to the nodes map self.nodes.insert(index_value, node); + // if the calculated node was a leaf, remove it from leaves set. + if self.leaves.contains(&index_value) { + self.leaves.remove(&index_value); + } + let sibling_node = index_value.sibling(); // node became a leaf only if it is a new node (it wasn't in nodes map) if !self.nodes.contains_key(&sibling_node) { self.leaves.insert(sibling_node); } - // node stops being a leaf if the path contains a node which is a child of this leaf - let mut parent = index_value; - parent.move_up(); - if self.leaves.contains(&parent) { - self.leaves.remove(&parent); - } - // insert node from Merkle path to the nodes map self.nodes.insert(sibling_node, hash.into()); diff --git a/src/merkle/partial_mt/tests.rs b/src/merkle/partial_mt/tests.rs index 1efbb6a..35d41ef 100644 --- a/src/merkle/partial_mt/tests.rs +++ b/src/merkle/partial_mt/tests.rs @@ -133,6 +133,9 @@ fn check_leaf_depth() { assert_eq!(pmt.get_leaf_depth(7).unwrap(), 1); } +// TODO: add test for add_path function and check correctness of leaf determination (requires +// inner_nodes iter) + // HELPER FUNCTIONS // --------------------------------------------------------------------------------------------