From dcbf6f64d55ce30bd47381e0b49f6844de1bed55 Mon Sep 17 00:00:00 2001
From: Augusto Hack <hack.augusto@gmail.com>
Date: Thu, 25 Jan 2024 16:25:37 +0100
Subject: [PATCH] crypto: PartialMmr api changes (#437)

---
 mock/src/mock/chain.rs               |  2 +-
 objects/src/transaction/chain_mmr.rs | 30 +++-------------------------
 2 files changed, 4 insertions(+), 28 deletions(-)

diff --git a/mock/src/mock/chain.rs b/mock/src/mock/chain.rs
index 4ff8e02e5..5cc661918 100644
--- a/mock/src/mock/chain.rs
+++ b/mock/src/mock/chain.rs
@@ -661,7 +661,7 @@ fn mmr_to_chain_mmr(mmr: &Mmr, blocks: &[BlockHeader]) -> ChainMmr {
     for i in 0..num_leaves {
         let node = mmr.get(i).unwrap();
         let path = mmr.open(i, mmr.forest()).unwrap().merkle_path;
-        partial_mmr.add(i, node, &path).unwrap();
+        partial_mmr.track(i, node, &path).unwrap();
     }
 
     ChainMmr::new(partial_mmr, blocks.to_vec()).unwrap()
diff --git a/objects/src/transaction/chain_mmr.rs b/objects/src/transaction/chain_mmr.rs
index 498bd9534..6686ad347 100644
--- a/objects/src/transaction/chain_mmr.rs
+++ b/objects/src/transaction/chain_mmr.rs
@@ -1,5 +1,5 @@
 use crate::{
-    crypto::merkle::{InnerNodeInfo, MerklePath, MmrDelta, MmrPeaks, PartialMmr},
+    crypto::merkle::{InnerNodeInfo, MmrPeaks, PartialMmr},
     utils::collections::{BTreeMap, Vec},
     BlockHeader, ChainMmrError,
 };
@@ -96,32 +96,8 @@ impl ChainMmr {
     /// Panics if the `block_header.block_num` is not equal to the current chain length (i.e., the
     /// provided block header is not the next block in the chain).
     pub fn add_block(&mut self, block_header: BlockHeader, track: bool) {
-        let block_num = block_header.block_num();
-        let block_hash = block_header.hash();
-
-        assert_eq!(block_num, self.chain_length() as u32);
-
-        // save the original peaks so that we can construct a Merkle path from them later
-        let mut original_peaks = self.mmr.peaks().peaks().to_vec();
-        original_peaks.reverse();
-
-        // update the partial MMR
-        let delta = MmrDelta {
-            forest: self.mmr.forest() + 1,
-            data: vec![block_header.hash()],
-        };
-        self.mmr.apply(delta).expect("failed to add a block to the partial MMR");
-
-        // if we want to track authentication path for this block, add it to the partial MMR and
-        // also add the block header to the block map
-        if track {
-            // path depth is the depth of the smallest tree after the update; this is defined by
-            // number of trailing zeros in the forest (ideally, we'd use a Forest struct for this)
-            let path_depth = self.mmr.forest().trailing_zeros() as usize;
-            let block_path = MerklePath::new(original_peaks[..path_depth].to_vec());
-            self.mmr.add(block_num as usize, block_hash, &block_path).expect("msg");
-            self.blocks.insert(block_num, block_header);
-        }
+        assert_eq!(block_header.block_num(), self.chain_length() as u32);
+        self.mmr.add(block_header.hash(), track);
     }
 
     // ITERATORS