feat: support payload distribution

wallet/core/src/account/mod.rs

@@ -22,6 +22,7 @@ use crate::imports::*;
 use crate::storage::account::AccountSettings;
 use crate::storage::AccountMetadata;
 use crate::storage::{PrvKeyData, PrvKeyDataId};
+use crate::tx::generator::settings::PayloadDistribution;
 use crate::tx::PaymentOutput;
 use crate::tx::{Fees, Generator, GeneratorSettings, GeneratorSummary, PaymentDestination, PendingTransaction, Signer};
 use crate::utxo::balance::{AtomicBalance, BalanceStrings};
@@ -310,8 +311,13 @@ pub trait Account: AnySync + Send + Sync + 'static {
     ) -> Result<(GeneratorSummary, Vec<kaspa_hashes::Hash>)> {
         let keydata = self.prv_key_data(wallet_secret).await?;
         let signer = Arc::new(Signer::new(self.clone().as_dyn_arc(), keydata, payment_secret));
-        let settings =
-            GeneratorSettings::try_new_with_account(self.clone().as_dyn_arc(), PaymentDestination::Change, Fees::None, None)?;
+        let settings = GeneratorSettings::try_new_with_account(
+            self.clone().as_dyn_arc(),
+            PaymentDestination::Change,
+            Fees::None,
+            None,
+            PayloadDistribution::default(),
+        )?;
         let generator = Generator::try_new(settings, Some(signer), Some(abortable))?;
 
         let mut stream = generator.stream();
@@ -344,7 +350,13 @@ pub trait Account: AnySync + Send + Sync + 'static {
         let keydata = self.prv_key_data(wallet_secret).await?;
         let signer = Arc::new(Signer::new(self.clone().as_dyn_arc(), keydata, payment_secret));
 
-        let settings = GeneratorSettings::try_new_with_account(self.clone().as_dyn_arc(), destination, priority_fee_sompi, payload)?;
+        let settings = GeneratorSettings::try_new_with_account(
+            self.clone().as_dyn_arc(),
+            destination,
+            priority_fee_sompi,
+            payload,
+            PayloadDistribution::default(),
+        )?;
 
         let generator = Generator::try_new(settings, Some(signer), Some(abortable))?;
 
@@ -372,7 +384,13 @@ pub trait Account: AnySync + Send + Sync + 'static {
         payment_secret: Option<Secret>,
         abortable: &Abortable,
     ) -> Result<Bundle, Error> {
-        let settings = GeneratorSettings::try_new_with_account(self.clone().as_dyn_arc(), destination, priority_fee_sompi, payload)?;
+        let settings = GeneratorSettings::try_new_with_account(
+            self.clone().as_dyn_arc(),
+            destination,
+            priority_fee_sompi,
+            payload,
+            PayloadDistribution::default(),
+        )?;
         let keydata = self.prv_key_data(wallet_secret).await?;
         let signer = Arc::new(PSKBSigner::new(self.clone().as_dyn_arc(), keydata, payment_secret));
         let generator = Generator::try_new(settings, None, Some(abortable))?;
@@ -453,6 +471,7 @@ pub trait Account: AnySync + Send + Sync + 'static {
             final_transaction_destination,
             priority_fee_sompi,
             final_transaction_payload,
+            PayloadDistribution::default(),
         )?
         .utxo_context_transfer(destination_account.utxo_context());
 
@@ -480,7 +499,13 @@ pub trait Account: AnySync + Send + Sync + 'static {
         payload: Option<Vec<u8>>,
         abortable: &Abortable,
     ) -> Result<GeneratorSummary> {
-        let settings = GeneratorSettings::try_new_with_account(self.as_dyn_arc(), destination, priority_fee_sompi, payload)?;
+        let settings = GeneratorSettings::try_new_with_account(
+            self.as_dyn_arc(),
+            destination,
+            priority_fee_sompi,
+            payload,
+            PayloadDistribution::default(),
+        )?;
 
         let generator = Generator::try_new(settings, None, Some(abortable))?;
 
@@ -607,6 +632,7 @@ pub trait DerivationCapableAccount: Account {
                         PaymentDestination::Change,
                         Fees::None,
                         None,
+                        PayloadDistribution::default(),
                         None,
                     )?;
 

wallet/core/src/account/pskb.rs

@@ -341,6 +341,7 @@ pub fn pskt_to_pending_transaction(
         final_transaction_priority_fee: fee_u.into(),
         final_transaction_destination,
         final_transaction_payload: None,
+        payload_distribution: crate::tx::generator::settings::PayloadDistribution::default(),
     };
 
     // Create the Generator

wallet/core/src/tx/generator/generator.rs

@@ -59,6 +59,7 @@
 
 use crate::imports::*;
 use crate::result::Result;
+use crate::tx::generator::settings::PayloadDistribution;
 use crate::tx::{
     mass::*, Fees, GeneratorSettings, GeneratorSummary, PaymentDestination, PendingTransaction, PendingTransactionIterator,
     PendingTransactionStream,
@@ -301,6 +302,8 @@ struct Inner {
     final_transaction_payload_mass: u64,
     // execution context
     context: Mutex<Context>,
+    // payload distribution
+    payload_distribution: PayloadDistribution,
 }
 
 impl std::fmt::Debug for Inner {
@@ -324,6 +327,7 @@ impl std::fmt::Debug for Inner {
             .field("final_transaction_payload", &self.final_transaction_payload)
             .field("final_transaction_payload_mass", &self.final_transaction_payload_mass)
             // .field("context", &self.context)
+            .field("payload_distribution", &self.payload_distribution)
             .finish()
     }
 }
@@ -352,6 +356,7 @@ impl Generator {
             final_transaction_destination,
             final_transaction_payload,
             destination_utxo_context,
+            payload_distribution,
         } = settings;
 
         let network_type = NetworkType::from(network_id);
@@ -460,6 +465,7 @@ impl Generator {
             final_transaction_payload,
             final_transaction_payload_mass,
             destination_utxo_context,
+            payload_distribution,
         };
 
         Ok(Self { inner: Arc::new(inner) })
@@ -567,7 +573,12 @@ impl Generator {
 
     /// Calculate relay transaction mass for the current transaction `data`
     fn calc_relay_transaction_mass(&self, data: &Data) -> u64 {
-        data.aggregate_mass + self.inner.standard_change_output_compute_mass
+        let payload_mass = match self.inner.payload_distribution {
+            PayloadDistribution::AllTxs => self.inner.final_transaction_payload_mass,
+            PayloadDistribution::FinalOnly => 0,
+        };
+
+        data.aggregate_mass + self.inner.standard_change_output_compute_mass + payload_mass
     }
 
     /// Calculate relay transaction fees for the current transaction `data`
@@ -1063,14 +1074,24 @@ impl Generator {
                 let output_value = aggregate_input_value - transaction_fees;
                 let script_public_key = pay_to_address_script(&self.inner.change_address);
                 let output = TransactionOutput::new(output_value, script_public_key.clone());
-                let tx = Transaction::new(0, inputs, vec![output], 0, SUBNETWORK_ID_NATIVE, 0, vec![]);
+
+                // Get payload based on payload distribution
+                let payload = match self.inner.payload_distribution {
+                    PayloadDistribution::AllTxs => self.inner.final_transaction_payload.clone(),
+                    PayloadDistribution::FinalOnly => vec![],
+                };
+
+                let tx = Transaction::new(0, inputs, vec![output], 0, SUBNETWORK_ID_NATIVE, 0, payload);
 
                 let mut transaction_mass = self.inner.mass_calculator.calc_overall_mass_for_unsigned_consensus_transaction(
                     &tx,
                     &utxo_entry_references,
                     self.inner.minimum_signatures,
                 )?;
+
+                // Add additional mass for compound transactions
                 transaction_mass = transaction_mass.saturating_add(self.inner.network_params.additional_compound_transaction_mass());
+
                 if transaction_mass > MAXIMUM_STANDARD_TRANSACTION_MASS {
                     // this should never occur as we should not produce transactions higher than the mass limit
                     return Err(Error::MassCalculationError);

wallet/core/src/tx/generator/settings.rs

@@ -11,6 +11,22 @@ use crate::utxo::{UtxoContext, UtxoEntryReference, UtxoIterator};
 use kaspa_addresses::Address;
 use workflow_core::channel::Multiplexer;
 
+/// Defines how payload data is distributed across transactions in a multi-transaction sequence.
+/// This is particularly relevant when a large transaction is split into multiple smaller ones.
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum PayloadDistribution {
+    /// Payload is included only in the final transaction.
+    FinalOnly,
+    /// Payload is included in all transactions.
+    AllTxs,
+}
+
+impl Default for PayloadDistribution {
+    fn default() -> Self {
+        Self::FinalOnly
+    }
+}
+
 pub struct GeneratorSettings {
     // Network type
     pub network_id: NetworkId,
@@ -36,6 +52,7 @@ pub struct GeneratorSettings {
     pub final_transaction_payload: Option<Vec<u8>>,
     // transaction is a transfer between accounts
     pub destination_utxo_context: Option<UtxoContext>,
+    pub payload_distribution: PayloadDistribution,
 }
 
 // impl std::fmt::Debug for GeneratorSettings {
@@ -62,6 +79,7 @@ impl GeneratorSettings {
         final_transaction_destination: PaymentDestination,
         final_priority_fee: Fees,
         final_transaction_payload: Option<Vec<u8>>,
+        payload_distribution: PayloadDistribution,
     ) -> Result<Self> {
         let network_id = account.utxo_context().processor().network_id()?;
         let change_address = account.change_address()?;
@@ -85,6 +103,7 @@ impl GeneratorSettings {
             final_transaction_destination,
             final_transaction_payload,
             destination_utxo_context: None,
+            payload_distribution,
         };
 
         Ok(settings)
@@ -99,6 +118,7 @@ impl GeneratorSettings {
         final_transaction_destination: PaymentDestination,
         final_priority_fee: Fees,
         final_transaction_payload: Option<Vec<u8>>,
+        payload_distribution: PayloadDistribution,
         multiplexer: Option<Multiplexer<Box<Events>>>,
     ) -> Result<Self> {
         let network_id = utxo_context.processor().network_id()?;
@@ -118,6 +138,7 @@ impl GeneratorSettings {
             final_transaction_destination,
             final_transaction_payload,
             destination_utxo_context: None,
+            payload_distribution,
         };
 
         Ok(settings)
@@ -133,6 +154,7 @@ impl GeneratorSettings {
         final_transaction_destination: PaymentDestination,
         final_priority_fee: Fees,
         final_transaction_payload: Option<Vec<u8>>,
+        payload_distribution: PayloadDistribution,
         multiplexer: Option<Multiplexer<Box<Events>>>,
     ) -> Result<Self> {
         let settings = GeneratorSettings {
@@ -149,6 +171,7 @@ impl GeneratorSettings {
             final_transaction_destination,
             final_transaction_payload,
             destination_utxo_context: None,
+            payload_distribution,
         };
 
         Ok(settings)

wallet/core/src/tx/generator/test.rs

@@ -2,16 +2,20 @@
 
 use crate::error::Error;
 use crate::result::Result;
+use crate::tx::generator::settings::PayloadDistribution;
 use crate::tx::{Fees, MassCalculator, PaymentDestination};
 use crate::utxo::UtxoEntryReference;
 use crate::{tx::PaymentOutputs, utils::kaspa_to_sompi};
 use kaspa_addresses::Address;
+use kaspa_consensus_client::UtxoEntry;
 use kaspa_consensus_core::network::{NetworkId, NetworkType};
-use kaspa_consensus_core::tx::Transaction;
+use kaspa_consensus_core::tx::{Transaction, TransactionId, TransactionOutpoint};
+use kaspa_txscript::pay_to_address_script;
 use rand::prelude::*;
 use std::cell::RefCell;
 use std::fmt::Debug;
 use std::rc::Rc;
+use std::sync::Arc;
 use workflow_log::style;
 
 use super::*;
@@ -430,6 +434,7 @@ where
         final_transaction_priority_fee: final_priority_fee,
         final_transaction_destination,
         final_transaction_payload,
+        payload_distribution: PayloadDistribution::default(),
     };
 
     Generator::try_new(settings, None, None)
@@ -725,3 +730,84 @@ fn test_generator_inputs_250k_outputs_2_sweep() -> Result<()> {
     generator.unwrap().harness().accumulate(2875).finalize();
     Ok(())
 }
+
+#[test]
+fn test_generator_payload_distribution() -> Result<()> {
+    let network_id = NetworkId::new(NetworkType::Mainnet);
+    let change_address = output_address(NetworkType::Mainnet);
+    let destination = output_address(NetworkType::Mainnet);
+
+    // Create some test UTXO entries that will require multiple transactions
+    let utxo_entries: Vec<UtxoEntryReference> = (0..10)
+        .map(|i| {
+            let outpoint = TransactionOutpoint::new(TransactionId::from_bytes([i as u8; 32]), 0);
+            let utxo = UtxoEntry {
+                address: Some(change_address.clone()),
+                outpoint: outpoint.into(),
+                amount: 1_000_000_000, // 1 KAS each
+                script_public_key: pay_to_address_script(&change_address),
+                block_daa_score: 0,
+                is_coinbase: false,
+            };
+            UtxoEntryReference { utxo: Arc::new(utxo) }
+        })
+        .collect();
+
+    let test_payload = vec![1, 2, 3, 4, 5]; // Test payload
+
+    // Test case 1: FinalOnly behavior (default)
+    {
+        let settings = GeneratorSettings::try_new_with_iterator(
+            network_id,
+            Box::new(utxo_entries.clone().into_iter()),
+            None,
+            change_address.clone(),
+            1,
+            1,
+            PaymentDestination::PaymentOutputs(PaymentOutputs::from((destination.clone(), 5_000_000_000))),
+            Fees::SenderPays(1000),
+            Some(test_payload.clone()),
+            PayloadDistribution::FinalOnly,
+            None,
+        )?;
+
+        let generator = Generator::try_new(settings, None, None)?;
+        let transactions: Vec<_> = generator.iter().collect::<Result<Vec<_>>>()?;
+
+        // Verify only final transaction has payload
+        for (i, tx) in transactions.iter().enumerate() {
+            if i == transactions.len() - 1 {
+                assert_eq!(tx.transaction().payload, test_payload, "Final transaction should have payload");
+            } else {
+                assert!(tx.transaction().payload.is_empty(), "Intermediate transaction should not have payload");
+            }
+        }
+    }
+
+    // Test case 2: AllTxs behavior
+    {
+        let settings = GeneratorSettings::try_new_with_iterator(
+            network_id,
+            Box::new(utxo_entries.into_iter()),
+            None,
+            change_address,
+            1,
+            1,
+            PaymentDestination::PaymentOutputs(PaymentOutputs::from((destination, 5_000_000_000))),
+            Fees::SenderPays(1000),
+            Some(test_payload.clone()),
+            PayloadDistribution::AllTxs,
+            None,
+        )?;
+
+        let generator = Generator::try_new(settings, None, None)?;
+        let transactions: Vec<_> = generator.iter().collect::<Result<Vec<_>>>()?;
+
+        // Verify all transactions have payload
+        for tx in transactions.iter() {
+            assert_eq!(tx.transaction().payload, test_payload, "Every transaction should have payload");
+        }
+    }
+
+    Ok(())
+}