Make Allotments unique sorted (#989)

sdk/src/types/block/error.rs

@@ -24,12 +24,12 @@ use crate::types::block::{
 #[derive(Debug, PartialEq, Eq)]
 #[allow(missing_docs)]
 pub enum Error {
+    AllotmentsNotUniqueSorted,
     ConsumedAmountOverflow,
     ConsumedNativeTokensAmountOverflow,
     CreatedAmountOverflow,
     CreatedNativeTokensAmountOverflow,
     Crypto(CryptoError),
-    DuplicateAllotment(AccountId),
     DuplicateSignatureUnlock(u16),
     DuplicateUtxo(UtxoInput),
     ExpirationUnlockConditionZero,
@@ -122,12 +122,12 @@ impl std::error::Error for Error {}
 impl fmt::Display for Error {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         match self {
+            Self::AllotmentsNotUniqueSorted => write!(f, "allotments are not unique and/or sorted"),
             Self::ConsumedAmountOverflow => write!(f, "consumed amount overflow"),
             Self::ConsumedNativeTokensAmountOverflow => write!(f, "consumed native tokens amount overflow"),
             Self::CreatedAmountOverflow => write!(f, "created amount overflow"),
             Self::CreatedNativeTokensAmountOverflow => write!(f, "created native tokens amount overflow"),
             Self::Crypto(e) => write!(f, "cryptographic error: {e}"),
-            Self::DuplicateAllotment(id) => write!(f, "duplicate allotment, account ID: {id}"),
             Self::DuplicateSignatureUnlock(index) => {
                 write!(f, "duplicate signature unlock at index: {index}")
             }

sdk/src/types/block/mana/allotment.rs

@@ -21,6 +21,18 @@ pub struct Allotment {
     pub(crate) mana: u64,
 }
 
+impl PartialOrd for Allotment {
+    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
+        Some(self.cmp(other))
+    }
+}
+
+impl Ord for Allotment {
+    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
+        self.account_id.cmp(&other.account_id)
+    }
+}
+
 impl Allotment {
     pub fn new(account_id: AccountId, mana: u64) -> Result<Self, Error> {
         if mana > MAX_THEORETICAL_MANA {

sdk/src/types/block/mana/mod.rs

@@ -3,11 +3,11 @@
 
 mod allotment;
 
-use alloc::{boxed::Box, vec::Vec};
+use alloc::{boxed::Box, collections::BTreeSet, vec::Vec};
 use core::ops::RangeInclusive;
 
 use derive_more::Deref;
-use hashbrown::HashSet;
+use iterator_sorted::is_unique_sorted;
 use packable::{bounded::BoundedU16, prefix::BoxedSlicePrefix, Packable};
 
 pub use self::allotment::Allotment;
@@ -27,20 +27,24 @@ pub struct Allotments(#[packable(verify_with = verify_allotments)] BoxedSlicePre
 
 fn verify_allotments<const VERIFY: bool>(allotments: &[Allotment], _visitor: &()) -> Result<(), Error> {
     if VERIFY {
-        let mut mana_sum: u64 = 0;
-        let mut unique_ids = HashSet::with_capacity(allotments.len());
-        for Allotment { account_id, mana } in allotments.iter() {
-            mana_sum = mana_sum
-                .checked_add(*mana)
-                .ok_or(Error::InvalidAllotmentManaSum(mana_sum as u128 + *mana as u128))?;
-
-            if mana_sum > MAX_THEORETICAL_MANA {
-                return Err(Error::InvalidAllotmentManaSum(mana_sum as u128));
-            }
-
-            if !unique_ids.insert(account_id) {
-                return Err(Error::DuplicateAllotment(*account_id));
-            }
+        if !is_unique_sorted(allotments.iter().map(|a| a.account_id)) {
+            return Err(Error::AllotmentsNotUniqueSorted);
+        }
+        verify_allotments_sum(allotments)?;
+    }
+
+    Ok(())
+}
+
+fn verify_allotments_sum<'a>(allotments: impl IntoIterator<Item = &'a Allotment>) -> Result<(), Error> {
+    let mut mana_sum: u64 = 0;
+    for Allotment { mana, .. } in allotments {
+        mana_sum = mana_sum
+            .checked_add(*mana)
+            .ok_or(Error::InvalidAllotmentManaSum(mana_sum as u128 + *mana as u128))?;
+
+        if mana_sum > MAX_THEORETICAL_MANA {
+            return Err(Error::InvalidAllotmentManaSum(mana_sum as u128));
         }
     }
 
@@ -56,6 +60,15 @@ impl TryFrom<Vec<Allotment>> for Allotments {
     }
 }
 
+impl TryFrom<BTreeSet<Allotment>> for Allotments {
+    type Error = Error;
+
+    #[inline(always)]
+    fn try_from(allotments: BTreeSet<Allotment>) -> Result<Self, Self::Error> {
+        Self::from_set(allotments)
+    }
+}
+
 impl IntoIterator for Allotments {
     type Item = Allotment;
     type IntoIter = alloc::vec::IntoIter<Self::Item>;
@@ -83,6 +96,17 @@ impl Allotments {
         Ok(Self(allotments))
     }
 
+    /// Creates a new [`Allotments`] from an ordered set.
+    pub fn from_set(allotments: BTreeSet<Allotment>) -> Result<Self, Error> {
+        let allotments =
+            BoxedSlicePrefix::<Allotment, AllotmentCount>::try_from(allotments.into_iter().collect::<Box<[_]>>())
+                .map_err(Error::InvalidAllotmentCount)?;
+
+        verify_allotments_sum(allotments.as_ref())?;
+
+        Ok(Self(allotments))
+    }
+
     /// Gets a reference to an [`Allotment`], if one exists, using an [`AccountId`].
     #[inline(always)]
     pub fn get(&self, account_id: &AccountId) -> Option<&Allotment> {

sdk/src/types/block/payload/transaction/essence/regular.rs

@@ -1,7 +1,7 @@
 // Copyright 2020-2021 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
-use alloc::vec::Vec;
+use alloc::{collections::BTreeSet, vec::Vec};
 
 use hashbrown::HashSet;
 use packable::{bounded::BoundedU16, prefix::BoxedSlicePrefix, Packable};
@@ -26,7 +26,7 @@ pub struct RegularTransactionEssenceBuilder {
     inputs: Vec<Input>,
     inputs_commitment: InputsCommitment,
     outputs: Vec<Output>,
-    allotments: Vec<Allotment>,
+    allotments: BTreeSet<Allotment>,
     payload: OptionalPayload,
     creation_time: Option<u64>,
 }
@@ -39,7 +39,7 @@ impl RegularTransactionEssenceBuilder {
             inputs: Vec::new(),
             inputs_commitment,
             outputs: Vec::new(),
-            allotments: Vec::new(),
+            allotments: BTreeSet::new(),
             payload: OptionalPayload::default(),
             creation_time: None,
         }
@@ -81,9 +81,15 @@ impl RegularTransactionEssenceBuilder {
         self
     }
 
-    /// Add an allotment to a [`RegularTransactionEssenceBuilder`].
+    /// Add an [`Allotment`] to a [`RegularTransactionEssenceBuilder`].
     pub fn add_allotment(mut self, allotment: Allotment) -> Self {
-        self.allotments.push(allotment);
+        self.allotments.insert(allotment);
+        self
+    }
+
+    /// Replaces an [`Allotment`] of the [`RegularTransactionEssenceBuilder`] with a new one, or adds it.
+    pub fn replace_allotment(mut self, allotment: Allotment) -> Self {
+        self.allotments.replace(allotment);
         self
     }
 
@@ -126,7 +132,7 @@ impl RegularTransactionEssenceBuilder {
             verify_outputs::<true>(&outputs, protocol_parameters)?;
         }
 
-        let allotments = Allotments::from_vec(self.allotments)?;
+        let allotments = Allotments::from_set(self.allotments)?;
 
         verify_payload::<true>(&self.payload)?;