Skip to content

Commit

Permalink
Make Allotments unique sorted
Browse files Browse the repository at this point in the history
  • Loading branch information
Alex Coats committed Aug 3, 2023
1 parent a1e2802 commit b63d2ae
Show file tree
Hide file tree
Showing 4 changed files with 66 additions and 24 deletions.
4 changes: 2 additions & 2 deletions sdk/src/types/block/error.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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}")
}
Expand Down
12 changes: 12 additions & 0 deletions sdk/src/types/block/mana/allotment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
56 changes: 40 additions & 16 deletions sdk/src/types/block/mana/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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));
}
}

Expand All @@ -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>;
Expand Down Expand Up @@ -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> {
Expand Down
18 changes: 12 additions & 6 deletions sdk/src/types/block/payload/transaction/essence/regular.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand All @@ -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>,
}
Expand All @@ -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,
}
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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)?;

Expand Down

0 comments on commit b63d2ae

Please sign in to comment.