update work score computations

iotaledger · Oct 9, 2023 · 830ff47 · 830ff47
1 parent 25f183c
commit 830ff47
Show file tree

Hide file tree

Showing 18 changed files with 158 additions and 65 deletions.
diff --git a/sdk/src/types/block/core/basic.rs b/sdk/src/types/block/core/basic.rs
@@ -157,15 +157,15 @@ impl BasicBlock {
 }
 
 impl WorkScore for BasicBlock {
-    fn work_score(&self, workscore_structure: WorkScoreStructure) -> u32 {
-        let mut score = workscore_structure.block;
-        let min_strong_parents_threshold = workscore_structure.min_strong_parents_threshold as usize;
+    fn work_score(&self, work_score_struct: WorkScoreStructure) -> u32 {
+        let mut score = work_score_struct.block;
+        let min_strong_parents_threshold = work_score_struct.min_strong_parents_threshold as usize;
         if self.strong_parents.len() < min_strong_parents_threshold {
             let missing_parents_count = min_strong_parents_threshold - self.strong_parents.len();
-            score += workscore_structure.missing_parent * missing_parents_count as u32;
+            score += work_score_struct.missing_parent * missing_parents_count as u32;
         }
         if let Some(payload) = &*self.payload {
-            score += payload.work_score(workscore_structure);
+            score += payload.work_score(work_score_struct);
         }
         score
     }

diff --git a/sdk/src/types/block/core/mod.rs b/sdk/src/types/block/core/mod.rs
@@ -132,10 +132,10 @@ impl Block {
 }
 
 impl WorkScore for Block {
-    fn work_score(&self, workscore_structure: WorkScoreStructure) -> u32 {
+    fn work_score(&self, work_score_struct: WorkScoreStructure) -> u32 {
         match self {
-            Self::Basic(basic) => basic.work_score(workscore_structure),
-            Self::Validation(validation) => validation.work_score(workscore_structure),
+            Self::Basic(basic) => basic.work_score(work_score_struct),
+            Self::Validation(validation) => validation.work_score(work_score_struct),
         }
     }
 }

diff --git a/sdk/src/types/block/core/validation.rs b/sdk/src/types/block/core/validation.rs
@@ -162,7 +162,7 @@ impl ValidationBlock {
 }
 
 impl WorkScore for ValidationBlock {
-    fn work_score(&self, workscore_structure: WorkScoreStructure) -> u32 {
+    fn work_score(&self, _: WorkScoreStructure) -> u32 {
         // The work score of a validation block is `0`.
         0
     }

diff --git a/sdk/src/types/block/core/wrapper.rs b/sdk/src/types/block/core/wrapper.rs
@@ -110,7 +110,7 @@ impl BlockHeader {
 }
 
 impl WorkScore for BlockHeader {
-    fn work_score(&self, workscore_structure: WorkScoreStructure) -> u32 {
+    fn work_score(&self, _: WorkScoreStructure) -> u32 {
         // The work score of a block header is `0`.
         0
     }
@@ -298,11 +298,11 @@ impl BlockWrapper {
 }
 
 impl WorkScore for BlockWrapper {
-    fn work_score(&self, workscore_structure: WorkScoreStructure) -> u32 {
-        let mut score = workscore_structure.data_byte * self.packed_len() as u32 / 1024;
-        score += self.header.work_score(workscore_structure);
-        score += self.block.work_score(workscore_structure);
-        score += self.signature.work_score(workscore_structure);
+    fn work_score(&self, work_score_struct: WorkScoreStructure) -> u32 {
+        let mut score = work_score_struct.data_byte * self.packed_len() as u32 / 1024;
+        score += self.header.work_score(work_score_struct);
+        score += self.block.work_score(work_score_struct);
+        score += self.signature.work_score(work_score_struct);
         score
     }
 }

diff --git a/sdk/src/types/block/output/account.rs b/sdk/src/types/block/output/account.rs
@@ -25,7 +25,7 @@ use crate::types::{
             NativeTokens, Output, OutputBuilderAmount, OutputId, Rent, RentStructure, StateTransitionError,
             StateTransitionVerifier,
         },
-        protocol::ProtocolParameters,
+        protocol::{ProtocolParameters, WorkScore, WorkScoreStructure},
         semantic::{TransactionFailureReason, ValidationContext},
         unlock::Unlock,
         Error,
@@ -651,6 +651,15 @@ impl StateTransitionVerifier for AccountOutput {
     }
 }
 
+impl WorkScore for AccountOutput {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        let native_token_score = self.native_tokens().work_score(work_score_params);
+        let features_score = self.features().work_score(work_score_params);
+        let immutable_features_score = self.immutable_features().work_score(work_score_params);
+        work_score_params.output + native_token_score + features_score + immutable_features_score
+    }
+}
+
 impl Packable for AccountOutput {
     type UnpackError = Error;
     type UnpackVisitor = ProtocolParameters;

diff --git a/sdk/src/types/block/output/basic.rs b/sdk/src/types/block/output/basic.rs
@@ -16,7 +16,7 @@ use crate::types::{
             verify_output_amount_min, verify_output_amount_packable, verify_output_amount_supply, NativeToken,
             NativeTokens, Output, OutputBuilderAmount, OutputId, Rent, RentStructure,
         },
-        protocol::ProtocolParameters,
+        protocol::{ProtocolParameters, WorkScore, WorkScoreStructure},
         semantic::{TransactionFailureReason, ValidationContext},
         unlock::Unlock,
         Error,
@@ -323,6 +323,14 @@ impl BasicOutput {
     }
 }
 
+impl WorkScore for BasicOutput {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        let native_token_score = self.native_tokens().work_score(work_score_params);
+        let features_score = self.features().work_score(work_score_params);
+        work_score_params.output + native_token_score + features_score
+    }
+}
+
 fn verify_unlock_conditions<const VERIFY: bool>(unlock_conditions: &UnlockConditions) -> Result<(), Error> {
     if VERIFY {
         if unlock_conditions.address().is_none() {

diff --git a/sdk/src/types/block/output/delegation.rs b/sdk/src/types/block/output/delegation.rs
@@ -21,7 +21,7 @@ use crate::types::{
             verify_output_amount_min, verify_output_amount_packable, verify_output_amount_supply, Output,
             OutputBuilderAmount, OutputId, Rent, RentStructure, StateTransitionError, StateTransitionVerifier,
         },
-        protocol::ProtocolParameters,
+        protocol::{ProtocolParameters, WorkScore, WorkScoreStructure},
         semantic::{TransactionFailureReason, ValidationContext},
         slot::EpochIndex,
         unlock::Unlock,
@@ -405,6 +405,12 @@ impl StateTransitionVerifier for DelegationOutput {
     }
 }
 
+impl WorkScore for DelegationOutput {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        work_score_params.output
+    }
+}
+
 impl Packable for DelegationOutput {
     type UnpackError = Error;
     type UnpackVisitor = ProtocolParameters;

diff --git a/sdk/src/types/block/output/feature/mod.rs b/sdk/src/types/block/output/feature/mod.rs
@@ -28,7 +28,11 @@ pub use self::{
     staking::StakingFeature,
     tag::TagFeature,
 };
-use crate::types::block::{create_bitflags, Error};
+use crate::types::block::{
+    create_bitflags,
+    protocol::{WorkScore, WorkScoreStructure},
+    Error,
+};
 
 ///
 #[derive(Clone, Eq, PartialEq, Hash, From, Packable)]
@@ -313,6 +317,20 @@ impl Features {
     }
 }
 
+impl WorkScore for Features {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        let features_score = self
+            .iter()
+            .map(|f| match f {
+                Feature::BlockIssuer(_) => work_score_params.block_issuer,
+                Feature::Staking(_) => work_score_params.staking,
+                _ => 0,
+            })
+            .sum::<u32>();
+        features_score
+    }
+}
+
 #[inline]
 fn verify_unique_sorted<const VERIFY: bool>(features: &[Feature], _: &()) -> Result<(), Error> {
     if VERIFY && !is_unique_sorted(features.iter().map(Feature::kind)) {

diff --git a/sdk/src/types/block/output/foundry.rs b/sdk/src/types/block/output/foundry.rs
@@ -25,7 +25,7 @@ use crate::types::{
             NativeTokens, Output, OutputBuilderAmount, OutputId, Rent, RentStructure, StateTransitionError,
             StateTransitionVerifier, TokenId, TokenScheme,
         },
-        protocol::ProtocolParameters,
+        protocol::{ProtocolParameters, WorkScore, WorkScoreStructure},
         semantic::{TransactionFailureReason, ValidationContext},
         unlock::Unlock,
         Error,
@@ -598,6 +598,20 @@ impl StateTransitionVerifier for FoundryOutput {
     }
 }
 
+impl WorkScore for FoundryOutput {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        let native_token_score = self.native_tokens().work_score(work_score_params);
+        let features_score = self.features().work_score(work_score_params);
+        let immutable_features_score = self.immutable_features().work_score(work_score_params);
+        let token_scheme_score = self
+            .token_scheme()
+            .is_simple()
+            .then_some(work_score_params.native_token)
+            .unwrap_or(0);
+        work_score_params.output + native_token_score + features_score + immutable_features_score + token_scheme_score
+    }
+}
+
 impl Packable for FoundryOutput {
     type UnpackError = Error;
     type UnpackVisitor = ProtocolParameters;

diff --git a/sdk/src/types/block/output/mod.rs b/sdk/src/types/block/output/mod.rs
@@ -58,7 +58,7 @@ pub use self::{
     token_scheme::{SimpleTokenScheme, TokenScheme},
     unlock_condition::{UnlockCondition, UnlockConditions},
 };
-use super::protocol::ProtocolParameters;
+use super::protocol::{ProtocolParameters, WorkScore, WorkScoreStructure};
 use crate::types::block::{address::Address, semantic::ValidationContext, slot::SlotIndex, Error};
 
 /// The maximum number of outputs of a transaction.
@@ -474,6 +474,30 @@ impl Rent for Output {
     }
 }
 
+impl WorkScore for Output {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        match self {
+            Self::Basic(basic) => basic.work_score(work_score_params),
+            Self::Account(account) => account.work_score(work_score_params),
+            Self::Foundry(foundry) => foundry.work_score(work_score_params),
+            Self::Nft(nft) => nft.work_score(work_score_params),
+            Self::Delegation(delegation) => delegation.work_score(work_score_params),
+        }
+    }
+}
+
+impl<T: WorkScore, const N: usize> WorkScore for [T; N] {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        self.as_slice().work_score(work_score_params)
+    }
+}
+
+impl<T: WorkScore> WorkScore for [T] {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        self.iter().map(|o| o.work_score(work_score_params)).sum()
+    }
+}
+
 pub(crate) fn verify_output_amount_min(amount: u64) -> Result<(), Error> {
     if amount < Output::AMOUNT_MIN {
         Err(Error::InvalidOutputAmount(amount))

diff --git a/sdk/src/types/block/output/native_token.rs b/sdk/src/types/block/output/native_token.rs
@@ -12,7 +12,11 @@ use iterator_sorted::is_unique_sorted;
 use packable::{bounded::BoundedU8, prefix::BoxedSlicePrefix, Packable};
 use primitive_types::U256;
 
-use crate::types::block::{output::foundry::FoundryId, Error};
+use crate::types::block::{
+    output::foundry::FoundryId,
+    protocol::{WorkScore, WorkScoreStructure},
+    Error,
+};
 
 impl_id!(pub TokenId, 38, "Unique identifiers of native tokens. The TokenId of native tokens minted by a specific foundry is the same as the FoundryId.");
 
@@ -247,6 +251,13 @@ impl NativeTokens {
     }
 }
 
+// TODO: should we also impl `WorkScore` for `NativeToken` for plain consistency?
+impl WorkScore for NativeTokens {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        self.len() as u32 * work_score_params.native_token
+    }
+}
+
 #[inline]
 fn verify_unique_sorted<const VERIFY: bool>(native_tokens: &[NativeToken], _: &()) -> Result<(), Error> {
     if VERIFY && !is_unique_sorted(native_tokens.iter().map(NativeToken::token_id)) {

diff --git a/sdk/src/types/block/output/nft.rs b/sdk/src/types/block/output/nft.rs
@@ -22,7 +22,7 @@ use crate::types::{
             NativeTokens, Output, OutputBuilderAmount, OutputId, Rent, RentStructure, StateTransitionError,
             StateTransitionVerifier,
         },
-        protocol::ProtocolParameters,
+        protocol::{ProtocolParameters, WorkScore, WorkScoreStructure},
         semantic::{TransactionFailureReason, ValidationContext},
         unlock::Unlock,
         Error,
@@ -461,6 +461,15 @@ impl StateTransitionVerifier for NftOutput {
     }
 }
 
+impl WorkScore for NftOutput {
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        let native_token_score = self.native_tokens().work_score(work_score_params);
+        let features_score = self.features().work_score(work_score_params);
+        let immutable_features_score = self.immutable_features().work_score(work_score_params);
+        work_score_params.output + native_token_score + features_score + immutable_features_score
+    }
+}
+
 impl Packable for NftOutput {
     type UnpackError = Error;
     type UnpackVisitor = ProtocolParameters;

diff --git a/sdk/src/types/block/payload/mod.rs b/sdk/src/types/block/payload/mod.rs
@@ -7,7 +7,7 @@ pub mod tagged_data;
 pub mod transaction;
 
 use alloc::boxed::Box;
-use core::ops::Deref;
+use core::{mem::size_of, ops::Deref};
 
 use packable::{
     error::{UnpackError, UnpackErrorExt},
@@ -65,11 +65,14 @@ impl Payload {
 }
 
 impl WorkScore for Payload {
-    fn work_score(&self, workscore_structure: WorkScoreStructure) -> u32 {
-        match self {
-            Self::Transaction(transaction_payload) => transaction_payload.work_score(workscore_structure),
-            Self::TaggedData(tagged_data_payload) => tagged_data_payload.work_score(workscore_structure),
-        }
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        // 1 byte for the payload kind
+        let score = size_of::<u8>() as u32
+            + match self {
+                Self::Transaction(transaction_payload) => transaction_payload.work_score(work_score_params),
+                Self::TaggedData(tagged_data_payload) => tagged_data_payload.work_score(work_score_params),
+            };
+        score
     }
 }
 

diff --git a/sdk/src/types/block/payload/tagged_data/mod.rs b/sdk/src/types/block/payload/tagged_data/mod.rs
@@ -9,7 +9,7 @@ use core::ops::RangeInclusive;
 use packable::{
     bounded::{BoundedU32, BoundedU8},
     prefix::BoxedSlicePrefix,
-    Packable,
+    Packable, PackableExt,
 };
 
 use crate::types::block::{
@@ -64,9 +64,9 @@ impl TaggedDataPayload {
 }
 
 impl WorkScore for TaggedDataPayload {
-    fn work_score(&self, workscore_structure: WorkScoreStructure) -> u32 {
-        // The work score of a tagged data payload is `0`.
-        0
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        let size_score = self.packed_len() as u32 * work_score_params.data_byte;
+        size_score
     }
 }
 

diff --git a/sdk/src/types/block/payload/transaction/essence/regular.rs b/sdk/src/types/block/payload/transaction/essence/regular.rs
@@ -275,29 +275,12 @@ impl RegularTransactionEssence {
 }
 
 impl WorkScore for RegularTransactionEssence {
-    fn work_score(&self, workscore_structure: WorkScoreStructure) -> u32 {
-        let mut score = self.inputs().len() as u32 * workscore_structure.input;
-        score += self.context_inputs().len() as u32 * workscore_structure.context_input;
-        for output in self.outputs() {
-            score += workscore_structure.output;
-            if let Some(nts) = output.native_tokens() {
-                score += nts.len() as u32 * workscore_structure.native_token;
-            }
-            if matches!(output, Output::Foundry(foundry) if foundry.token_scheme().is_simple()) {
-                score += workscore_structure.native_token
-            }
-            if let Some(features) = output.features() {
-                for feature in features.iter() {
-                    score += match feature {
-                        Feature::BlockIssuer(_) => workscore_structure.block_issuer,
-                        Feature::Staking(_) => workscore_structure.staking,
-                        _ => 0,
-                    }
-                }
-            }
-        }
-        score += self.mana_allotments().len() as u32 * workscore_structure.allotment;
-        score
+    fn work_score(&self, work_score_params: WorkScoreStructure) -> u32 {
+        let input_score = self.inputs().len() as u32 * work_score_params.input;
+        let context_input_score = self.context_inputs().len() as u32 * work_score_params.context_input;
+        let outputs_score = self.outputs().work_score(work_score_params);
+        let allotment_score = self.mana_allotments().len() as u32 * work_score_params.allotment;
+        input_score + context_input_score + outputs_score + allotment_score
     }
 }