feat(iota-e2e-tests): Include object snapshot generation and address …

…swap in migration test (#4389)

* feat(iota-genesis-builder): add address swap map for swaping origin addresses to destination during the migration process

* fix(iota-genesis-builder): relative path supporting, remove mut modifier

* fix(iota-genesis-builder): address_swap_map become field of Migration struct, minor renaming refactoring

* fix(iota-genesis-builder): add additional hashmap that track if address has been swapped at least once
Create dedicated function to convert and swap addresses

* fix(iota-genesis-builder): review's cmments fix

* fix(iota-genesis-builder): anyhow error has been added instead of panic

* cargo clippy

* typo spell fix

* fix(iota-genesis-builder): add comments to the public functions

* fix(iota-genesis-builder): remove mut modificator, remove extra hashmap

* fix(iota-types): parse first address from bech32 format, csv headers support

* fix(iota-types): example of CSV file in comments above initialization function

* Minor fix

* refactor(iota-genesis-builder): move output processing out of main

* feat(iota-genesis-builder): create a dedicated run migration function for iota snapshot

* feat(iota-e2e-tests): include snapshot generation and address swap in migration test

* Make some improvements to the address swap map usage (#4376)

* Make some improvements to the address swap map usage

* fix(iota-e2e-tests): update import

* fix typo

---------

Co-authored-by: Dkwcs <pavlo.botnar@gmail.com>
Co-authored-by: DaughterOfMars <chloedaughterofmars@gmail.com>

crates/iota-e2e-tests/tests/full_node_migration_tests.rs

@@ -1,11 +1,24 @@
 // Copyright (c) 2024 IOTA Stiftung
 // SPDX-License-Identifier: Apache-2.0
 
-use std::{path::PathBuf, str::FromStr};
+use std::{
+    fs::File,
+    io::{BufWriter, Write},
+    path::PathBuf,
+    str::FromStr,
+};
 
 use anyhow::anyhow;
 use bip32::DerivationPath;
-use iota_genesis_builder::SnapshotSource;
+use iota_genesis_builder::{
+    SnapshotSource,
+    stardust::{
+        migration::{Migration, MigrationTargetNetwork},
+        parse::HornetSnapshotParser,
+        process_outputs::scale_amount_for_iota,
+        types::address_swap_map::AddressSwapMap,
+    },
+};
 use iota_json_rpc_types::{
     IotaObjectDataFilter, IotaObjectDataOptions, IotaObjectResponseQuery,
     IotaTransactionBlockResponse, IotaTransactionBlockResponseOptions,
@@ -21,15 +34,18 @@ use iota_types::{
     gas_coin::GAS,
     programmable_transaction_builder::ProgrammableTransactionBuilder,
     quorum_driver_types::ExecuteTransactionRequestType,
-    stardust::output::NftOutput,
+    stardust::{coin_type::CoinType, output::NftOutput},
     transaction::{Argument, ObjectArg, Transaction, TransactionData},
 };
 use move_core_types::ident_str;
 use shared_crypto::intent::Intent;
 use tempfile::tempdir;
 use test_cluster::TestClusterBuilder;
 
-const MIGRATION_DATA_PATH: &str = "tests/migration/stardust_object_snapshot.bin";
+const HORNET_SNAPSHOT_PATH: &str = "tests/migration/test_hornet_full_snapshot.bin";
+const ADDRESS_SWAP_MAP_PATH: &str = "tests/migration/address_swap.csv";
+const TEST_TARGET_NETWORK: &str = "alphanet-test";
+const MIGRATION_DATA_FILE_NAME: &str = "stardust_object_snapshot.bin";
 
 /// Got from iota-genesis-builder/src/stardust/test_outputs/alias_ownership.rs
 const MAIN_ADDRESS_MNEMONIC: &str = "few hood high omit camp keep burger give happy iron evolve draft few dawn pulp jazz box dash load snake gown bag draft car";
@@ -39,26 +55,67 @@ const SPONSOR_ADDRESS_MNEMONIC: &str = "okay pottery arch air egg very cave cash
 #[sim_test]
 async fn test_full_node_load_migration_data() -> Result<(), anyhow::Error> {
     telemetry_subscribers::init_for_testing();
-    let snapshot_source = SnapshotSource::Local(PathBuf::from_str(MIGRATION_DATA_PATH).unwrap());
+
+    // Setup the temporary dir and create the writer for the stardust object
+    // snapshot
+    let dir = tempdir()?;
+    let stardudst_object_snapshot_file_path = dir.path().join(MIGRATION_DATA_FILE_NAME);
+    let object_snapshot_writer =
+        BufWriter::new(File::create(&stardudst_object_snapshot_file_path)?);
+
+    // Generate the stardust object snapshot
+    genesis_builder_snapshot_generation(object_snapshot_writer)?;
+    // Then load it
+    let snapshot_source = SnapshotSource::Local(stardudst_object_snapshot_file_path);
+
+    // A new test cluster can be spawn with the stardust object snapshot
     let test_cluster = TestClusterBuilder::new()
         .with_migration_data(vec![snapshot_source])
         .build()
         .await;
 
+    // Use a client to issue a test transaction
     let client = test_cluster.wallet.get_client().await.unwrap();
-
     let tx_response = address_unlock_condition(client).await?;
-
     let IotaTransactionBlockResponse {
         confirmed_local_execution,
         errors,
         ..
     } = tx_response;
+
+    // The transaction must be successful
     assert!(confirmed_local_execution.unwrap());
     assert!(errors.is_empty());
     Ok(())
 }
 
+fn genesis_builder_snapshot_generation(
+    object_snapshot_writer: impl Write,
+) -> Result<(), anyhow::Error> {
+    let mut snapshot_parser =
+        HornetSnapshotParser::new::<false>(File::open(HORNET_SNAPSHOT_PATH)?)?;
+    let total_supply = scale_amount_for_iota(snapshot_parser.total_supply()?)?;
+    let target_network = MigrationTargetNetwork::from_str(TEST_TARGET_NETWORK)?;
+    let coin_type = CoinType::Iota;
+    let address_swap_map = AddressSwapMap::from_csv(ADDRESS_SWAP_MAP_PATH)?;
+
+    // Migrate using the parser output stream
+    Migration::new(
+        snapshot_parser.target_milestone_timestamp(),
+        total_supply,
+        target_network,
+        coin_type,
+        address_swap_map,
+    )?
+    .run_for_iota(
+        snapshot_parser.target_milestone_timestamp(),
+        snapshot_parser.outputs(),
+        object_snapshot_writer,
+    )?;
+
+    Ok(())
+}
+
 async fn address_unlock_condition(
     iota_client: IotaClient,
 ) -> Result<IotaTransactionBlockResponse, anyhow::Error> {

crates/iota-e2e-tests/tests/migration/address_swap.csv

@@ -0,0 +1,2 @@
+Origin,Destination
+iota1qp8h9augeh6tk3uvlxqfapuwv93atv63eqkpru029p6sgvr49eufyz7katr,0x4f72f788cdf4bb478cf9809e878e6163d5b351c82c11f1ea28750430752e7892

crates/iota-genesis-builder/src/main.rs

@@ -4,26 +4,20 @@
 //! Creating a stardust objects snapshot out of a Hornet snapshot.
 //! TIP that defines the Hornet snapshot file format:
 //! https://github.com/iotaledger/tips/blob/main/tips/TIP-0035/tip-0035.md
-use std::{collections::BTreeMap, fs::File, io::BufWriter};
+use std::{fs::File, io::BufWriter};
 
-use anyhow::{Result, anyhow};
+use anyhow::Result;
 use clap::{Parser, Subcommand};
 use iota_genesis_builder::{
     OBJECT_SNAPSHOT_FILE_PATH,
     stardust::{
         migration::{Migration, MigrationTargetNetwork},
         parse::HornetSnapshotParser,
-        types::{address_swap_map::AddressSwapMap, output_header::OutputHeader},
+        process_outputs::scale_amount_for_iota,
+        types::address_swap_map::AddressSwapMap,
     },
 };
-use iota_sdk::types::block::{
-    address::Address,
-    output::{
-        AliasOutputBuilder, BasicOutputBuilder, FoundryOutputBuilder, NftOutputBuilder, Output,
-        unlock_condition::{AddressUnlockCondition, StorageDepositReturnUnlockCondition},
-    },
-};
-use iota_types::{stardust::coin_type::CoinType, timelock::timelock::is_vested_reward};
+use iota_types::stardust::coin_type::CoinType;
 use tracing::Level;
 use tracing_subscriber::FmtSubscriber;
 
@@ -104,182 +98,13 @@ fn main() -> Result<()> {
 
     match coin_type {
         CoinType::Iota => {
-            struct MergingIterator<I> {
-                unlocked_address_balances: BTreeMap<Address, OutputHeaderWithBalance>,
-                snapshot_timestamp_s: u32,
-                outputs: I,
-            }
-
-            impl<I> MergingIterator<I> {
-                fn new(snapshot_timestamp_s: u32, outputs: I) -> Self {
-                    Self {
-                        unlocked_address_balances: Default::default(),
-                        snapshot_timestamp_s,
-                        outputs,
-                    }
-                }
-            }
-
-            impl<I: Iterator<Item = Result<(OutputHeader, Output)>>> Iterator for MergingIterator<I> {
-                type Item = I::Item;
-
-                fn next(&mut self) -> Option<Self::Item> {
-                    // First process all the outputs, building the unlocked_address_balances map as
-                    // we go.
-                    for res in self.outputs.by_ref() {
-                        if let Ok((header, output)) = res {
-                            fn mergeable_address(
-                                header: &OutputHeader,
-                                output: &Output,
-                                snapshot_timestamp_s: u32,
-                            ) -> Option<Address> {
-                                // ignore all non-basic outputs and non vesting outputs
-                                if !output.is_basic()
-                                    || !is_vested_reward(header.output_id(), output.as_basic())
-                                {
-                                    return None;
-                                }
-
-                                if let Some(unlock_conditions) = output.unlock_conditions() {
-                                    // check if vesting unlock period is already done
-                                    if unlock_conditions.is_time_locked(snapshot_timestamp_s) {
-                                        return None;
-                                    }
-                                    unlock_conditions.address().map(|uc| *uc.address())
-                                } else {
-                                    None
-                                }
-                            }
-
-                            if let Some(address) =
-                                mergeable_address(&header, &output, self.snapshot_timestamp_s)
-                            {
-                                // collect the unlocked vesting balances
-                                self.unlocked_address_balances
-                                    .entry(address)
-                                    .and_modify(|x| x.balance += output.amount())
-                                    .or_insert(OutputHeaderWithBalance {
-                                        output_header: header,
-                                        balance: output.amount(),
-                                    });
-                                continue;
-                            } else {
-                                return Some(Ok((header, output)));
-                            }
-                        } else {
-                            return Some(res);
-                        }
-                    }
-
-                    // Now that we are out
-                    self.unlocked_address_balances.pop_first().map(
-                        |(address, output_header_with_balance)| {
-                            // create a new basic output which holds the aggregated balance from
-                            // unlocked vesting outputs for this address
-                            let basic = BasicOutputBuilder::new_with_amount(
-                                output_header_with_balance.balance,
-                            )
-                            .add_unlock_condition(AddressUnlockCondition::new(address))
-                            .finish()
-                            .expect("should be able to create a basic output");
-
-                            Ok((output_header_with_balance.output_header, basic.into()))
-                        },
-                    )
-                }
-            }
-
-            let merged_outputs = MergingIterator::new(
+            migration.run_for_iota(
                 snapshot_parser.target_milestone_timestamp(),
                 snapshot_parser.outputs(),
-            )
-            .map(|res| {
-                let (header, mut output) = res?;
-                scale_output_amount_for_iota(&mut output)?;
-
-                Ok::<_, anyhow::Error>((header, output))
-            });
-            itertools::process_results(merged_outputs, |outputs| {
-                migration.run(outputs, object_snapshot_writer)
-            })??;
+                object_snapshot_writer,
+            )?;
         }
     }
 
     Ok(())
 }
-
-struct OutputHeaderWithBalance {
-    output_header: OutputHeader,
-    balance: u64,
-}
-
-fn scale_output_amount_for_iota(output: &mut Output) -> Result<()> {
-    *output = match output {
-        Output::Basic(ref basic_output) => {
-            // Update amount
-            let mut builder = BasicOutputBuilder::from(basic_output)
-                .with_amount(scale_amount_for_iota(basic_output.amount())?);
-
-            // Update amount in potential storage deposit return unlock condition
-            if let Some(sdr_uc) = basic_output
-                .unlock_conditions()
-                .get(StorageDepositReturnUnlockCondition::KIND)
-            {
-                let sdr_uc = sdr_uc.as_storage_deposit_return();
-                builder = builder.replace_unlock_condition(
-                    StorageDepositReturnUnlockCondition::new(
-                        sdr_uc.return_address(),
-                        scale_amount_for_iota(sdr_uc.amount())?,
-                        u64::MAX,
-                    )
-                    .unwrap(),
-                );
-            };
-
-            Output::from(builder.finish()?)
-        }
-        Output::Alias(ref alias_output) => Output::from(
-            AliasOutputBuilder::from(alias_output)
-                .with_amount(scale_amount_for_iota(alias_output.amount())?)
-                .finish()?,
-        ),
-        Output::Foundry(ref foundry_output) => Output::from(
-            FoundryOutputBuilder::from(foundry_output)
-                .with_amount(scale_amount_for_iota(foundry_output.amount())?)
-                .finish()?,
-        ),
-        Output::Nft(ref nft_output) => {
-            // Update amount
-            let mut builder = NftOutputBuilder::from(nft_output)
-                .with_amount(scale_amount_for_iota(nft_output.amount())?);
-
-            // Update amount in potential storage deposit return unlock condition
-            if let Some(sdr_uc) = nft_output
-                .unlock_conditions()
-                .get(StorageDepositReturnUnlockCondition::KIND)
-            {
-                let sdr_uc = sdr_uc.as_storage_deposit_return();
-                builder = builder.replace_unlock_condition(
-                    StorageDepositReturnUnlockCondition::new(
-                        sdr_uc.return_address(),
-                        scale_amount_for_iota(sdr_uc.amount())?,
-                        u64::MAX,
-                    )
-                    .unwrap(),
-                );
-            };
-
-            Output::from(builder.finish()?)
-        }
-        Output::Treasury(_) => return Ok(()),
-    };
-    Ok(())
-}
-
-fn scale_amount_for_iota(amount: u64) -> Result<u64> {
-    const IOTA_MULTIPLIER: u64 = 1000;
-
-    amount
-        .checked_mul(IOTA_MULTIPLIER)
-        .ok_or_else(|| anyhow!("overflow multiplying amount {amount} by {IOTA_MULTIPLIER}"))
-}

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

@@ -32,6 +32,7 @@ use crate::stardust::{
         verification::{created_objects::CreatedObjects, verify_outputs},
     },
     native_token::package_data::NativeTokenPackageData,
+    process_outputs::get_merged_outputs_for_iota,
     types::{address_swap_map::AddressSwapMap, output_header::OutputHeader},
 };
 
@@ -163,6 +164,20 @@ impl Migration {
         Ok(())
     }
 
+    /// Run all stages of the migration coming from a Hornet snapshot with IOTA
+    /// coin type.
+    pub fn run_for_iota<'a>(
+        self,
+        target_milestone_timestamp: u32,
+        outputs: impl Iterator<Item = Result<(OutputHeader, Output)>> + 'a,
+        writer: impl Write,
+    ) -> Result<()> {
+        itertools::process_results(
+            get_merged_outputs_for_iota(target_milestone_timestamp, outputs),
+            |outputs| self.run(outputs, writer),
+        )?
+    }
+
     /// The migration objects.
     ///
     /// The system packages and underlying `init` objects

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

@@ -8,6 +8,7 @@
 pub mod migration;
 pub mod native_token;
 pub mod parse;
+pub mod process_outputs;
 #[cfg(feature = "test-outputs")]
 pub mod test_outputs;
 pub mod types;