Better transitions

sdk/src/client/api/block_builder/options.rs

@@ -6,7 +6,7 @@ use alloc::collections::{BTreeMap, BTreeSet};
 use serde::{Deserialize, Serialize};
 
 use crate::{
-    client::api::transaction_builder::Burn,
+    client::api::transaction_builder::{transition::Transitions, Burn},
     types::block::{
         address::Address,
         output::{AccountId, OutputId},
@@ -25,6 +25,8 @@ pub struct TransactionOptions {
     pub tagged_data_payload: Option<TaggedDataPayload>,
     /// Inputs that must be used for the transaction.
     pub required_inputs: BTreeSet<OutputId>,
+    /// Specifies what needs to be transitioned in the transaction and how.
+    pub transitions: Option<Transitions>,
     /// Specifies what needs to be burned in the transaction.
     pub burn: Option<Burn>,
     /// A string attached to the transaction.
@@ -45,6 +47,7 @@ impl Default for TransactionOptions {
             remainder_value_strategy: Default::default(),
             tagged_data_payload: Default::default(),
             required_inputs: Default::default(),
+            transitions: Default::default(),
             burn: Default::default(),
             note: Default::default(),
             allow_micro_amount: false,

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

@@ -11,10 +11,11 @@ use super::Requirement;
 use crate::types::block::{
     context_input::ContextInputError,
     mana::ManaError,
-    output::{ChainId, NativeTokenError, OutputError, OutputId, TokenId},
+    output::{feature::FeatureError, AccountId, ChainId, NativeTokenError, OutputError, OutputId, TokenId},
     payload::PayloadError,
     semantic::TransactionFailureReason,
     signature::SignatureError,
+    slot::EpochIndex,
     unlock::UnlockError,
     BlockError,
 };
@@ -25,9 +26,16 @@ use crate::types::block::{
 pub enum TransactionBuilderError {
     #[error("additional inputs required for {0:?}, but additional input selection is disabled")]
     AdditionalInputsRequired(Requirement),
+    #[error("account {0} is already staking")]
+    AlreadyStaking(AccountId),
     /// Can't burn and transition an output at the same time.
     #[error("can't burn and transition an output at the same time, chain ID: {0}")]
     BurnAndTransition(ChainId),
+    #[error("account id {account_id} cannot end staking until {end_epoch}")]
+    CannotEndStaking {
+        account_id: AccountId,
+        end_epoch: EpochIndex,
+    },
     #[error("mana rewards provided without an associated burn or custom input, output ID: {0}")]
     ExtraManaRewards(OutputId),
     /// Insufficient amount provided.
@@ -68,9 +76,15 @@ pub enum TransactionBuilderError {
     /// No available inputs were provided to transaction builder.
     #[error("no available inputs provided")]
     NoAvailableInputsProvided,
+    #[error("account {0} is not staking")]
+    NotStaking(AccountId),
     /// Required input is not available.
     #[error("required input {0} is not available")]
     RequiredInputIsNotAvailable(OutputId),
+    #[error("new staking period {additional_epochs} is less than the minimum {min}")]
+    StakingPeriodLessThanMin { additional_epochs: u32, min: u32 },
+    #[error("cannot transition non-implicit-account output {0}")]
+    TransitionNonImplicitAccount(OutputId),
     /// Unfulfillable requirement.
     #[error("unfulfillable requirement {0:?}")]
     UnfulfillableRequirement(Requirement),
@@ -102,6 +116,9 @@ pub enum TransactionBuilderError {
     /// Unlock errors.
     #[error("{0}")]
     Unlock(#[from] UnlockError),
+    /// Feature errors.
+    #[error("{0}")]
+    Feature(#[from] FeatureError),
     /// Semantic errors.
     #[error("{0}")]
     Semantic(#[from] TransactionFailureReason),

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

@@ -15,7 +15,7 @@ use std::collections::{HashMap, HashSet};
 use crypto::keys::bip44::Bip44;
 use packable::PackableExt;
 
-pub use self::{burn::Burn, error::TransactionBuilderError, requirement::Requirement};
+pub use self::{burn::Burn, error::TransactionBuilderError, requirement::Requirement, transition::Transitions};
 use crate::{
     client::{
         api::{
@@ -32,8 +32,8 @@ use crate::{
         input::{Input, UtxoInput, INPUT_COUNT_RANGE},
         mana::ManaAllotment,
         output::{
-            AccountId, AccountOutputBuilder, AnchorOutputBuilder, BasicOutputBuilder, NftOutputBuilder, Output,
-            OutputId, OUTPUT_COUNT_RANGE,
+            AccountId, AccountOutputBuilder, BasicOutputBuilder, ChainId, NftOutputBuilder, Output, OutputId,
+            OUTPUT_COUNT_RANGE,
         },
         payload::{
             signed_transaction::{Transaction, TransactionCapabilities, TransactionCapabilityFlag},
@@ -188,6 +188,7 @@ pub struct TransactionBuilder {
     provided_outputs: Vec<Output>,
     added_outputs: Vec<Output>,
     addresses: HashSet<Address>,
+    transitions: Option<Transitions>,
     burn: Option<Burn>,
     remainders: Remainders,
     creation_slot: SlotIndex,
@@ -215,7 +216,7 @@ pub(crate) struct Remainders {
     address: Option<Address>,
     data: Vec<RemainderData>,
     storage_deposit_returns: Vec<Output>,
-    added_mana: u64,
+    added_mana: HashMap<Option<ChainId>, u64>,
 }
 
 impl TransactionBuilder {
@@ -260,6 +261,7 @@ impl TransactionBuilder {
             provided_outputs: outputs.into_iter().collect(),
             added_outputs: Vec::new(),
             addresses,
+            transitions: None,
             burn: None,
             remainders: Default::default(),
             creation_slot: creation_slot_index.into(),
@@ -377,19 +379,28 @@ impl TransactionBuilder {
             return Err(TransactionBuilderError::InvalidInputCount(self.selected_inputs.len()));
         }
 
-        if self.remainders.added_mana > 0 {
-            let remainder_address = self
-                .get_remainder_address()?
-                .ok_or(TransactionBuilderError::MissingInputWithEd25519Address)?
-                .0;
-            let added_mana = self.remainders.added_mana;
-            if let Some(output) = self.get_output_for_added_mana(&remainder_address) {
+        let remainder_address = self
+            .get_remainder_address()?
+            .ok_or(TransactionBuilderError::MissingInputWithEd25519Address)?
+            .0;
+        for (chain_id, added_mana) in core::mem::take(&mut self.remainders.added_mana) {
+            if let Some(output) = self.get_output_for_added_mana(chain_id, &remainder_address) {
+                log::debug!(
+                    "Adding {added_mana} excess input mana to output with address {remainder_address} and {chain_id:?}"
+                );
+                let new_mana = output.mana() + added_mana;
+                *output = match output {
+                    Output::Basic(b) => BasicOutputBuilder::from(&*b).with_mana(new_mana).finish_output()?,
+                    Output::Account(a) => AccountOutputBuilder::from(&*a).with_mana(new_mana).finish_output()?,
+                    Output::Nft(n) => NftOutputBuilder::from(&*n).with_mana(new_mana).finish_output()?,
+                    _ => unreachable!(),
+                };
+            } else if let Some(output) = self.get_output_for_added_mana(None, &remainder_address) {
                 log::debug!("Adding {added_mana} excess input mana to output with address {remainder_address}");
                 let new_mana = output.mana() + added_mana;
                 *output = match output {
                     Output::Basic(b) => BasicOutputBuilder::from(&*b).with_mana(new_mana).finish_output()?,
                     Output::Account(a) => AccountOutputBuilder::from(&*a).with_mana(new_mana).finish_output()?,
-                    Output::Anchor(a) => AnchorOutputBuilder::from(&*a).with_mana(new_mana).finish_output()?,
                     Output::Nft(n) => NftOutputBuilder::from(&*n).with_mana(new_mana).finish_output()?,
                     _ => unreachable!(),
                 };
@@ -512,25 +523,31 @@ impl TransactionBuilder {
         Ok(added_output)
     }
 
-    /// Sets the required inputs of an [`TransactionBuilder`].
+    /// Sets the required inputs of a [`TransactionBuilder`].
     pub fn with_required_inputs(mut self, inputs: impl IntoIterator<Item = OutputId>) -> Self {
         self.required_inputs = inputs.into_iter().collect();
         self
     }
 
-    /// Sets the burn of an [`TransactionBuilder`].
+    /// Sets the transitions of a [`TransactionBuilder`].
+    pub fn with_transitions(mut self, transitions: impl Into<Option<Transitions>>) -> Self {
+        self.transitions = transitions.into();
+        self
+    }
+
+    /// Sets the burn of a [`TransactionBuilder`].
     pub fn with_burn(mut self, burn: impl Into<Option<Burn>>) -> Self {
         self.burn = burn.into();
         self
     }
 
-    /// Sets the remainder address of an [`TransactionBuilder`].
+    /// Sets the remainder address of a [`TransactionBuilder`].
     pub fn with_remainder_address(mut self, address: impl Into<Option<Address>>) -> Self {
         self.remainders.address = address.into();
         self
     }
 
-    /// Sets the mana allotments of an [`TransactionBuilder`].
+    /// Sets the mana allotments of a [`TransactionBuilder`].
     pub fn with_mana_allotments(mut self, mana_allotments: impl IntoIterator<Item = (AccountId, u64)>) -> Self {
         self.mana_allotments = mana_allotments.into_iter().collect();
         self

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

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use alloc::collections::BTreeMap;
-use std::collections::HashMap;
 
 use crypto::keys::bip44::Bip44;
 use primitive_types::U256;
@@ -13,8 +12,8 @@ use crate::{
     types::block::{
         address::{Address, Ed25519Address},
         output::{
-            unlock_condition::AddressUnlockCondition, AccountOutput, BasicOutput, BasicOutputBuilder, NativeToken,
-            NftOutput, Output, StorageScoreParameters, TokenId,
+            unlock_condition::AddressUnlockCondition, BasicOutputBuilder, ChainId, NativeToken, Output,
+            StorageScoreParameters, TokenId,
         },
     },
 };
@@ -110,9 +109,26 @@ impl TransactionBuilder {
             .ok_or(TransactionBuilderError::MissingInputWithEd25519Address)?;
 
         // If there is a mana remainder, try to fit it in an existing output
-        if mana_diff > 0 && self.output_for_added_mana_exists(&remainder_address) {
-            log::debug!("Allocating {mana_diff} excess input mana for output with address {remainder_address}");
-            self.remainders.added_mana = std::mem::take(&mut mana_diff);
+        if mana_diff > 0 {
+            for (chain_id, (input_mana, output_mana)) in self.mana_chains()? {
+                if input_mana > output_mana
+                    && (self.output_for_added_mana_exists(Some(chain_id), &remainder_address)
+                        || self.output_for_added_mana_exists(None, &remainder_address))
+                {
+                    // Get the lowest of the total diff or the diff for this chain
+                    let mana_to_add = mana_diff.min(input_mana - output_mana);
+                    log::debug!(
+                        "Allocating {mana_to_add} excess input mana for output with address {remainder_address} and chain id {chain_id}"
+                    );
+                    mana_diff -= mana_to_add;
+                    self.remainders.added_mana.insert(Some(chain_id), mana_to_add);
+                }
+            }
+            // Any leftover mana diff can go in any basic output with the right address
+            if mana_diff > 0 && self.output_for_added_mana_exists(None, &remainder_address) {
+                log::debug!("Allocating {mana_diff} excess input mana for output with address {remainder_address}");
+                self.remainders.added_mana.insert(None, std::mem::take(&mut mana_diff));
+            }
         }
 
         if input_amount == output_amount && mana_diff == 0 && native_tokens_diff.is_empty() {
@@ -132,10 +148,10 @@ impl TransactionBuilder {
         Ok((storage_deposit_returns, remainder_outputs))
     }
 
-    fn output_for_added_mana_exists(&self, remainder_address: &Address) -> bool {
+    fn output_for_added_mana_exists(&self, chain_id: Option<ChainId>, remainder_address: &Address) -> bool {
         // Find the first value that matches the remainder address
-        self.non_remainder_outputs().any(|o| {
-            (o.is_basic() || o.is_account() || o.is_anchor() || o.is_nft())
+        self.added_outputs.iter().any(|o| {
+            (o.chain_id() == chain_id || chain_id.is_none() && (o.is_basic() || o.is_account() || o.is_nft()))
                 && o.unlock_conditions()
                     .map_or(true, |uc| uc.expiration().is_none() && uc.timelock().is_none())
                 && matches!(o.required_address(
@@ -145,30 +161,19 @@ impl TransactionBuilder {
         })
     }
 
-    pub(crate) fn get_output_for_added_mana(&mut self, remainder_address: &Address) -> Option<&mut Output> {
-        // Establish the order in which we want to pick an output
-        let sort_order = [AccountOutput::KIND, BasicOutput::KIND, NftOutput::KIND]
-            .into_iter()
-            .zip(0..)
-            .collect::<HashMap<_, _>>();
-        // Remove those that do not have an ordering and sort
-        let ordered_outputs = self
-            .provided_outputs
-            .iter_mut()
-            .chain(&mut self.added_outputs)
-            .filter(|o| {
-                o.unlock_conditions()
+    pub(crate) fn get_output_for_added_mana(
+        &mut self,
+        chain_id: Option<ChainId>,
+        remainder_address: &Address,
+    ) -> Option<&mut Output> {
+        self.added_outputs.iter_mut().find(|o| {
+            (o.chain_id() == chain_id || chain_id.is_none() && (o.is_basic() || o.is_account() || o.is_nft()))
+                && o.unlock_conditions()
                     .map_or(true, |uc| uc.expiration().is_none() && uc.timelock().is_none())
-            })
-            .filter_map(|o| sort_order.get(&o.kind()).map(|order| (*order, o)))
-            .collect::<BTreeMap<_, _>>();
-
-        // Find the first value that matches the remainder address
-        ordered_outputs.into_values().find(|o| {
-            matches!(o.required_address(
-                self.latest_slot_commitment_id.slot_index(),
-                self.protocol_parameters.committable_age_range(),
-            ), Ok(Some(address)) if &address == remainder_address)
+                && matches!(o.required_address(
+                            self.latest_slot_commitment_id.slot_index(),
+                            self.protocol_parameters.committable_age_range(),
+                        ), Ok(Some(address)) if &address == remainder_address)
         })
     }
 
@@ -203,16 +208,30 @@ impl TransactionBuilder {
 
         let (selected_mana, required_mana) = self.mana_sums(false)?;
 
+        log::debug!("selected mana: {selected_mana}, required: {required_mana}");
+
         let remainder_address = self.get_remainder_address()?.map(|v| v.0);
 
+        let mana_chains = self.mana_chains()?;
+
         // Mana can potentially be added to an appropriate existing output instead of a new remainder output
         let mut mana_remainder = selected_mana > required_mana
             && remainder_address.map_or(true, |remainder_address| {
-                !self.output_for_added_mana_exists(&remainder_address)
+                let mut mana_diff = selected_mana - required_mana;
+                for (chain_id, (mana_in, mana_out)) in mana_chains {
+                    if mana_in > mana_out && self.output_for_added_mana_exists(Some(chain_id), &remainder_address) {
+                        mana_diff -= mana_diff.min(mana_in - mana_out);
+                        if mana_diff == 0 {
+                            return false;
+                        }
+                    }
+                }
+                mana_diff > 0 && !self.output_for_added_mana_exists(None, &remainder_address)
             });
         // If we are burning mana, we may not need a mana remainder
         if self.burn.as_ref().map_or(false, |b| b.mana()) {
             let initial_excess = self.initial_mana_excess()?;
+            log::debug!("initial_mana_excess: {initial_excess}");
             mana_remainder &= selected_mana > required_mana + initial_excess;
         }
 

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

@@ -13,7 +13,7 @@ use crate::{
         address::Address,
         input::{Input, UtxoInput},
         mana::ManaAllotment,
-        output::{AccountOutput, AccountOutputBuilder, BasicOutput, FoundryOutput, NftOutput, Output},
+        output::{AccountOutput, AccountOutputBuilder, BasicOutput, ChainId, FoundryOutput, NftOutput, Output},
         payload::{signed_transaction::Transaction, SignedTransactionPayload},
         signature::Ed25519Signature,
         unlock::{AccountUnlock, NftUnlock, ReferenceUnlock, SignatureUnlock, Unlock, Unlocks},
@@ -138,9 +138,8 @@ impl TransactionBuilder {
             .as_mut()
             .ok_or(TransactionBuilderError::UnfulfillableRequirement(Requirement::Mana))?;
         if let Some(output) = self
-            .provided_outputs
+            .added_outputs
             .iter_mut()
-            .chain(&mut self.added_outputs)
             .filter(|o| o.is_account() && o.mana() != 0)
             .find(|o| o.as_account().account_id() == issuer_id)
         {
@@ -296,7 +295,8 @@ impl TransactionBuilder {
         if include_remainders {
             // Add the remainder outputs mana as well as the excess mana we've allocated to add to existing outputs
             // later.
-            required_mana += self.remainder_outputs().map(|o| o.mana()).sum::<u64>() + self.remainders.added_mana;
+            required_mana += self.remainder_outputs().map(|o| o.mana()).sum::<u64>()
+                + self.remainders.added_mana.values().sum::<u64>();
         }
 
         Ok((self.total_selected_mana(None)?, required_mana))
@@ -330,6 +330,24 @@ impl TransactionBuilder {
                 input.output.mana()
             })
     }
+
+    pub(crate) fn mana_chains(&self) -> Result<HashMap<ChainId, (u64, u64)>, TransactionBuilderError> {
+        let include_generated = self.burn.as_ref().map_or(true, |b| !b.generated_mana());
+        let mut res = self
+            .non_remainder_outputs()
+            .filter_map(|o| o.chain_id().map(|id| (id, (0, o.mana()))))
+            .collect::<HashMap<_, _>>();
+        for input in &self.selected_inputs {
+            if let Some(chain_id) = input
+                .output
+                .chain_id()
+                .map(|id| id.or_from_output_id(input.output_id()))
+            {
+                res.entry(chain_id).or_default().0 += self.total_mana(input, include_generated)?;
+            }
+        }
+        Ok(res)
+    }
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]