refactor(wallet): introduce a NullifierIndex struct

lez/wallet/src/lib.rs

@@ -7,10 +7,7 @@
     reason = "Most of the shadows come from args parsing which is ok"
 )]
 
-use std::{
-    collections::{HashMap, HashSet},
-    path::PathBuf,
-};
+use std::{collections::HashSet, path::PathBuf};
 
 pub use account_manager::AccountIdentity;
 use anyhow::{Context as _, Result};
@@ -25,7 +22,7 @@ use lee::{
     },
 };
 use lee_core::{
-    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, Nullifier, SharedSecretKey,
+    Commitment, CommitmentSetDigest, DUMMY_COMMITMENT, MembershipProof, SharedSecretKey,
     account::Nonce, compute_digest_for_path, program::InstructionData,
 };
 use log::info;
@@ -37,7 +34,7 @@ use crate::{
     account::{AccountIdWithPrivacy, Label},
     config::WalletConfigOverrides,
     poller::TxPoller,
-    storage::key_chain::SharedAccountEntry,
+    storage::key_chain::{NullifierIndex, SharedAccountEntry},
 };
 
 pub mod account;
@@ -758,7 +755,7 @@ impl WalletCore {
     fn sync_private_accounts_with_tx(
         &mut self,
         tx: LeeTransaction,
-        index: &mut HashMap<Nullifier, AccountId>,
+        index: &mut NullifierIndex,
         handled: &HashSet<usize>,
     ) {
         let LeeTransaction::PrivacyPreserving(tx) = tx else {
@@ -804,28 +801,24 @@ impl WalletCore {
                             let npk = &key_chain.nullifier_public_key;
                             let account_id = lee::AccountId::for_private_account(npk, &kind);
                             let nsk = key_chain.private_key_holder.nullifier_secret_key;
-                            let nullifier = Nullifier::for_account_update(
-                                &Commitment::new(&account_id, &res_acc),
-                                &nsk,
-                            );
-                            (account_id, kind, res_acc, nullifier)
+                            (account_id, kind, res_acc, nsk)
                         })
                     })
                     .collect::<Vec<_>>()
             })
             .collect::<Vec<_>>();
 
-        for (affected_account_id, kind, new_acc, nullifier) in affected_accounts {
+        for (affected_account_id, kind, new_acc, nsk) in affected_accounts {
             info!(
                 "Received new account for account_id {affected_account_id:#?} with account object {new_acc:#?}"
             );
+            // Await the account's next update by its nullifier, so later updates
+            // to it are caught without tag matching.
+            index.track(affected_account_id, &new_acc, &nsk);
             self.storage
                 .key_chain_mut()
                 .insert_private_account(affected_account_id, kind, new_acc)
                 .expect("Account Id should exist");
-            // Await the account's next update by its nullifier, so later updates
-            // to it are caught without tag matching.
-            index.insert(nullifier, affected_account_id);
         }
 
         // Scan for updates to shared accounts (GMS-derived).
@@ -835,7 +828,7 @@ impl WalletCore {
     fn sync_shared_private_accounts_with_tx(
         &mut self,
         tx: &PrivacyPreservingTransaction,
-        index: &mut HashMap<Nullifier, AccountId>,
+        index: &mut NullifierIndex,
         handled: &HashSet<usize>,
     ) {
         let shared_keys: Vec<_> = self
@@ -882,14 +875,10 @@ impl WalletCore {
                         .expect("Ciphertext ID is expected to fit in u32"),
                 ) {
                     info!("Synced shared account {account_id:#?} with new state {new_acc:#?}");
-                    let nullifier = Nullifier::for_account_update(
-                        &Commitment::new(&account_id, &new_acc),
-                        &nsk,
-                    );
+                    index.track(account_id, &new_acc, &nsk);
                     self.storage
                         .key_chain_mut()
                         .update_shared_private_account_state(&account_id, new_acc);
-                    index.insert(nullifier, account_id);
                 }
             }
         }

lez/wallet/src/storage/key_chain.rs

@@ -10,7 +10,8 @@ use key_protocol::key_management::{
 };
 use lee::{Account, AccountId, privacy_preserving_transaction::message::Message};
 use lee_core::{
-    Commitment, EncryptionScheme, Identifier, Nullifier, PrivateAccountKind, SharedSecretKey,
+    Commitment, EncryptionScheme, Identifier, Nullifier, NullifierSecretKey, PrivateAccountKind,
+    SharedSecretKey,
 };
 use log::{debug, warn};
 use serde::{Deserialize, Serialize};
@@ -61,6 +62,41 @@ pub struct SharedAccountEntry {
     pub account: Account,
 }
 
+/// Maps each owned or shared private account to the nullifier its next update will publish,
+/// so sync can spot updates by nullifier rather than view tag.
+#[derive(Default)]
+pub struct NullifierIndex(HashMap<Nullifier, AccountId>);
+
+impl NullifierIndex {
+    fn next_update_nullifier(
+        account_id: AccountId,
+        account: &Account,
+        nsk: &NullifierSecretKey,
+    ) -> Nullifier {
+        Nullifier::for_account_update(&Commitment::new(&account_id, account), nsk)
+    }
+
+    /// Returns the account whose next update would publish `nullifier`.
+    #[must_use]
+    pub fn account_for(&self, nullifier: &Nullifier) -> Option<AccountId> {
+        self.0.get(nullifier).copied()
+    }
+
+    /// Indexes `account_id` by the nullifier its next update will publish.
+    pub fn track(&mut self, account_id: AccountId, account: &Account, nsk: &NullifierSecretKey) {
+        self.0.insert(
+            Self::next_update_nullifier(account_id, account, nsk),
+            account_id,
+        );
+    }
+
+    /// Replaces a spent nullifier with the account's `next` one.
+    pub fn update(&mut self, spent: &Nullifier, next: Nullifier, account_id: AccountId) {
+        self.0.remove(spent);
+        self.0.insert(next, account_id);
+    }
+}
+
 #[derive(Clone, Debug)]
 #[cfg_attr(test, derive(PartialEq, Eq))]
 pub struct UserKeyChain {
@@ -304,8 +340,8 @@ impl UserKeyChain {
     /// Maps each owned and shared account's current-state update nullifier to its `account_id`,
     /// so co-owner updates are found during sync by nullifier rather than view tag.
     #[must_use]
-    pub fn build_latest_nullifier_index(&self) -> HashMap<Nullifier, AccountId> {
-        let mut index = HashMap::new();
+    pub fn build_latest_nullifier_index(&self) -> NullifierIndex {
+        let mut index = NullifierIndex::default();
 
         // For each (regular) found account the user owns, compute its nullifier and put
         // into the map. This is the next nullifier it will look for.
@@ -313,8 +349,7 @@ impl UserKeyChain {
             let account_id =
                 AccountId::for_private_account(&found.key_chain.nullifier_public_key, found.kind);
             let nsk = found.key_chain.private_key_holder.nullifier_secret_key;
-            let commitment = Commitment::new(&account_id, found.account);
-            index.insert(Nullifier::for_account_update(&commitment, &nsk), account_id);
+            index.track(account_id, found.account, &nsk);
         }
 
         // Same for the shared accounts.
@@ -323,8 +358,7 @@ impl UserKeyChain {
                 continue;
             };
             let nsk = keys.nullifier_secret_key;
-            let commitment = Commitment::new(&account_id, &entry.account);
-            index.insert(Nullifier::for_account_update(&commitment, &nsk), account_id);
+            index.track(account_id, &entry.account, &nsk);
         }
 
         index
@@ -336,14 +370,16 @@ impl UserKeyChain {
     pub fn sync_updates_via_nullifiers(
         &mut self,
         message: &Message,
-        index: &mut HashMap<Nullifier, AccountId>,
+        index: &mut NullifierIndex,
     ) -> HashSet<usize> {
         // Get the nullifier information if awaiting the nullifier.
         let hits: Vec<(usize, Nullifier, AccountId)> = message
             .new_nullifiers
             .iter()
             .enumerate()
-            .filter_map(|(i, (nullifier, _))| index.get(nullifier).map(|&id| (i, *nullifier, id)))
+            .filter_map(|(i, (nullifier, _))| {
+                index.account_for(nullifier).map(|id| (i, *nullifier, id))
+            })
             .collect();
 
         let mut handled = HashSet::new();
@@ -351,10 +387,9 @@ impl UserKeyChain {
             // Try decrypting the commitment connectted to the nullifier and get the next
             // nullifier to await.
             if let Some(new_nullifier) = self.apply_nullifier_update(account_id, message, i) {
-                index.remove(&old_nullifier);
                 // Update the index to await for the new state of the account, i.e.
                 // the new nullifier.
-                index.insert(new_nullifier, account_id);
+                index.update(&old_nullifier, new_nullifier, account_id);
                 // Record that this nullifier's position can be skipped for scanning.
                 handled.insert(i);
             }
@@ -397,8 +432,7 @@ impl UserKeyChain {
 
         let (kind, new_account) =
             EncryptionScheme::decrypt(&encrypted.ciphertext, &secret, commitment, ciph_id)?;
-        let new_nullifier =
-            Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk);
+        let new_nullifier = NullifierIndex::next_update_nullifier(account_id, &new_account, &nsk);
 
         if is_shared {
             self.update_shared_private_account_state(&account_id, new_account);
@@ -832,7 +866,7 @@ mod tests {
         let old_nullifier =
             Nullifier::for_account_update(&Commitment::new(&account_id, &old_account), &nsk);
         let mut index = kc.build_latest_nullifier_index();
-        assert_eq!(index.get(&old_nullifier), Some(&account_id));
+        assert_eq!(index.account_for(&old_nullifier), Some(account_id));
 
         let new_account = Account {
             balance: 150,
@@ -870,8 +904,8 @@ mod tests {
         );
         let new_nullifier =
             Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk);
-        assert_eq!(index.get(&new_nullifier), Some(&account_id));
-        assert!(!index.contains_key(&old_nullifier));
+        assert_eq!(index.account_for(&new_nullifier), Some(account_id));
+        assert!(index.account_for(&old_nullifier).is_none());
     }
 
     #[test]
@@ -903,7 +937,7 @@ mod tests {
         let old_nullifier =
             Nullifier::for_account_update(&Commitment::new(&account_id, &old_account), &nsk);
         let mut index = kc.build_latest_nullifier_index();
-        assert_eq!(index.get(&old_nullifier), Some(&account_id));
+        assert_eq!(index.account_for(&old_nullifier), Some(account_id));
 
         let new_account = Account {
             balance: 250,
@@ -935,8 +969,8 @@ mod tests {
         );
         let new_nullifier =
             Nullifier::for_account_update(&Commitment::new(&account_id, &new_account), &nsk);
-        assert_eq!(index.get(&new_nullifier), Some(&account_id));
-        assert!(!index.contains_key(&old_nullifier));
+        assert_eq!(index.account_for(&new_nullifier), Some(account_id));
+        assert!(index.account_for(&old_nullifier).is_none());
     }
 
     #[test]