Add transaction capabilities

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

@@ -23,8 +23,9 @@ use crate::{
         input::INPUT_COUNT_RANGE,
         output::{
             AccountOutput, AccountTransition, ChainId, FoundryOutput, NativeTokensBuilder, NftOutput, Output, OutputId,
-            OUTPUT_COUNT_RANGE,
+            TokenId, OUTPUT_COUNT_RANGE,
         },
+        payload::transaction::{TransactionCapabilities, TransactionCapabilityFlag},
         protocol::ProtocolParameters,
         slot::SlotIndex,
     },
@@ -487,19 +488,32 @@ impl InputSelection {
                     }
                 }
                 Output::Foundry(foundry_output) => {
+                    let foundry_id = foundry_output.id();
                     let foundry_input = input_foundries.iter().find(|i| {
                         if let Output::Foundry(foundry_input) = &i.output {
-                            foundry_output.id() == foundry_input.id()
+                            foundry_id == foundry_input.id()
                         } else {
                             false
                         }
                     });
                     if let Some(foundry_input) = foundry_input {
+                        let token_id = TokenId::from(foundry_id);
+                        let mut capabilities = TransactionCapabilities::default();
+                        // TODO is this really the right approach?
+                        if self
+                            .burn
+                            .as_ref()
+                            .map(|burn| burn.native_tokens.contains_key(&token_id))
+                            .unwrap_or_default()
+                        {
+                            capabilities.add_capability(TransactionCapabilityFlag::BurnNativeTokens);
+                        }
                         if let Err(err) = FoundryOutput::transition_inner(
                             foundry_input.output.as_foundry(),
                             foundry_output,
                             input_native_tokens_builder.deref(),
                             output_native_tokens_builder.deref(),
+                            &capabilities,
                         ) {
                             log::debug!("validate_transitions error {err:?}");
                             return Err(Error::UnfulfillableRequirement(Requirement::Foundry(

sdk/src/types/block/address/restricted.rs

@@ -1,14 +1,14 @@
 // Copyright 2023 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
-use alloc::boxed::Box;
-
-use derive_more::Deref;
 use getset::Getters;
-use packable::{error::UnpackErrorExt, prefix::BoxedSlicePrefix, Packable, PackableExt};
+use packable::{Packable, PackableExt};
 
 use super::Address;
-use crate::types::block::Error;
+use crate::types::block::{
+    capabilities::{Capabilities, CapabilityFlag},
+    Error,
+};
 
 #[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd, Hash, Getters, Packable)]
 #[getset(get = "pub")]
@@ -94,9 +94,12 @@ impl AddressCapabilityFlag {
     const ACCOUNT_OUTPUTS: u8 = 0b00100000;
     const NFT_OUTPUTS: u8 = 0b01000000;
     const DELEGATION_OUTPUTS: u8 = 0b10000000;
+}
 
-    /// Converts the flag into the byte representation.
-    pub fn as_byte(&self) -> u8 {
+impl CapabilityFlag for AddressCapabilityFlag {
+    type Iterator = core::array::IntoIter<Self, 8>;
+
+    fn as_byte(&self) -> u8 {
         match self {
             Self::OutputsWithNativeTokens => Self::OUTPUTS_WITH_NATIVE_TOKENS,
             Self::OutputsWithMana => Self::OUTPUTS_WITH_MANA,
@@ -109,8 +112,7 @@ impl AddressCapabilityFlag {
         }
     }
 
-    /// Returns the index in [`AddressCapabilities`] to which this flag is applied.
-    pub fn index(&self) -> usize {
+    fn index(&self) -> usize {
         match self {
             Self::OutputsWithNativeTokens
             | Self::OutputsWithMana
@@ -123,8 +125,7 @@ impl AddressCapabilityFlag {
         }
     }
 
-    /// Returns an iterator over all flags.
-    pub fn all() -> impl Iterator<Item = Self> {
+    fn all() -> Self::Iterator {
         [
             Self::OutputsWithNativeTokens,
             Self::OutputsWithMana,
@@ -139,139 +140,7 @@ impl AddressCapabilityFlag {
     }
 }
 
-/// A list of bitflags that represent the capabilities of an [`Address`].
-/// If an output is created by a transaction with an
-/// [`UnlockCondition`](crate::types::block::output::UnlockCondition) containing a [`RestrictedAddress`], the
-/// transaction is valid only if the specified conditions, corresponding to the [`AddressCapabilityFlag`]s, hold for
-/// that output.
-#[derive(Clone, Debug, Default, Eq, PartialEq, Ord, PartialOrd, Hash, Deref)]
-#[repr(transparent)]
-pub struct AddressCapabilities(BoxedSlicePrefix<u8, u8>);
-
-impl AddressCapabilities {
-    /// Returns an [`AddressCapabilities`] with every possible flag enabled.
-    pub fn all() -> Self {
-        let mut res = Self::default();
-        res.set_all();
-        res
-    }
-
-    /// Returns an [`AddressCapabilities`] with every possible flag disabled.
-    pub fn none() -> Self {
-        Self::default()
-    }
-
-    /// Returns whether every possible [`AddressCapabilityFlag`] is enabled.
-    pub fn is_all(&self) -> bool {
-        AddressCapabilityFlag::all().all(|flag| self.has_capability(flag))
-    }
-
-    /// Returns whether every possible [`AddressCapabilityFlag`] is disabled.
-    pub fn is_none(&self) -> bool {
-        self.0.iter().all(|b| 0.eq(b))
-    }
-
-    /// Enables every possible [`AddressCapabilityFlag`].
-    pub fn set_all(&mut self) -> &mut Self {
-        for flag in AddressCapabilityFlag::all() {
-            self.add_capability(flag);
-        }
-        self
-    }
-
-    /// Disabled every possible [`AddressCapabilityFlag`].
-    pub fn set_none(&mut self) -> &mut Self {
-        *self = Default::default();
-        self
-    }
-
-    /// Enables a given [`AddressCapabilityFlag`].
-    pub fn add_capability(&mut self, flag: AddressCapabilityFlag) -> &mut Self {
-        if self.0.len() <= flag.index() {
-            let mut v = Box::<[_]>::from(self.0.clone()).into_vec();
-            v.resize(flag.index() + 1, 0);
-            // Unwrap: safe because the indexes are within u8 bounds
-            self.0 = v.into_boxed_slice().try_into().unwrap();
-        }
-        self.0[flag.index()] |= flag.as_byte();
-        self
-    }
-
-    /// Enables a given set of [`AddressCapabilityFlag`]s.
-    pub fn add_capabilities(&mut self, flags: impl IntoIterator<Item = AddressCapabilityFlag>) -> &mut Self {
-        for flag in flags {
-            self.add_capability(flag);
-        }
-        self
-    }
-
-    /// Enables a given set of [`AddressCapabilityFlag`]s.
-    pub fn with_capabilities(mut self, flags: impl IntoIterator<Item = AddressCapabilityFlag>) -> Self {
-        self.add_capabilities(flags);
-        self
-    }
-
-    /// Enables a given set of [`AddressCapabilityFlag`]s.
-    pub fn set_capabilities(&mut self, flags: impl IntoIterator<Item = AddressCapabilityFlag>) -> &mut Self {
-        *self = Self::default().with_capabilities(flags);
-        self
-    }
-
-    /// Returns whether a given [`AddressCapabilityFlag`] is enabled.
-    pub fn has_capability(&self, flag: AddressCapabilityFlag) -> bool {
-        self.0
-            .get(flag.index())
-            .map(|byte| byte & flag.as_byte() == flag.as_byte())
-            .unwrap_or_default()
-    }
-
-    /// Returns whether a given set of [`AddressCapabilityFlag`]s are enabled.
-    pub fn has_capabilities(&self, flags: impl IntoIterator<Item = AddressCapabilityFlag>) -> bool {
-        flags.into_iter().all(|flag| self.has_capability(flag))
-    }
-
-    /// Returns an iterator over all enabled [`AddressCapabilityFlag`]s.
-    pub fn capabilities_iter(&self) -> impl Iterator<Item = AddressCapabilityFlag> + '_ {
-        self.0.iter().enumerate().flat_map(|(idx, byte)| {
-            AddressCapabilityFlag::all().filter(move |f| (idx == f.index() && byte & f.as_byte() == f.as_byte()))
-        })
-    }
-}
-
-impl<I: IntoIterator<Item = AddressCapabilityFlag>> From<I> for AddressCapabilities {
-    fn from(value: I) -> Self {
-        Self::default().with_capabilities(value)
-    }
-}
-
-impl Packable for AddressCapabilities {
-    type UnpackError = Error;
-    type UnpackVisitor = ();
-
-    fn pack<P: packable::packer::Packer>(&self, packer: &mut P) -> Result<(), P::Error> {
-        if !self.is_none() {
-            self.0.pack(packer)?;
-        } else {
-            0_u8.pack(packer)?;
-        }
-        Ok(())
-    }
-
-    fn unpack<U: packable::unpacker::Unpacker, const VERIFY: bool>(
-        unpacker: &mut U,
-        visitor: &Self::UnpackVisitor,
-    ) -> Result<Self, packable::error::UnpackError<Self::UnpackError, U::Error>> {
-        use packable::prefix::UnpackPrefixError;
-        Ok(Self(
-            BoxedSlicePrefix::unpack::<_, VERIFY>(unpacker, visitor)
-                // TODO: not sure if this is the best way to do this
-                .map_packable_err(|e| match e {
-                    UnpackPrefixError::Item(i) | UnpackPrefixError::Prefix(i) => i,
-                })
-                .coerce()?,
-        ))
-    }
-}
+pub type AddressCapabilities = Capabilities<AddressCapabilityFlag>;
 
 #[cfg(feature = "serde")]
 pub(crate) mod dto {
@@ -314,12 +183,14 @@ pub(crate) mod dto {
         type Error = Error;
 
         fn try_from(value: RestrictedAddressDto) -> Result<Self, Self::Error> {
-            Ok(Self::new(value.address)?.with_allowed_capabilities(AddressCapabilities(
-                value
-                    .allowed_capabilities
-                    .try_into()
-                    .map_err(Error::InvalidAddressCapabilitiesCount)?,
-            )))
+            Ok(
+                Self::new(value.address)?.with_allowed_capabilities(AddressCapabilities::from_bytes(
+                    value
+                        .allowed_capabilities
+                        .try_into()
+                        .map_err(Error::InvalidCapabilitiesCount)?,
+                )),
+            )
         }
     }