moved transaction_records_by_id methods into it

zingolib/src/blaze/block_management_reorg_detection.rs

@@ -334,7 +334,7 @@ impl BlockManagementData {
             transaction_metadata_set
                 .write()
                 .await
-                .remove_txns_at_height(reorg_height);
+                .invalidate_all_transactions_after_or_at_height(reorg_height);
         }
     }
 

zingolib/src/wallet/transaction_records_by_id.rs

@@ -1,8 +1,18 @@
 use std::collections::HashMap;
 
+use orchard::note_encryption::OrchardDomain;
+use sapling_crypto::note_encryption::SaplingDomain;
+
+use zcash_note_encryption::Domain;
+use zcash_primitives::consensus::BlockHeight;
+
 use zcash_primitives::transaction::TxId;
 
-use crate::wallet::data::TransactionRecord;
+use crate::wallet::{
+    data::TransactionRecord,
+    notes::NoteInterface,
+    traits::{DomainWalletExt, Recipient},
+};
 
 #[derive(Debug)]
 pub struct TransactionRecordsById(pub HashMap<TxId, TransactionRecord>);
@@ -20,9 +30,115 @@ impl std::ops::DerefMut for TransactionRecordsById {
         &mut self.0
     }
 }
+
+/// Constructors
 impl TransactionRecordsById {
-    // Associated function to create a TransactionRecordMap from a HashMap
+    /// Constructs a new TransactionRecordsById with an empty map.
+    pub fn new() -> Self {
+        TransactionRecordsById(HashMap::new())
+    }
+    // Constructs a TransactionRecordMap from a HashMap
     pub fn from_map(map: HashMap<TxId, TransactionRecord>) -> Self {
         TransactionRecordsById(map)
     }
 }
+/// Methods to modify the map.
+impl TransactionRecordsById {
+    /// Invalidates all transactions from a given height including the block with block height `reorg_height`
+    ///
+    /// All information above a certain height is invalidated during a reorg.
+    pub fn invalidate_all_transactions_after_or_at_height(&mut self, reorg_height: BlockHeight) {
+        // First, collect txids that need to be removed
+        let txids_to_remove = self
+            .values()
+            .filter_map(|transaction_metadata| {
+                if transaction_metadata
+                    .status
+                    .is_confirmed_after_or_at(&reorg_height)
+                    || transaction_metadata
+                        .status
+                        .is_broadcast_after_or_at(&reorg_height)
+                // tODo: why dont we only remove confirmed transactions. unconfirmed transactions may still be valid in the mempool and may later confirm or expire.
+                {
+                    Some(transaction_metadata.txid)
+                } else {
+                    None
+                }
+            })
+            .collect::<Vec<_>>();
+        self.invalidate_transactions(txids_to_remove);
+    }
+    /// Invalidiates a vec of txids by removing them and then all references to them.
+    ///
+    /// A transaction can be invalidated either by a reorg or if it was never confirmed by a miner.
+    /// This is required in the case that a note was spent in a invalidated transaction.
+    /// Takes a slice of txids corresponding to the invalidated transactions, searches all notes for being spent in one of those txids, and resets them to unspent.
+    pub(crate) fn invalidate_transactions(&mut self, txids_to_remove: Vec<TxId>) {
+        for txid in &txids_to_remove {
+            self.remove(txid);
+        }
+
+        self.invalidate_transaction_specific_transparent_spends(&txids_to_remove);
+        // roll back any sapling spends in each invalidated tx
+        self.invalidate_transaction_specific_domain_spends::<SaplingDomain>(&txids_to_remove);
+        // roll back any orchard spends in each invalidated tx
+        self.invalidate_transaction_specific_domain_spends::<OrchardDomain>(&txids_to_remove);
+    }
+    /// Reverts any spent transparent notes in the given transactions to unspent.
+    pub(crate) fn invalidate_transaction_specific_transparent_spends(
+        &mut self,
+        invalidated_txids: &[TxId],
+    ) {
+        self.values_mut().for_each(|transaction_metadata| {
+            // Update UTXOs to roll back any spent utxos
+            transaction_metadata
+                .transparent_notes
+                .iter_mut()
+                .for_each(|utxo| {
+                    if utxo.is_spent() && invalidated_txids.contains(&utxo.spent().unwrap().0) {
+                        *utxo.spent_mut() = None;
+                    }
+
+                    if utxo.unconfirmed_spent.is_some()
+                        && invalidated_txids.contains(&utxo.unconfirmed_spent.unwrap().0)
+                    {
+                        utxo.unconfirmed_spent = None;
+                    }
+                })
+        });
+    }
+    /// Reverts any spent shielded notes in the given transactions to unspent.
+    pub(crate) fn invalidate_transaction_specific_domain_spends<D: DomainWalletExt>(
+        &mut self,
+        invalidated_txids: &[TxId],
+    ) where
+        <D as Domain>::Recipient: Recipient,
+        <D as Domain>::Note: PartialEq + Clone,
+    {
+        self.values_mut().for_each(|transaction_metadata| {
+            // Update notes to rollback any spent notes
+            D::to_notes_vec_mut(transaction_metadata)
+                .iter_mut()
+                .for_each(|nd| {
+                    // Mark note as unspent if the txid being removed spent it.
+                    if nd.spent().is_some() && invalidated_txids.contains(&nd.spent().unwrap().0) {
+                        *nd.spent_mut() = None;
+                    }
+
+                    // Remove unconfirmed spends too
+                    if nd.pending_spent().is_some()
+                        && invalidated_txids.contains(&nd.pending_spent().unwrap().0)
+                    {
+                        *nd.pending_spent_mut() = None;
+                    }
+                });
+        });
+    }
+}
+
+impl Default for TransactionRecordsById {
+    /// Default constructor
+    fn default() -> Self {
+        Self::new()
+    }
+}

zingolib/src/wallet/transactions.rs

@@ -1,5 +1,3 @@
-use std::collections::HashMap;
-
 use crate::wallet::{data::WitnessTrees, transaction_records_by_id::TransactionRecordsById};
 
 /// HashMap of all transactions in a wallet, keyed by txid.
@@ -17,13 +15,13 @@ pub mod recording;
 impl TxMapAndMaybeTrees {
     pub(crate) fn new_with_witness_trees() -> TxMapAndMaybeTrees {
         Self {
-            transaction_records_by_id: TransactionRecordsById(HashMap::new()),
+            transaction_records_by_id: TransactionRecordsById::new(),
             witness_trees: Some(WitnessTrees::default()),
         }
     }
     pub(crate) fn new_treeless() -> TxMapAndMaybeTrees {
         Self {
-            transaction_records_by_id: TransactionRecordsById(HashMap::new()),
+            transaction_records_by_id: TransactionRecordsById::new(),
             witness_trees: None,
         }
     }

zingolib/src/wallet/transactions/recording.rs

@@ -24,85 +24,13 @@ use crate::{
 
 use super::TxMapAndMaybeTrees;
 impl TxMapAndMaybeTrees {
-    pub fn remove_txids(&mut self, txids_to_remove: Vec<TxId>) {
-        for txid in &txids_to_remove {
-            self.transaction_records_by_id.remove(txid);
-        }
-        self.transaction_records_by_id
-            .values_mut()
-            .for_each(|transaction_metadata| {
-                // Update UTXOs to rollback any spent utxos
-                transaction_metadata
-                    .transparent_notes
-                    .iter_mut()
-                    .for_each(|utxo| {
-                        if utxo.is_spent() && txids_to_remove.contains(&utxo.spent().unwrap().0) {
-                            *utxo.spent_mut() = None;
-                        }
-
-                        if utxo.unconfirmed_spent.is_some()
-                            && txids_to_remove.contains(&utxo.unconfirmed_spent.unwrap().0)
-                        {
-                            utxo.unconfirmed_spent = None;
-                        }
-                    })
-            });
-        self.remove_domain_specific_txids::<SaplingDomain>(&txids_to_remove);
-        self.remove_domain_specific_txids::<OrchardDomain>(&txids_to_remove);
-    }
+    /// During reorgs, we need to remove all txns at a given height, and all spends that refer to any removed txns.
+    pub fn invalidate_all_transactions_after_or_at_height(&mut self, reorg_height: u64) {
+        let reorg_height = BlockHeight::from_u32(reorg_height as u32);
 
-    fn remove_domain_specific_txids<D: DomainWalletExt>(&mut self, txids_to_remove: &[TxId])
-    where
-        <D as Domain>::Recipient: Recipient,
-        <D as Domain>::Note: PartialEq + Clone,
-    {
         self.transaction_records_by_id
-            .values_mut()
-            .for_each(|transaction_metadata| {
-                // Update notes to rollback any spent notes
-                D::to_notes_vec_mut(transaction_metadata)
-                    .iter_mut()
-                    .for_each(|nd| {
-                        // Mark note as unspent if the txid being removed spent it.
-                        if nd.spent().is_some() && txids_to_remove.contains(&nd.spent().unwrap().0)
-                        {
-                            *nd.spent_mut() = None;
-                        }
+            .invalidate_all_transactions_after_or_at_height(reorg_height);
 
-                        // Remove unconfirmed spends too
-                        if nd.pending_spent().is_some()
-                            && txids_to_remove.contains(&nd.pending_spent().unwrap().0)
-                        {
-                            *nd.pending_spent_mut() = None;
-                        }
-                    });
-            });
-    }
-
-    // During reorgs, we need to remove all txns at a given height, and all spends that refer to any removed txns.
-    pub fn remove_txns_at_height(&mut self, reorg_height: u64) {
-        let reorg_height = BlockHeight::from_u32(reorg_height as u32);
-
-        // First, collect txids that need to be removed
-        let txids_to_remove = self
-            .transaction_records_by_id
-            .values()
-            .filter_map(|transaction_metadata| {
-                if transaction_metadata
-                    .status
-                    .is_confirmed_after_or_at(&reorg_height)
-                    || transaction_metadata
-                        .status
-                        .is_broadcast_after_or_at(&reorg_height)
-                // tODo: why dont we only remove confirmed transactions. unconfirmed transactions may still be valid in the mempool and may later confirm or expire.
-                {
-                    Some(transaction_metadata.txid)
-                } else {
-                    None
-                }
-            })
-            .collect::<Vec<_>>();
-        self.remove_txids(txids_to_remove);
         if let Some(ref mut t) = self.witness_trees {
             t.witness_tree_sapling
                 .truncate_removing_checkpoint(&(reorg_height - 1))
@@ -114,6 +42,7 @@ impl TxMapAndMaybeTrees {
         }
     }
 
+    /// Invalidates all those transactions which were broadcast but never 'confirmed' accepted by a miner.
     pub(crate) fn clear_expired_mempool(&mut self, latest_height: u64) {
         let cutoff = BlockHeight::from_u32((latest_height.saturating_sub(MAX_REORG as u64)) as u32);
 
@@ -130,7 +59,8 @@ impl TxMapAndMaybeTrees {
             .iter()
             .for_each(|t| println!("Removing expired mempool tx {}", t));
 
-        self.remove_txids(txids_to_remove);
+        self.transaction_records_by_id
+            .invalidate_transactions(txids_to_remove);
     }
 
     // Check this transaction to see if it is an outgoing transaction, and if it is, mark all received notes with non-textual memos in this