Merge pull request #927 from fluidvanadium/methodify_witness_trees

Methodify witness trees
zingolabs · Apr 10, 2024 · 870270e · 870270e
2 parents 36caa15 + 5a96a4d
commit 870270e
Show file tree

Hide file tree

Showing 5 changed files with 25 additions and 34 deletions.
diff --git a/integration-tests/tests/integrations.rs b/integration-tests/tests/integrations.rs
@@ -778,8 +778,7 @@ mod slow {
             .transaction_metadata_set
             .read()
             .await
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .unwrap()
             .witness_tree_orchard
             .marked_positions()
@@ -811,8 +810,7 @@ mod slow {
             .transaction_metadata_set
             .read()
             .await
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .unwrap()
             .witness_tree_orchard
             .marked_positions()
@@ -825,8 +823,7 @@ mod slow {
                 .transaction_metadata_set
                 .read()
                 .await
-                .witness_trees
-                .as_ref()
+                .witness_trees()
                 .unwrap()
                 .witness_tree_orchard
         );
@@ -854,8 +851,7 @@ mod slow {
             .transaction_metadata_set
             .read()
             .await
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .unwrap()
             .witness_tree_orchard
             .marked_positions()
@@ -887,8 +883,7 @@ mod slow {
                 .transaction_metadata_set
                 .read()
                 .await
-                .witness_trees
-                .as_ref()
+                .witness_trees()
                 .unwrap()
                 .witness_tree_orchard
         );
@@ -898,8 +893,7 @@ mod slow {
             .transaction_metadata_set
             .read()
             .await
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .unwrap()
             .witness_tree_orchard
             .marked_positions()
@@ -2128,8 +2122,7 @@ mod slow {
                 .transaction_metadata_set
                 .read()
                 .await
-                .witness_trees
-                .as_ref()
+                .witness_trees()
                 .unwrap()
                 .witness_tree_orchard
                 .max_leaf_position(0),
@@ -2139,8 +2132,7 @@ mod slow {
                 .transaction_metadata_set
                 .read()
                 .await
-                .witness_trees
-                .as_ref()
+                .witness_trees()
                 .unwrap()
                 .witness_tree_orchard
                 .max_leaf_position(0)
@@ -2477,8 +2469,7 @@ mod slow {
                 .transaction_metadata_set
                 .read()
                 .await
-                .witness_trees
-                .as_ref()
+                .witness_trees()
                 .unwrap()
                 .witness_tree_orchard
         );
@@ -2595,7 +2586,7 @@ mod slow {
             .transaction_metadata_set
             .read()
             .await;
-        let wallet_trees = read_lock.witness_trees.as_ref().unwrap();
+        let wallet_trees = read_lock.witness_trees().unwrap();
         let last_leaf = wallet_trees
             .witness_tree_orchard
             .max_leaf_position(0)
@@ -3300,8 +3291,7 @@ mod slow {
             .transaction_metadata_set
             .read()
             .await
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .unwrap()
             .witness_tree_orchard
             .witness_at_checkpoint_depth(
@@ -3340,8 +3330,7 @@ mod slow {
             .transaction_metadata_set
             .read()
             .await
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .unwrap()
             .witness_tree_orchard
             .witness_at_checkpoint_depth(

diff --git a/zingolib/src/wallet/send.rs b/zingolib/src/wallet/send.rs
@@ -272,8 +272,7 @@ impl LightWallet {
             .read()
             .await;
         let witness_trees = txmds_readlock
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .ok_or("No spend capability.")?;
         let (tx_builder, total_shielded_receivers) = match self
             .create_and_populate_tx_builder(

diff --git a/zingolib/src/wallet/traits.rs b/zingolib/src/wallet/traits.rs
@@ -455,16 +455,14 @@ where
         txmds: &TxMapAndMaybeTrees,
     ) -> Option<&MemoryStoreShardTree<<Self::WalletNote as ShieldedNoteInterface>::Node>> {
         txmds
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .map(|trees| Self::get_shardtree(trees))
     }
     fn transaction_metadata_set_to_shardtree_mut(
         txmds: &mut TxMapAndMaybeTrees,
     ) -> Option<&mut MemoryStoreShardTree<<Self::WalletNote as ShieldedNoteInterface>::Node>> {
         txmds
-            .witness_trees
-            .as_mut()
+            .witness_trees_mut()
             .map(|trees| Self::get_shardtree_mut(trees))
     }
     fn get_shardtree(

diff --git a/zingolib/src/wallet/transactions.rs b/zingolib/src/wallet/transactions.rs
@@ -7,7 +7,7 @@ use crate::wallet::{data::WitnessTrees, transaction_records_by_id::TransactionRe
 /// this struct are threadsafe/locked properly.
 pub struct TxMapAndMaybeTrees {
     pub transaction_records_by_id: TransactionRecordsById,
-    pub witness_trees: Option<WitnessTrees>,
+    witness_trees: Option<WitnessTrees>,
 }
 
 pub mod get;
@@ -27,6 +27,12 @@ impl TxMapAndMaybeTrees {
             witness_trees: None,
         }
     }
+    pub fn witness_trees(&self) -> Option<&WitnessTrees> {
+        self.witness_trees.as_ref()
+    }
+    pub(crate) fn witness_trees_mut(&mut self) -> Option<&mut WitnessTrees> {
+        self.witness_trees.as_mut()
+    }
     pub fn clear(&mut self) {
         self.transaction_records_by_id.clear();
         self.witness_trees.as_mut().map(WitnessTrees::clear);

diff --git a/zingolib/src/wallet/witnesses.rs b/zingolib/src/wallet/witnesses.rs
@@ -13,7 +13,7 @@ impl LightWallet {
             .transaction_metadata_set
             .write()
             .await
-            .witness_trees
+            .witness_trees_mut()
         {
             trees.insert_all_frontier_nodes(legacy_sapling_frontier, legacy_orchard_frontier)
         };
@@ -25,7 +25,7 @@ impl LightWallet {
             .transaction_metadata_set
             .write()
             .await;
-        if let Some(ref mut trees) = txmds_writelock.witness_trees {
+        if let Some(ref mut trees) = txmds_writelock.witness_trees_mut() {
             trees
                 .witness_tree_sapling
                 .truncate_removing_checkpoint(&BlockHeight::from(last_synced_height as u32))
@@ -43,8 +43,7 @@ impl LightWallet {
             .transaction_metadata_set
             .read()
             .await
-            .witness_trees
-            .as_ref()
+            .witness_trees()
             .is_some_and(|trees| {
                 trees
                     .witness_tree_orchard