all tree size update

HerodotusDev · Jan 22, 2024 · e229de4 · e229de4
1 parent e87018e
commit e229de4
Showing 1 changed file with 8 additions and 4 deletions.
diff --git a/src/mmr/core.rs b/src/mmr/core.rs
@@ -424,8 +424,10 @@ impl MMR {
     }
 
     pub async fn get_peaks(&self, option: PeaksOptions) -> Result<Vec<String>, MMRError> {
-        let elements_count = self.elements_count.get().await?;
-        let tree_size = option.elements_count.unwrap_or(elements_count);
+        let tree_size = match option.elements_count {
+            Some(count) => count,
+            None => self.elements_count.get().await?,
+        };
 
         let peaks_idxs = find_peaks(tree_size);
         let peaks = self.retrieve_peaks_hashes(peaks_idxs, None).await?;
@@ -464,8 +466,10 @@ impl MMR {
     }
 
     pub async fn bag_the_peaks(&self, elements_count: Option<usize>) -> Result<String, MMRError> {
-        let element_count_result = self.elements_count.get().await;
-        let tree_size = elements_count.unwrap_or(element_count_result?);
+        let tree_size = match elements_count {
+            Some(count) => count,
+            None => self.elements_count.get().await?,
+        };
         let peaks_idxs = find_peaks(tree_size);
 
         let peaks_hashes = self.retrieve_peaks_hashes(peaks_idxs.clone(), None).await?;