From 9edcdd9a4e4170ce48ca656e0e8f36dcbe42ff07 Mon Sep 17 00:00:00 2001
From: edouard <m@edouard.paris>
Date: Tue, 29 Aug 2023 14:55:49 +0200
Subject: [PATCH] Add labels to change outputs according to main label

---
 gui/src/app/state/spend/step.rs | 38 +++++++++++++++++++++++++++++----
 gui/src/daemon/model.rs         |  2 +-
 2 files changed, 35 insertions(+), 5 deletions(-)

diff --git a/gui/src/app/state/spend/step.rs b/gui/src/app/state/spend/step.rs
index 63d98ba4..b43f8092 100644
--- a/gui/src/app/state/spend/step.rs
+++ b/gui/src/app/state/spend/step.rs
@@ -30,6 +30,7 @@ const DUST_OUTPUT_SATS: u64 = 5_000;
 pub struct TransactionDraft {
     network: Network,
     inputs: Vec<Coin>,
+    recipients: Vec<Recipient>,
     generated: Option<Psbt>,
     batch_label: Option<String>,
     labels: HashMap<String, String>,
@@ -40,6 +41,7 @@ impl TransactionDraft {
         Self {
             network,
             inputs: Vec::new(),
+            recipients: Vec::new(),
             generated: None,
             batch_label: None,
             labels: HashMap::new(),
@@ -373,6 +375,7 @@ impl Step for DefineSpend {
                 }
             }
         }
+        draft.recipients = self.recipients.clone();
         if self.recipients.len() > 1 {
             draft.batch_label = Some(self.batch_label.value.clone());
         }
@@ -407,7 +410,7 @@ impl Step for DefineSpend {
     }
 }
 
-#[derive(Default)]
+#[derive(Default, Clone)]
 struct Recipient {
     label: form::Value<String>,
     address: form::Value<String>,
@@ -520,9 +523,36 @@ impl Step for SaveSpend {
             draft.network,
         );
         tx.labels = draft.labels.clone();
-        if let Some(label) = &draft.batch_label {
-            tx.labels
-                .insert(tx.psbt.unsigned_tx.txid().to_string(), label.clone());
+
+        if tx.is_batch() {
+            if let Some(label) = &draft.batch_label {
+                tx.labels
+                    .insert(tx.psbt.unsigned_tx.txid().to_string(), label.clone());
+                for (i, output) in tx.psbt.unsigned_tx.output.iter().enumerate() {
+                    let address_str = Address::from_script(&output.script_pubkey, tx.network)
+                        .unwrap()
+                        .to_string();
+                    if tx.change_indexes.contains(&i) && tx.labels.contains_key(&address_str) {
+                        tx.labels
+                            .insert(address_str, format!("Change of {}", label.clone()));
+                    }
+                }
+            }
+        } else if let Some(recipient) = draft.recipients.first() {
+            if !recipient.label.value.is_empty() {
+                let label = recipient.label.value.clone();
+                tx.labels
+                    .insert(tx.psbt.unsigned_tx.txid().to_string(), label.clone());
+                for (i, output) in tx.psbt.unsigned_tx.output.iter().enumerate() {
+                    let address_str = Address::from_script(&output.script_pubkey, tx.network)
+                        .unwrap()
+                        .to_string();
+                    if tx.change_indexes.contains(&i) && tx.labels.contains_key(&address_str) {
+                        tx.labels
+                            .insert(address_str, format!("Change of {}", label.clone()));
+                    }
+                }
+            }
         }
 
         self.spend = Some(psbt::PsbtState::new(self.wallet.clone(), tx, false));
diff --git a/gui/src/daemon/model.rs b/gui/src/daemon/model.rs
index a20a58c2..ae5c36e7 100644
--- a/gui/src/daemon/model.rs
+++ b/gui/src/daemon/model.rs
@@ -28,7 +28,7 @@ pub fn remaining_sequence(coin: &Coin, blockheight: u32, timelock: u16) -> u32 {
 
 #[derive(Debug, Clone)]
 pub struct SpendTx {
-    network: Network,
+    pub network: Network,
     pub coins: Vec<Coin>,
     pub labels: HashMap<String, String>,
     pub psbt: Psbt,