feat(iota-genesis-builder): add address swap map for swapping origin …

…addresses to destination during the migration process (#4314)

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

---------

Co-authored-by: DaughterOfMars <[email protected]>

crates/iota-genesis-builder/Cargo.toml

@@ -16,6 +16,7 @@ bcs.workspace = true
 bigdecimal = "0.4"
 camino.workspace = true
 clap.workspace = true
+csv = "1.2"
 fastcrypto.workspace = true
 flate2.workspace = true
 fs_extra = "1.3"

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

@@ -13,7 +13,7 @@ use iota_genesis_builder::{
     stardust::{
         migration::{Migration, MigrationTargetNetwork},
         parse::HornetSnapshotParser,
-        types::output_header::OutputHeader,
+        types::{address_swap_map::AddressSwapMap, output_header::OutputHeader},
     },
 };
 use iota_sdk::types::block::{
@@ -42,6 +42,11 @@ enum Snapshot {
     Iota {
         #[clap(long, help = "Path to the Iota Hornet full-snapshot file")]
         snapshot_path: String,
+        #[clap(
+            long,
+            help = "Path to the address swap map file. This must be a CSV file with two columns, where an entry contains in the first column an IotaAddress present in the Hornet full-snapshot and in the second column an IotaAddress that will be used for the swap."
+        )]
+        address_swap_map_path: String,
         #[clap(long, value_parser = clap::value_parser!(MigrationTargetNetwork), help = "Target network for migration")]
         target_network: MigrationTargetNetwork,
     },
@@ -56,11 +61,17 @@ fn main() -> Result<()> {
 
     // Parse the CLI arguments
     let cli = Cli::parse();
-    let (snapshot_path, target_network, coin_type) = match cli.snapshot {
+    let (snapshot_path, address_swap_map_path, target_network, coin_type) = match cli.snapshot {
         Snapshot::Iota {
             snapshot_path,
+            address_swap_map_path,
             target_network,
-        } => (snapshot_path, target_network, CoinType::Iota),
+        } => (
+            snapshot_path,
+            address_swap_map_path,
+            target_network,
+            CoinType::Iota,
+        ),
     };
 
     // Start the Hornet snapshot parser
@@ -73,12 +84,14 @@ fn main() -> Result<()> {
         CoinType::Iota => scale_amount_for_iota(snapshot_parser.total_supply()?)?,
     };
 
+    let address_swap_map = AddressSwapMap::from_csv(&address_swap_map_path)?;
     // Prepare the migration using the parser output stream
     let migration = Migration::new(
         snapshot_parser.target_milestone_timestamp(),
         total_supply,
         target_network,
         coin_type,
+        address_swap_map,
     )?;
 
     // Prepare the writer for the objects snapshot

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

@@ -35,7 +35,6 @@ use iota_types::{
     stardust::{
         coin_type::CoinType,
         output::{Nft, foundry::create_foundry_amount_coin},
-        stardust_to_iota_address, stardust_to_iota_address_owner,
     },
     timelock::timelock,
     transaction::{
@@ -53,7 +52,10 @@ use crate::{
             MigrationTargetNetwork, PACKAGE_DEPS, create_migration_context, package_module_bytes,
             verification::created_objects::CreatedObjects,
         },
-        types::{output_header::OutputHeader, token_scheme::SimpleTokenSchemeU64},
+        types::{
+            address_swap_map::AddressSwapMap, output_header::OutputHeader,
+            token_scheme::SimpleTokenSchemeU64,
+        },
     },
 };
 
@@ -307,6 +309,7 @@ impl Executor {
         header: &OutputHeader,
         alias: &AliasOutput,
         coin_type: CoinType,
+        address_swap_map: &mut AddressSwapMap,
     ) -> Result<CreatedObjects> {
         let mut created_objects = CreatedObjects::default();
 
@@ -316,7 +319,8 @@ impl Executor {
         let move_alias = iota_types::stardust::output::Alias::try_from_stardust(alias_id, alias)?;
 
         // TODO: We should ensure that no circular ownership exists.
-        let alias_output_owner = stardust_to_iota_address_owner(alias.governor_address())?;
+        let alias_output_owner =
+            address_swap_map.swap_stardust_to_iota_address_owner(alias.governor_address())?;
 
         let package_deps = InputObjects::new(self.load_packages(PACKAGE_DEPS).collect());
         let version = package_deps.lamport_timestamp(&[]);
@@ -558,10 +562,14 @@ impl Executor {
         basic_output: &BasicOutput,
         target_milestone_timestamp_sec: u32,
         coin_type: &CoinType,
+        address_swap_map: &mut AddressSwapMap,
     ) -> Result<CreatedObjects> {
         let mut basic =
             iota_types::stardust::output::BasicOutput::new(header.new_object_id(), basic_output)?;
-        let owner: IotaAddress = basic_output.address().to_string().parse()?;
+
+        let basic_objects_owner =
+            address_swap_map.swap_stardust_to_iota_address(basic_output.address())?;
+
         let mut created_objects = CreatedObjects::default();
 
         // The minimum version of the manually created objects
@@ -570,11 +578,12 @@ impl Executor {
 
         let object = if basic.is_simple_coin(target_milestone_timestamp_sec) {
             if !basic_output.native_tokens().is_empty() {
-                let coins = self.create_native_token_coins(basic_output.native_tokens(), owner)?;
+                let coins = self
+                    .create_native_token_coins(basic_output.native_tokens(), basic_objects_owner)?;
                 created_objects.set_native_tokens(coins)?;
             }
             let amount_coin = basic.into_genesis_coin_object(
-                owner,
+                basic_objects_owner,
                 &self.protocol_config,
                 &self.tx_context,
                 version,
@@ -596,7 +605,7 @@ impl Executor {
                 basic.native_tokens.id = UID::new(self.tx_context.fresh_id());
             }
             let object = basic.to_genesis_object(
-                owner,
+                basic_objects_owner,
                 &self.protocol_config,
                 &self.tx_context,
                 version,
@@ -618,10 +627,12 @@ impl Executor {
         output_id: OutputId,
         basic_output: &BasicOutput,
         target_milestone_timestamp: u32,
+        address_swap_map: &mut AddressSwapMap,
     ) -> Result<CreatedObjects> {
         let mut created_objects = CreatedObjects::default();
 
-        let owner: IotaAddress = basic_output.address().to_string().parse()?;
+        let basic_output_owner =
+            address_swap_map.swap_stardust_to_iota_address(basic_output.address())?;
 
         let package_deps = InputObjects::new(self.load_packages(PACKAGE_DEPS).collect());
         let version = package_deps.lamport_timestamp(&[]);
@@ -631,7 +642,7 @@ impl Executor {
 
         let object = timelock::to_genesis_object(
             timelock,
-            owner,
+            basic_output_owner,
             &self.protocol_config,
             &self.tx_context,
             version,
@@ -648,6 +659,7 @@ impl Executor {
         header: &OutputHeader,
         nft: &NftOutput,
         coin_type: CoinType,
+        address_swap_map: &mut AddressSwapMap,
     ) -> Result<CreatedObjects> {
         let mut created_objects = CreatedObjects::default();
 
@@ -657,8 +669,11 @@ impl Executor {
         let move_nft = Nft::try_from_stardust(nft_id, nft)?;
 
         // TODO: We should ensure that no circular ownership exists.
-        let nft_output_owner_address = stardust_to_iota_address(nft.address())?;
-        let nft_output_owner = stardust_to_iota_address_owner(nft.address())?;
+        let nft_output_owner_address =
+            address_swap_map.swap_stardust_to_iota_address(nft.address())?;
+
+        let nft_output_owner =
+            address_swap_map.swap_stardust_to_iota_address_owner(nft.address())?;
 
         let package_deps = InputObjects::new(self.load_packages(PACKAGE_DEPS).collect());
         let version = package_deps.lamport_timestamp(&[]);

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

@@ -32,7 +32,7 @@ use crate::stardust::{
         verification::{created_objects::CreatedObjects, verify_outputs},
     },
     native_token::package_data::NativeTokenPackageData,
-    types::output_header::OutputHeader,
+    types::{address_swap_map::AddressSwapMap, output_header::OutputHeader},
 };
 
 /// We fix the protocol version used in the migration.
@@ -74,6 +74,7 @@ pub struct Migration {
     /// The coin type to use in order to migrate outputs. Can only be equal to
     /// `Iota` at the moment. Is fixed for the entire migration process.
     coin_type: CoinType,
+    address_swap_map: AddressSwapMap,
 }
 
 impl Migration {
@@ -84,6 +85,7 @@ impl Migration {
         total_supply: u64,
         target_network: MigrationTargetNetwork,
         coin_type: CoinType,
+        address_swap_map: AddressSwapMap,
     ) -> Result<Self> {
         let executor = Executor::new(
             ProtocolVersion::new(MIGRATION_PROTOCOL_VERSION),
@@ -96,6 +98,7 @@ impl Migration {
             executor,
             output_objects_map: Default::default(),
             coin_type,
+            address_swap_map,
         })
     }
 
@@ -135,7 +138,7 @@ impl Migration {
             .collect::<Vec<_>>();
         info!("Verifying ledger state...");
         self.verify_ledger_state(&outputs)?;
-
+        self.address_swap_map.verify_all_addresses_swapped()?;
         Ok(())
     }
 
@@ -194,14 +197,18 @@ impl Migration {
     ) -> Result<()> {
         for (header, output) in outputs {
             let created = match output {
-                Output::Alias(alias) => {
-                    self.executor
-                        .create_alias_objects(header, alias, self.coin_type)?
-                }
-                Output::Nft(nft) => {
-                    self.executor
-                        .create_nft_objects(header, nft, self.coin_type)?
-                }
+                Output::Alias(alias) => self.executor.create_alias_objects(
+                    header,
+                    alias,
+                    self.coin_type,
+                    &mut self.address_swap_map,
+                )?,
+                Output::Nft(nft) => self.executor.create_nft_objects(
+                    header,
+                    nft,
+                    self.coin_type,
+                    &mut self.address_swap_map,
+                )?,
                 Output::Basic(basic) => {
                     // All timelocked vested rewards(basic outputs with the specific ID format)
                     // should be migrated as TimeLock<Balance<IOTA>> objects.
@@ -214,13 +221,15 @@ impl Migration {
                             header.output_id(),
                             basic,
                             self.target_milestone_timestamp_sec,
+                            &mut self.address_swap_map,
                         )?
                     } else {
                         self.executor.create_basic_objects(
                             header,
                             basic,
                             self.target_milestone_timestamp_sec,
                             &self.coin_type,
+                            &mut self.address_swap_map,
                         )?
                     }
                 }
@@ -244,6 +253,7 @@ impl Migration {
             self.target_milestone_timestamp_sec,
             self.total_supply,
             self.executor.store(),
+            &self.address_swap_map,
         )?;
         Ok(())
     }

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

@@ -32,7 +32,7 @@ use crate::stardust::{
             random_output_header, unlock_object,
         },
     },
-    types::output_header::OutputHeader,
+    types::{address_swap_map::AddressSwapMap, output_header::OutputHeader},
 };
 
 /// Test the id of a `BasicOutput` that is transformed to a simple coin.
@@ -53,6 +53,7 @@ fn basic_simple_coin_id() {
         1_000_000,
         MigrationTargetNetwork::Mainnet,
         CoinType::Iota,
+        AddressSwapMap::default(),
     )
     .unwrap();
     migration
@@ -103,6 +104,7 @@ fn basic_simple_coin_id_with_expired_timelock() {
             1_000_000,
             MigrationTargetNetwork::Mainnet,
             CoinType::Iota,
+            AddressSwapMap::default(),
         )
         .unwrap();
         migration
@@ -139,6 +141,7 @@ fn basic_id() {
         1_000_000,
         MigrationTargetNetwork::Mainnet,
         CoinType::Iota,
+        AddressSwapMap::default(),
     )
     .unwrap();
     migration
@@ -183,6 +186,7 @@ fn basic_simple_coin_migration_with_native_token() {
         1_000_000,
         MigrationTargetNetwork::Mainnet,
         CoinType::Iota,
+        AddressSwapMap::default(),
     )
     .unwrap();
     migration.run_migration(outputs).unwrap();
@@ -225,6 +229,7 @@ fn basic_simple_coin_migration_with_native_tokens() {
         1_000_000,
         MigrationTargetNetwork::Mainnet,
         CoinType::Iota,
+        AddressSwapMap::default(),
     )
     .unwrap();
     migration.run_migration(outputs.clone()).unwrap();

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

@@ -39,7 +39,10 @@ use crate::stardust::{
         },
         verification::created_objects::CreatedObjects,
     },
-    types::{output_header::OutputHeader, output_index::random_output_index},
+    types::{
+        address_swap_map::AddressSwapMap, output_header::OutputHeader,
+        output_index::random_output_index,
+    },
 };
 
 mod alias;
@@ -63,8 +66,13 @@ fn run_migration(
     outputs: impl IntoIterator<Item = (OutputHeader, Output)>,
     coin_type: CoinType,
 ) -> anyhow::Result<(Executor, HashMap<OutputId, CreatedObjects>)> {
-    let mut migration =
-        Migration::new(1, total_supply, MigrationTargetNetwork::Mainnet, coin_type)?;
+    let mut migration = Migration::new(
+        1,
+        total_supply,
+        MigrationTargetNetwork::Mainnet,
+        coin_type,
+        AddressSwapMap::default(),
+    )?;
     migration.run_migration(outputs)?;
     Ok(migration.into_parts())
 }