From d2850654a2fda351557d8ebe72082bca782b972c Mon Sep 17 00:00:00 2001
From: Aleksandr Karbyshev <karbyshev@users.noreply.github.com>
Date: Thu, 7 Mar 2024 18:25:22 +0100
Subject: [PATCH] Atomic version of `insert_address`

---
 crates/apps/src/lib/cli/wallet.rs           |  3 +-
 crates/apps/src/lib/config/genesis/chain.rs | 24 +++++-----
 crates/sdk/src/tx.rs                        | 14 +++---
 crates/sdk/src/wallet/mod.rs                | 49 ++++++++++++++-------
 4 files changed, 59 insertions(+), 31 deletions(-)

diff --git a/crates/apps/src/lib/cli/wallet.rs b/crates/apps/src/lib/cli/wallet.rs
index 201e64dfc0..6f6fc053c6 100644
--- a/crates/apps/src/lib/cli/wallet.rs
+++ b/crates/apps/src/lib/cli/wallet.rs
@@ -1463,7 +1463,8 @@ fn transparent_address_add(
     let alias = alias.to_lowercase();
     let mut wallet = load_wallet(ctx);
     if wallet
-        .insert_address(&alias, address, alias_force)
+        .insert_address_atomic(&alias, address, alias_force)
+        .expect("Failed to update the wallet storage.")
         .is_none()
     {
         edisplay_line!(io, "Address not added");
diff --git a/crates/apps/src/lib/config/genesis/chain.rs b/crates/apps/src/lib/config/genesis/chain.rs
index 7004afc09c..7132289d02 100644
--- a/crates/apps/src/lib/config/genesis/chain.rs
+++ b/crates/apps/src/lib/config/genesis/chain.rs
@@ -129,11 +129,13 @@ impl Finalized {
     ) -> Wallet<CliWalletUtils> {
         let mut wallet = crate::wallet::load_or_new(base_dir);
         for (alias, config) in &self.tokens.token {
-            wallet.insert_address(
-                alias.normalize(),
-                config.address.clone(),
-                false,
-            );
+            wallet
+                .insert_address_atomic(
+                    alias.normalize(),
+                    config.address.clone(),
+                    false,
+                )
+                .expect("Failed to update the wallet storage.");
             wallet
                 .add_vp_type_to_address_atomic(
                     AddressVpType::Token,
@@ -170,11 +172,13 @@ impl Finalized {
             InternalAddress::Governance,
             InternalAddress::Pgf,
         ] {
-            wallet.insert_address(
-                int_add.to_string().to_lowercase(),
-                Address::Internal(int_add.clone()),
-                false,
-            );
+            wallet
+                .insert_address_atomic(
+                    int_add.to_string().to_lowercase(),
+                    Address::Internal(int_add.clone()),
+                    false,
+                )
+                .expect("Failed to update the wallet storage.");
         }
 
         wallet
diff --git a/crates/sdk/src/tx.rs b/crates/sdk/src/tx.rs
index fce97301ba..b569525765 100644
--- a/crates/sdk/src/tx.rs
+++ b/crates/sdk/src/tx.rs
@@ -536,11 +536,15 @@ pub async fn save_initialized_accounts<N: Namada>(
                 None => N::WalletUtils::read_alias(&encoded).into(),
             };
             let alias = alias.into_owned();
-            let added = context.wallet_mut().await.insert_address(
-                alias.clone(),
-                address.clone(),
-                args.wallet_alias_force,
-            );
+            let added = context
+                .wallet_mut()
+                .await
+                .insert_address_atomic(
+                    alias.clone(),
+                    address.clone(),
+                    args.wallet_alias_force,
+                )
+                .expect("Failed to update the wallet storage.");
             match added {
                 Some(new_alias) if new_alias != encoded => {
                     display_line!(
diff --git a/crates/sdk/src/wallet/mod.rs b/crates/sdk/src/wallet/mod.rs
index 45294e5893..eaf8b3b5a5 100644
--- a/crates/sdk/src/wallet/mod.rs
+++ b/crates/sdk/src/wallet/mod.rs
@@ -1224,21 +1224,21 @@ impl<U: WalletIo> Wallet<U> {
         }
     }
 
-    /// Add a new address with the given alias. If the alias is already used,
-    /// will ask whether the existing alias should be replaced, a different
-    /// alias is desired, or the alias creation should be cancelled. Return
-    /// the chosen alias if the address has been added, otherwise return
-    /// nothing.
-    pub fn insert_address(
-        &mut self,
-        alias: impl AsRef<str>,
-        address: Address,
-        force_alias: bool,
-    ) -> Option<String> {
-        self.store
-            .insert_address::<U>(alias.into(), address, force_alias)
-            .map(Into::into)
-    }
+    // /// Add a new address with the given alias. If the alias is already used,
+    // /// will ask whether the existing alias should be replaced, a different
+    // /// alias is desired, or the alias creation should be cancelled. Return
+    // /// the chosen alias if the address has been added, otherwise return
+    // /// nothing.
+    // pub fn insert_address(
+    //     &mut self,
+    //     alias: impl AsRef<str>,
+    //     address: Address,
+    //     force_alias: bool,
+    // ) -> Option<String> {
+    //     self.store
+    //         .insert_address::<U>(alias.into(), address, force_alias)
+    //         .map(Into::into)
+    // }
 
     /// Add a new keypair with the given alias. If the alias is already used,
     /// will ask whether the existing alias should be replaced, a different
@@ -1468,4 +1468,23 @@ impl<U: WalletIo + WalletStorage> Wallet<U> {
             },
         )
     }
+
+    /// Add a new address with the given alias. If the alias is already used,
+    /// will ask whether the existing alias should be replaced, a different
+    /// alias is desired, or the alias creation should be cancelled. Return
+    /// the chosen alias if the address has been added, otherwise return
+    /// nothing.
+    pub fn insert_address_atomic(
+        &mut self,
+        alias: impl AsRef<str>,
+        address: Address,
+        force_alias: bool,
+    ) -> Result<Option<String>, LoadStoreError> {
+        let mut addr_alias: Option<Alias> = Option::default();
+        self.utils.update_store(|store| {
+            addr_alias =
+                store.insert_address::<U>(alias.into(), address, force_alias);
+        })?;
+        Ok(addr_alias.map(Into::into))
+    }
 }