feat(iota-genesis-builder): Add native token circulating supply validation

crates/iota-genesis-builder/src/stardust/migration/tests/basic.rs

@@ -589,6 +589,6 @@ fn basic_migration_with_incorrect_total_supply() {
 
     assert_eq!(
         err.to_string(),
-        format!("total supply mismatch: found {total_value}, expected {total_supply}")
+        format!("base token total supply: found {total_value}, expected {total_supply}")
     );
 }

crates/iota-genesis-builder/src/stardust/migration/verification/alias.rs

@@ -23,8 +23,9 @@ use crate::stardust::{
         verification::{
             created_objects::CreatedObjects,
             util::{
-                verify_address_owner, verify_issuer_feature, verify_metadata_feature,
-                verify_native_tokens, verify_parent, verify_sender_feature,
+                TokensAmountCounter, verify_address_owner, verify_issuer_feature,
+                verify_metadata_feature, verify_native_tokens, verify_parent,
+                verify_sender_feature,
             },
         },
     },
@@ -37,7 +38,7 @@ pub(super) fn verify_alias_output(
     created_objects: &CreatedObjects,
     foundry_data: &HashMap<stardust::TokenId, FoundryLedgerData>,
     storage: &InMemoryStorage,
-    total_value: &mut u64,
+    tokens_counter: &mut TokensAmountCounter,
     address_swap_map: &AddressSwapMap,
 ) -> anyhow::Result<()> {
     let alias_id = ObjectID::new(*output.alias_id_non_null(&output_id));
@@ -95,7 +96,7 @@ pub(super) fn verify_alias_output(
         created_output.balance.value(),
         output.amount()
     );
-    *total_value += created_output.balance.value();
+    tokens_counter.update_total_value_for_iota(created_output.balance.value());
 
     // Native Tokens
     verify_native_tokens::<Field<String, Balance>>(
@@ -104,6 +105,7 @@ pub(super) fn verify_alias_output(
         created_output.native_tokens,
         created_objects.native_tokens().ok(),
         storage,
+        tokens_counter,
     )?;
 
     // Legacy State Controller

crates/iota-genesis-builder/src/stardust/migration/verification/basic.rs

@@ -24,10 +24,10 @@ use crate::stardust::{
         verification::{
             created_objects::CreatedObjects,
             util::{
-                verify_address_owner, verify_coin, verify_expiration_unlock_condition,
-                verify_metadata_feature, verify_native_tokens, verify_parent,
-                verify_sender_feature, verify_storage_deposit_unlock_condition, verify_tag_feature,
-                verify_timelock_unlock_condition,
+                TokensAmountCounter, verify_address_owner, verify_coin,
+                verify_expiration_unlock_condition, verify_metadata_feature, verify_native_tokens,
+                verify_parent, verify_sender_feature, verify_storage_deposit_unlock_condition,
+                verify_tag_feature, verify_timelock_unlock_condition,
             },
         },
     },
@@ -41,7 +41,7 @@ pub(super) fn verify_basic_output(
     foundry_data: &HashMap<TokenId, FoundryLedgerData>,
     target_milestone_timestamp: u32,
     storage: &InMemoryStorage,
-    total_value: &mut u64,
+    tokens_counter: &mut TokensAmountCounter,
     address_swap_map: &AddressSwapMap,
 ) -> Result<()> {
     // If this is a timelocked vested reward, a `Timelock<Balance>` is created.
@@ -73,7 +73,7 @@ pub(super) fn verify_basic_output(
             created_timelock.locked().value(),
             output.amount()
         );
-        *total_value += created_timelock.locked().value();
+        tokens_counter.update_total_value_for_iota(created_timelock.locked().value());
 
         // Label
         let label = created_timelock
@@ -139,7 +139,7 @@ pub(super) fn verify_basic_output(
             created_output.balance.value(),
             output.amount()
         );
-        *total_value += created_output.balance.value();
+        tokens_counter.update_total_value_for_iota(created_output.balance.value());
 
         // Native Tokens
         verify_native_tokens::<Field<String, Balance>>(
@@ -148,6 +148,7 @@ pub(super) fn verify_basic_output(
             created_output.native_tokens,
             created_objects.native_tokens().ok(),
             storage,
+            tokens_counter,
         )?;
 
         // Storage Deposit Return Unlock Condition
@@ -202,7 +203,7 @@ pub(super) fn verify_basic_output(
 
         verify_address_owner(output.address(), created_coin_obj, "coin", address_swap_map)?;
         verify_coin(output.amount(), &created_coin)?;
-        *total_value += created_coin.value();
+        tokens_counter.update_total_value_for_iota(created_coin.value());
 
         // Native Tokens
         verify_native_tokens::<(TypeTag, Coin)>(
@@ -211,6 +212,7 @@ pub(super) fn verify_basic_output(
             None,
             created_objects.native_tokens().ok(),
             storage,
+            tokens_counter,
         )?;
     }
 

crates/iota-genesis-builder/src/stardust/migration/verification/foundry.rs

@@ -17,8 +17,8 @@ use crate::stardust::{
         verification::{
             CreatedObjects,
             util::{
-                truncate_to_max_allowed_u64_supply, verify_address_owner, verify_coin,
-                verify_parent, verify_shared_object,
+                TokensAmountCounter, truncate_to_max_allowed_u64_supply, verify_address_owner,
+                verify_coin, verify_parent, verify_shared_object,
             },
         },
     },
@@ -32,7 +32,7 @@ pub(super) fn verify_foundry_output(
     created_objects: &CreatedObjects,
     foundry_data: &HashMap<TokenId, FoundryLedgerData>,
     storage: &InMemoryStorage,
-    total_value: &mut u64,
+    tokens_counter: &mut TokensAmountCounter,
     address_swap_map: &AddressSwapMap,
 ) -> Result<()> {
     let foundry_data = foundry_data
@@ -57,7 +57,7 @@ pub(super) fn verify_foundry_output(
 
     verify_address_owner(alias_address, created_coin_obj, "coin", address_swap_map)?;
     verify_coin(output.amount(), &created_coin)?;
-    *total_value += created_coin.value();
+    tokens_counter.update_total_value_for_iota(created_coin.value());
 
     // Native token coin value
     let native_token_coin_id = created_objects.native_token_coin()?;
@@ -234,6 +234,8 @@ pub(super) fn verify_foundry_output(
         coin_manager.treasury_cap.total_supply.value,
         circulating_supply
     );
+    tokens_counter
+        .update_total_value_max(&foundry_data.to_canonical_string(false), circulating_supply);
 
     // Alias Address Unlock Condition
     verify_address_owner(

crates/iota-genesis-builder/src/stardust/migration/verification/mod.rs

@@ -9,6 +9,8 @@ use std::collections::HashMap;
 use anyhow::{anyhow, ensure};
 use iota_sdk::types::block::output::{Output, OutputId, TokenId};
 use iota_types::in_memory_storage::InMemoryStorage;
+use tracing::warn;
+use util::{BASE_TOKEN_KEY, TokensAmountCounter};
 
 use self::created_objects::CreatedObjects;
 use crate::stardust::{
@@ -32,7 +34,7 @@ pub(crate) fn verify_outputs<'a>(
     storage: &InMemoryStorage,
     address_swap_map: &AddressSwapMap,
 ) -> anyhow::Result<()> {
-    let mut total_value = 0;
+    let mut tokens_counter = TokensAmountCounter::new(total_supply);
     for (header, output) in outputs {
         let created_objects = output_objects_map
             .get(&header.output_id())
@@ -44,14 +46,22 @@ pub(crate) fn verify_outputs<'a>(
             foundry_data,
             target_milestone_timestamp,
             storage,
-            &mut total_value,
+            &mut tokens_counter,
             address_swap_map,
         )?;
     }
-    ensure!(
-        total_supply == total_value,
-        "total supply mismatch: found {total_value}, expected {total_supply}"
-    );
+    for (key, (total_value, expected_value)) in tokens_counter.into_inner() {
+        if key == BASE_TOKEN_KEY {
+            ensure!(
+                total_value == expected_value,
+                "base token total supply: found {total_value}, expected {expected_value}"
+            )
+        } else if expected_value != total_value {
+            warn!(
+                "total supply mismatch for {key}: found {total_value}, expected {expected_value}"
+            );
+        }
+    }
     Ok(())
 }
 
@@ -62,7 +72,7 @@ fn verify_output(
     foundry_data: &HashMap<TokenId, FoundryLedgerData>,
     target_milestone_timestamp: u32,
     storage: &InMemoryStorage,
-    total_value: &mut u64,
+    tokens_counter: &mut TokensAmountCounter,
     address_swap_map: &AddressSwapMap,
 ) -> anyhow::Result<()> {
     match output {
@@ -72,7 +82,7 @@ fn verify_output(
             created_objects,
             foundry_data,
             storage,
-            total_value,
+            tokens_counter,
             address_swap_map,
         ),
         Output::Basic(output) => basic::verify_basic_output(
@@ -82,7 +92,7 @@ fn verify_output(
             foundry_data,
             target_milestone_timestamp,
             storage,
-            total_value,
+            tokens_counter,
             address_swap_map,
         ),
         Output::Foundry(output) => foundry::verify_foundry_output(
@@ -91,7 +101,7 @@ fn verify_output(
             created_objects,
             foundry_data,
             storage,
-            total_value,
+            tokens_counter,
             address_swap_map,
         ),
         Output::Nft(output) => nft::verify_nft_output(
@@ -100,7 +110,7 @@ fn verify_output(
             created_objects,
             foundry_data,
             storage,
-            total_value,
+            tokens_counter,
             address_swap_map,
         ),
         // Treasury outputs aren't used since Stardust, so no need to verify anything here.

crates/iota-genesis-builder/src/stardust/migration/verification/nft.rs

@@ -21,10 +21,10 @@ use crate::stardust::{
         verification::{
             created_objects::CreatedObjects,
             util::{
-                verify_address_owner, verify_expiration_unlock_condition, verify_issuer_feature,
-                verify_metadata_feature, verify_native_tokens, verify_parent,
-                verify_sender_feature, verify_storage_deposit_unlock_condition, verify_tag_feature,
-                verify_timelock_unlock_condition,
+                TokensAmountCounter, verify_address_owner, verify_expiration_unlock_condition,
+                verify_issuer_feature, verify_metadata_feature, verify_native_tokens,
+                verify_parent, verify_sender_feature, verify_storage_deposit_unlock_condition,
+                verify_tag_feature, verify_timelock_unlock_condition,
             },
         },
     },
@@ -37,7 +37,7 @@ pub(super) fn verify_nft_output(
     created_objects: &CreatedObjects,
     foundry_data: &HashMap<TokenId, FoundryLedgerData>,
     storage: &InMemoryStorage,
-    total_value: &mut u64,
+    tokens_counter: &mut TokensAmountCounter,
     address_swap_map: &AddressSwapMap,
 ) -> anyhow::Result<()> {
     let created_output_obj = created_objects.output().and_then(|id| {
@@ -100,7 +100,7 @@ pub(super) fn verify_nft_output(
         created_output.balance.value(),
         output.amount()
     );
-    *total_value += created_output.balance.value();
+    tokens_counter.update_total_value_for_iota(created_output.balance.value());
 
     // Native Tokens
     verify_native_tokens::<Field<String, Balance>>(
@@ -109,6 +109,7 @@ pub(super) fn verify_nft_output(
         created_output.native_tokens,
         created_objects.native_tokens().ok(),
         storage,
+        tokens_counter,
     )?;
 
     // Storage Deposit Return Unlock Condition

crates/iota-genesis-builder/src/stardust/migration/verification/util.rs

@@ -32,12 +32,58 @@ use crate::stardust::{
     types::{address_swap_map::AddressSwapMap, token_scheme::MAX_ALLOWED_U64_SUPPLY},
 };
 
+pub const BASE_TOKEN_KEY: &str = "base_token";
+
+/// Counter used to count the generated tokens amounts.
+pub(super) struct TokensAmountCounter {
+    // A map of token type -> (real_generated_supply, expected_circulating_supply)
+    inner: HashMap<String, (u64, u64)>,
+}
+
+impl TokensAmountCounter {
+    /// Setup the tokens amount counter.
+    pub(super) fn new(initial_iota_supply: u64) -> Self {
+        let mut res = TokensAmountCounter {
+            inner: HashMap::new(),
+        };
+        res.update_total_value_max_for_iota(initial_iota_supply);
+        res
+    }
+
+    pub(super) fn into_inner(self) -> impl IntoIterator<Item = (String, (u64, u64))> {
+        self.inner
+    }
+
+    pub(super) fn update_total_value_for_iota(&mut self, value: u64) {
+        self.update_total_value(BASE_TOKEN_KEY, value);
+    }
+
+    fn update_total_value_max_for_iota(&mut self, max: u64) {
+        self.update_total_value_max(BASE_TOKEN_KEY, max);
+    }
+
+    pub(super) fn update_total_value(&mut self, key: &str, value: u64) {
+        self.inner
+            .entry(key.to_string())
+            .and_modify(|v| v.0 += value)
+            .or_insert((value, 0));
+    }
+
+    pub(super) fn update_total_value_max(&mut self, key: &str, max: u64) {
+        self.inner
+            .entry(key.to_string())
+            .and_modify(|v| v.1 = max)
+            .or_insert((0, max));
+    }
+}
+
 pub(super) fn verify_native_tokens<NtKind: NativeTokenKind>(
     native_tokens: &NativeTokens,
     foundry_data: &HashMap<TokenId, FoundryLedgerData>,
     native_tokens_bag: impl Into<Option<Bag>>,
     created_native_tokens: Option<&[ObjectID]>,
     storage: &InMemoryStorage,
+    tokens_counter: &mut TokensAmountCounter,
 ) -> Result<()> {
     // Token types should be unique as the token ID is guaranteed unique within
     // NativeTokens
@@ -83,11 +129,12 @@ pub(super) fn verify_native_tokens<NtKind: NativeTokenKind>(
             .token_scheme_u64
             .adjust_tokens(native_token.amount());
 
-        if let Some(created_value) = created_native_tokens.get(&expected_bag_key) {
+        if let Some(&created_value) = created_native_tokens.get(&expected_bag_key) {
             ensure!(
-                *created_value == reduced_amount,
+                created_value == reduced_amount,
                 "created token amount mismatch: found {created_value}, expected {reduced_amount}"
             );
+            tokens_counter.update_total_value(&expected_bag_key, created_value);
         } else {
             bail!(
                 "native token object was not created for token: {}",