Skip to content

Commit

Permalink
Use a custom collection type for order
Browse files Browse the repository at this point in the history
  • Loading branch information
Alex Coats committed Mar 13, 2024
1 parent 8f48a49 commit 0472c8f
Show file tree
Hide file tree
Showing 2 changed files with 121 additions and 95 deletions.
214 changes: 120 additions & 94 deletions sdk/src/client/api/block_builder/transaction_builder/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ pub(crate) mod remainder;
pub(crate) mod requirement;
pub(crate) mod transition;

use alloc::collections::BTreeMap;
use alloc::collections::{BTreeMap, VecDeque};
use std::collections::{HashMap, HashSet};

use crypto::keys::bip44::Bip44;
Expand Down Expand Up @@ -180,7 +180,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>,
Expand Down Expand Up @@ -252,7 +252,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(),
Expand Down Expand Up @@ -456,7 +456,7 @@ impl TransactionBuilder {

let data = PreparedTransactionData {
transaction,
inputs_data: self.selected_inputs,
inputs_data: self.selected_inputs.into_iter().collect(),
remainders: self.remainders.data,
mana_rewards: self.mana_rewards.into_iter().collect(),
};
Expand Down Expand Up @@ -485,7 +485,16 @@ impl TransactionBuilder {
self.requirements.push(requirement);
}

self.insert_input_sorted(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(input, required_address);

// New inputs/outputs may need context inputs
if !self.requirements.contains(&Requirement::ContextInputs) {
Expand Down Expand Up @@ -671,102 +680,119 @@ impl TransactionBuilder {
}
})
}
}

fn insert_input_sorted(&mut self, input: InputSigningData) {
let input_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");
if input_required_address.is_ed25519() {
if let Some(position) = self
.selected_inputs
.iter()
.position(|i| i.output_id() > input.output_id())
{
self.selected_inputs.insert(position, input);
} else {
self.selected_inputs.push(input);
}
#[derive(Clone, Debug, Default)]
pub(crate) struct OrderedInputs {
ed25519: VecDeque<InputSigningData>,
other: BTreeMap<Address, Vec<InputSigningData>>,
len: usize,
}

impl OrderedInputs {
pub(crate) fn iter(&self) -> <&Self as IntoIterator>::IntoIter {
self.into_iter()
}

pub(crate) fn insert(&mut self, input: InputSigningData, required_address: Address) {
if required_address.is_ed25519() {
self.ed25519.push_back(input);
} else {
match self.selected_inputs.iter().position(|input_signing_data| {
let required_address = input_signing_data
.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");
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
}
}
_ => false,
self.other.entry(required_address).or_default().push(input);
}
self.len += 1;
}

pub(crate) fn len(&self) -> usize {
self.len
}

pub(crate) fn is_empty(&self) -> bool {
self.len() == 0
}
}

impl<'a> IntoIterator for &'a OrderedInputs {
type Item = &'a InputSigningData;
type IntoIter = alloc::vec::IntoIter<Self::Item>;

fn into_iter(self) -> Self::IntoIter {
let mut other = self
.other
.iter()
.map(|(k, v)| (k, v.iter().collect::<VecDeque<_>>()))
.collect::<BTreeMap<_, _>>();
let mut inputs = Vec::new();
let mut queue = self.ed25519.iter().collect::<VecDeque<_>>();
while let Some(input) = queue.pop_front() {
// Add associated inputs to the front of the queue
match &input.output {
Output::Account(account_output) => {
queue = other
.remove(&Address::Account(AccountAddress::new(
account_output.account_id_non_null(input.output_id()),
)))
.into_iter()
.flatten()
.chain(queue)
.collect()
}
}) {
Some(position) => {
// Insert after the output we need
self.selected_inputs.insert(position + 1, input);
Output::Nft(nft_output) => {
queue = other
.remove(&Address::Nft(NftAddress::new(
nft_output.nft_id_non_null(input.output_id()),
)))
.into_iter()
.flatten()
.chain(queue)
.collect()
}
None => {
// insert before address
let account_or_nft_address = match &input.output {
Output::Account(account_output) => Some(Address::Account(AccountAddress::new(
_ => (),
};
inputs.push(input);
}
inputs.extend(other.into_values().flatten());
inputs.into_iter()
}
}

impl IntoIterator for OrderedInputs {
type Item = InputSigningData;
type IntoIter = alloc::vec::IntoIter<Self::Item>;

fn into_iter(mut self) -> Self::IntoIter {
let mut inputs = Vec::new();
let mut queue = self.ed25519;
while let Some(input) = queue.pop_front() {
// Add associated inputs to the front of the queue
match &input.output {
Output::Account(account_output) => {
queue = self
.other
.remove(&Address::Account(AccountAddress::new(
account_output.account_id_non_null(input.output_id()),
))),
Output::Nft(nft_output) => Some(Address::Nft(NftAddress::new(
)))
.into_iter()
.flatten()
.chain(queue)
.collect()
}
Output::Nft(nft_output) => {
queue = self
.other
.remove(&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 self.selected_inputs.iter().position(|input_signing_data| {
let required_address = input_signing_data
.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");
required_address == account_or_nft_address
}) {
Some(position) => {
// Insert before the output with this address required for unlocking
self.selected_inputs.insert(position, input);
}
// just push output
None => self.selected_inputs.push(input),
}
} else {
// just push basic or foundry output
self.selected_inputs.push(input);
}
)))
.into_iter()
.flatten()
.chain(queue)
.collect()
}
}
_ => (),
};
inputs.push(input);
}
inputs.extend(self.other.into_values().flatten());
inputs.into_iter()
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit 0472c8f

Please sign in to comment.