Better input ordering (#2175)

* Order inputs on selection instead of sorting later

* update comments

* Use a custom collection type for order

* optimize iters

* use others in iter

* reverse params

* remove address clones

* back it up

* cleanup and don't sort all the time

* comments

* remove IntoIter impl

sdk/examples/wallet/logger.rs

@@ -9,11 +9,7 @@
 //! ```
 
 use iota_sdk::{
-    client::{
-        api::GetAddressesOptions,
-        constants::SHIMMER_COIN_TYPE,
-        secret::{mnemonic::MnemonicSecretManager, SecretManager},
-    },
+    client::{constants::SHIMMER_COIN_TYPE, secret::mnemonic::MnemonicSecretManager},
     crypto::keys::bip44::Bip44,
     wallet::{ClientOptions, Wallet},
 };
@@ -39,9 +35,6 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     // Restore a wallet
     let client_options = ClientOptions::new().with_node(&std::env::var("NODE_URL").unwrap())?;
-    let secret_manager = SecretManager::Mnemonic(MnemonicSecretManager::try_from_mnemonic(
-        std::env::var("MNEMONIC").unwrap(),
-    )?);
 
     let secret_manager = MnemonicSecretManager::try_from_mnemonic(std::env::var("MNEMONIC").unwrap())?;
 

sdk/examples/wallet/offline_signing/1_prepare_transaction.rs

@@ -9,10 +9,7 @@
 //! ```
 
 use iota_sdk::{
-    client::{
-        api::PreparedTransactionDataDto, constants::SHIMMER_COIN_TYPE, secret::SecretManager,
-        stronghold::StrongholdAdapter,
-    },
+    client::{api::PreparedTransactionDataDto, constants::SHIMMER_COIN_TYPE, secret::SecretManager},
     crypto::keys::bip44::Bip44,
     types::block::address::Bech32Address,
     wallet::{ClientOptions, SendParams, Wallet},

sdk/examples/wallet/spammer.rs

@@ -12,14 +12,10 @@ use iota_sdk::{
     client::{
         constants::SHIMMER_COIN_TYPE,
         request_funds_from_faucet,
-        secret::{mnemonic::MnemonicSecretManager, SecretManage, SecretManager},
+        secret::{mnemonic::MnemonicSecretManager, SecretManager},
     },
     crypto::keys::bip44::Bip44,
-    types::block::{
-        address::{Address, Bech32Address, Hrp},
-        output::BasicOutput,
-        payload::signed_transaction::TransactionId,
-    },
+    types::block::{address::Bech32Address, output::BasicOutput, payload::signed_transaction::TransactionId},
     wallet::{ClientOptions, FilterOptions, SendParams, Wallet},
 };
 

sdk/src/client/api/block_builder/transaction_builder/mod.rs

@@ -10,11 +10,11 @@ pub(crate) mod remainder;
 pub(crate) mod requirement;
 pub(crate) mod transition;
 
-use alloc::collections::BTreeMap;
+use alloc::collections::{BTreeMap, VecDeque};
+use core::borrow::Borrow;
 use std::collections::{HashMap, HashSet};
 
 use crypto::keys::bip44::Bip44;
-use packable::PackableExt;
 
 pub use self::{burn::Burn, error::TransactionBuilderError, requirement::Requirement};
 use crate::{
@@ -40,7 +40,7 @@ use crate::{
             signed_transaction::{Transaction, TransactionCapabilities, TransactionCapabilityFlag},
             TaggedDataPayload,
         },
-        protocol::{CommittableAgeRange, ProtocolParameters},
+        protocol::ProtocolParameters,
         slot::{SlotCommitmentId, SlotIndex},
     },
 };
@@ -178,7 +178,7 @@ impl Client {
 pub struct TransactionBuilder {
     available_inputs: Vec<InputSigningData>,
     required_inputs: HashSet<OutputId>,
-    selected_inputs: Vec<InputSigningData>,
+    selected_inputs: OrderedInputs,
     bic_context_inputs: HashSet<BlockIssuanceCreditContextInput>,
     commitment_context_input: Option<CommitmentContextInput>,
     reward_context_inputs: HashSet<OutputId>,
@@ -250,7 +250,7 @@ impl TransactionBuilder {
         Self {
             available_inputs,
             required_inputs: HashSet::new(),
-            selected_inputs: Vec::new(),
+            selected_inputs: Default::default(),
             bic_context_inputs: HashSet::new(),
             commitment_context_input: None,
             reward_context_inputs: HashSet::new(),
@@ -308,7 +308,6 @@ impl TransactionBuilder {
         }
 
         // Gets requirements from outputs.
-        // TODO this may re-evaluate outputs added by inputs
         self.outputs_requirements();
 
         // Gets requirements from burn.
@@ -426,12 +425,6 @@ impl TransactionBuilder {
             }
         }
 
-        let inputs_data = Self::sort_input_signing_data(
-            self.selected_inputs,
-            self.latest_slot_commitment_id.slot_index(),
-            self.protocol_parameters.committable_age_range(),
-        )?;
-
         let mut inputs: Vec<Input> = Vec::new();
         let mut context_inputs = self
             .bic_context_inputs
@@ -440,7 +433,7 @@ impl TransactionBuilder {
             .chain(self.commitment_context_input.map(ContextInput::from))
             .collect::<Vec<_>>();
 
-        for (idx, input) in inputs_data.iter().enumerate() {
+        for (idx, input) in self.selected_inputs.iter().enumerate() {
             inputs.push(Input::Utxo(UtxoInput::from(*input.output_id())));
             if self.reward_context_inputs.contains(input.output_id()) {
                 context_inputs.push(RewardContextInput::new(idx as u16).unwrap().into());
@@ -471,7 +464,7 @@ impl TransactionBuilder {
 
         let data = PreparedTransactionData {
             transaction,
-            inputs_data,
+            inputs_data: self.selected_inputs.into_sorted_iter().collect(),
             remainders: self.remainders.data,
             mana_rewards: self.mana_rewards.into_iter().collect(),
         };
@@ -503,9 +496,18 @@ impl TransactionBuilder {
         // New input may need context inputs
         self.fulfill_context_inputs_requirements(&input);
 
-        self.selected_inputs.push(input);
+        let required_address = input
+            .output
+            .required_address(
+                self.latest_slot_commitment_id.slot_index(),
+                self.protocol_parameters.committable_age_range(),
+            )
+            // PANIC: safe to unwrap as non basic/account/foundry/nft/delegation outputs are already filtered out.
+            .unwrap()
+            .expect("expiration unlockable outputs already filtered out");
+        self.selected_inputs.insert(required_address, input);
 
-        Ok(if added_output { self.added_outputs.last() } else { None })
+        Ok(added_output.then(|| self.added_outputs.last().unwrap()))
     }
 
     /// Sets the required inputs of an [`TransactionBuilder`].
@@ -681,99 +683,113 @@ impl TransactionBuilder {
             }
         })
     }
+}
 
-    // Inputs need to be sorted before signing, because the reference unlock conditions can only reference a lower index
-    pub(crate) fn sort_input_signing_data(
-        mut inputs: Vec<InputSigningData>,
-        commitment_slot_index: SlotIndex,
-        committable_age_range: CommittableAgeRange,
-    ) -> Result<Vec<InputSigningData>, TransactionBuilderError> {
-        // initially sort by output to make it deterministic
-        // TODO: rethink this, we only need it deterministic for tests, for the protocol it doesn't matter, also there
-        // might be a more efficient way to do this
-        inputs.sort_by_key(|i| i.output.pack_to_vec());
-        // filter for ed25519 address first
-        let (mut sorted_inputs, account_nft_address_inputs): (Vec<InputSigningData>, Vec<InputSigningData>) =
-            inputs.into_iter().partition(|input_signing_data| {
-                let required_address = input_signing_data
-                    .output
-                    .required_address(commitment_slot_index, committable_age_range)
-                    // PANIC: safe to unwrap as non basic/account/foundry/nft outputs are already filtered out.
-                    .unwrap()
-                    .expect("expiration unlockable outputs already filtered out");
+#[derive(Clone, Debug, Default)]
+pub(crate) struct OrderedInputs {
+    ed25519: VecDeque<InputSigningData>,
+    other: BTreeMap<Address, VecDeque<InputSigningData>>,
+    len: usize,
+}
 
-                required_address.is_ed25519()
-            });
+impl OrderedInputs {
+    pub(crate) fn sorted_iter(&self) -> OrderedInputsIter<&Address, &InputSigningData> {
+        OrderedInputsIter {
+            queue: self.ed25519.iter().collect(),
+            other: self.other.iter().map(|(k, v)| (k, v.iter().collect())).collect(),
+        }
+    }
 
-        for input in account_nft_address_inputs {
-            let required_address = input
-                .output
-                .required_address(commitment_slot_index, committable_age_range)?
-                .expect("expiration unlockable outputs already filtered out");
+    pub(crate) fn into_sorted_iter(self) -> OrderedInputsIter<Address, InputSigningData> {
+        OrderedInputsIter {
+            queue: self.ed25519,
+            other: self.other,
+        }
+    }
 
-            match sorted_inputs
-                .iter()
-                .position(|input_signing_data| match required_address {
-                    Address::Account(unlock_address) => {
-                        if let Output::Account(account_output) = &input_signing_data.output {
-                            *unlock_address.account_id()
-                                == account_output.account_id_non_null(input_signing_data.output_id())
-                        } else {
-                            false
-                        }
-                    }
-                    Address::Nft(unlock_address) => {
-                        if let Output::Nft(nft_output) = &input_signing_data.output {
-                            *unlock_address.nft_id() == nft_output.nft_id_non_null(input_signing_data.output_id())
-                        } else {
-                            false
-                        }
+    pub(crate) fn iter(&self) -> impl Iterator<Item = &InputSigningData> + Clone {
+        self.ed25519.iter().chain(self.other.values().flatten())
+    }
+
+    pub(crate) fn insert(&mut self, required_address: Address, input: InputSigningData) {
+        if required_address.is_ed25519_backed() {
+            self.ed25519.push_back(input);
+        } else {
+            self.other.entry(required_address).or_default().push_back(input);
+        }
+        self.len += 1;
+    }
+
+    pub(crate) fn len(&self) -> usize {
+        self.len
+    }
+
+    pub(crate) fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}
+
+#[derive(Clone, Debug)]
+pub(crate) struct OrderedInputsIter<A: Borrow<Address> + Ord + core::hash::Hash, I: Borrow<InputSigningData>> {
+    queue: VecDeque<I>,
+    other: BTreeMap<A, VecDeque<I>>,
+}
+
+impl<A: Borrow<Address> + Ord + core::hash::Hash, I: Borrow<InputSigningData>> Iterator for OrderedInputsIter<A, I> {
+    type Item = I;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // Inputs that are unlocked by Ed25519 addresses go in the queue first
+        // because they do not need to reference other inputs for their unlocks.
+        // Each one may have additional dependents which are added to the front of
+        // the queue to be sorted immediately after the input they depend upon.
+        // Those can also have dependents which will go after them.
+        // This creates a tree structure with many to one relationship, which is
+        // flattened by this loop in insertion order.
+        if let Some(input) = self.queue.pop_front() {
+            // Add associated inputs to the front of the queue
+            match &input.borrow().output {
+                Output::Account(account_output) => {
+                    for input in self
+                        .other
+                        .remove(&Address::Account(AccountAddress::new(
+                            account_output.account_id_non_null(input.borrow().output_id()),
+                        )))
+                        .into_iter()
+                        .flatten()
+                        .rev()
+                    {
+                        self.queue.push_front(input);
                     }
-                    _ => false,
-                }) {
-                Some(position) => {
-                    // Insert after the output we need
-                    sorted_inputs.insert(position + 1, input);
                 }
-                None => {
-                    // insert before address
-                    let account_or_nft_address = match &input.output {
-                        Output::Account(account_output) => Some(Address::Account(AccountAddress::new(
-                            account_output.account_id_non_null(input.output_id()),
-                        ))),
-                        Output::Nft(nft_output) => Some(Address::Nft(NftAddress::new(
-                            nft_output.nft_id_non_null(input.output_id()),
-                        ))),
-                        _ => None,
-                    };
-
-                    if let Some(account_or_nft_address) = account_or_nft_address {
-                        // Check for existing outputs for this address, and insert before
-                        match sorted_inputs.iter().position(|input_signing_data| {
-                            let required_address = input_signing_data
-                                .output
-                                .required_address(commitment_slot_index, committable_age_range)
-                                // PANIC: safe to unwrap as non basic/alias/foundry/nft outputs are already filtered
-                                .unwrap()
-                                .expect("expiration unlockable outputs already filtered out");
-
-                            required_address == account_or_nft_address
-                        }) {
-                            Some(position) => {
-                                // Insert before the output with this address required for unlocking
-                                sorted_inputs.insert(position, input);
-                            }
-                            // just push output
-                            None => sorted_inputs.push(input),
-                        }
-                    } else {
-                        // just push basic or foundry output
-                        sorted_inputs.push(input);
+                Output::Nft(nft_output) => {
+                    for input in self
+                        .other
+                        .remove(&Address::Nft(NftAddress::new(
+                            nft_output.nft_id_non_null(input.borrow().output_id()),
+                        )))
+                        .into_iter()
+                        .flatten()
+                        .rev()
+                    {
+                        self.queue.push_front(input);
                     }
                 }
+                _ => (),
+            };
+            return Some(input);
+        }
+        // When the queue is empty, just add anything that is left over to the end of the list.
+        if let Some(mut entry) = self.other.first_entry() {
+            if let Some(input) = entry.get_mut().pop_front() {
+                // Since the structure is a list-of-lists, we need to pop
+                // the inner list if it's empty.
+                if entry.get().is_empty() {
+                    self.other.pop_first();
+                }
+                return Some(input);
             }
         }
-
-        Ok(sorted_inputs)
+        None
     }
 }

sdk/src/client/api/block_builder/transaction_builder/remainder.rs

@@ -30,7 +30,7 @@ impl TransactionBuilder {
         Ok(())
     }
 
-    /// Gets the remainder address from configuration of finds one from the inputs.
+    /// Gets the remainder address from configuration or finds one from the inputs.
     pub(crate) fn get_remainder_address(&self) -> Result<Option<(Address, Option<Bip44>)>, TransactionBuilderError> {
         if let Some(remainder_address) = &self.remainders.address {
             // Search in inputs for the Bip44 chain for the remainder address, so the ledger can regenerate it
@@ -50,7 +50,7 @@ impl TransactionBuilder {
             return Ok(Some((remainder_address.clone(), None)));
         }
 
-        for input in &self.selected_inputs {
+        for input in self.selected_inputs.iter() {
             let required_address = input
                 .output
                 .required_address(

sdk/src/client/api/block_builder/transaction_builder/requirement/amount.rs

@@ -86,7 +86,7 @@ impl TransactionBuilder {
         let mut inputs_sdr = HashMap::new();
         let mut outputs_sdr = HashMap::new();
 
-        for selected_input in &self.selected_inputs {
+        for selected_input in self.selected_inputs.iter() {
             inputs_sum += selected_input.output.amount();
 
             if let Some(sdruc) = sdruc_not_expired(&selected_input.output, self.latest_slot_commitment_id.slot_index())