Merge pull request #50 from HerodotusDev/optimization

optimization
HerodotusDev · Jan 22, 2024 · ae389ad · ae389ad
2 parents 4aaab0d + e229de4
commit ae389ad
Showing 1 changed file with 20 additions and 10 deletions.
diff --git a/src/mmr/core.rs b/src/mmr/core.rs
@@ -226,8 +226,10 @@ impl MMR {
         }
 
         let options = options.unwrap_or_default();
-        let element_count = self.elements_count.get().await?;
-        let tree_size = options.elements_count.unwrap_or(element_count);
+        let tree_size = match options.elements_count {
+            Some(count) => count,
+            None => self.elements_count.get().await?,
+        };
 
         if element_index > tree_size {
             return Err(MMRError::InvalidElementIndex);
@@ -284,8 +286,10 @@ impl MMR {
         options: Option<ProofOptions>,
     ) -> Result<Vec<Proof>, MMRError> {
         let options = options.unwrap_or_default();
-        let element_count = self.elements_count.get().await?;
-        let tree_size = options.elements_count.unwrap_or(element_count);
+        let tree_size = match options.elements_count {
+            Some(count) => count,
+            None => self.elements_count.get().await?,
+        };
 
         for &element_index in &elements_indexes {
             if element_index == 0 {
@@ -349,8 +353,10 @@ impl MMR {
         options: Option<ProofOptions>,
     ) -> Result<bool, MMRError> {
         let options = options.unwrap_or_default();
-        let element_count = self.elements_count.get().await?;
-        let tree_size = options.elements_count.unwrap_or(element_count);
+        let tree_size = match options.elements_count {
+            Some(count) => count,
+            None => self.elements_count.get().await?,
+        };
 
         let leaf_count = mmr_size_to_leaf_count(tree_size);
         let peaks_count = leaf_count_to_peaks_count(leaf_count);
@@ -418,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?;
@@ -458,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?;