From ce407df925bd5b92c43adeb15b490d335b36d791 Mon Sep 17 00:00:00 2001 From: miker83z Date: Wed, 15 Jan 2025 11:56:58 +0100 Subject: [PATCH] feat(iota-genesis-builder): add support for the SwapSplit --- .../tests/full_node_migration_tests.rs | 3 +- crates/iota-genesis-builder/src/main.rs | 25 +++- .../src/stardust/migration/migration.rs | 8 +- .../src/stardust/migration/tests/basic.rs | 2 +- .../src/stardust/migration/tests/mod.rs | 2 +- .../stardust/types/address_swap_split_map.rs | 110 ++++++++++++++++++ .../src/stardust/types/mod.rs | 1 + .../src/stardust/types/output_header.rs | 2 +- 8 files changed, 145 insertions(+), 8 deletions(-) create mode 100644 crates/iota-genesis-builder/src/stardust/types/address_swap_split_map.rs diff --git a/crates/iota-e2e-tests/tests/full_node_migration_tests.rs b/crates/iota-e2e-tests/tests/full_node_migration_tests.rs index 26726e5bb6b..199a772e046 100644 --- a/crates/iota-e2e-tests/tests/full_node_migration_tests.rs +++ b/crates/iota-e2e-tests/tests/full_node_migration_tests.rs @@ -16,7 +16,7 @@ use iota_genesis_builder::{ migration::{Migration, MigrationTargetNetwork}, parse::HornetSnapshotParser, process_outputs::scale_amount_for_iota, - types::address_swap_map::AddressSwapMap, + types::{address_swap_map::AddressSwapMap, address_swap_split_map::AddressSwapSplitMap}, }, }; use iota_json_rpc_types::{ @@ -111,6 +111,7 @@ fn genesis_builder_snapshot_generation( )? .run_for_iota( snapshot_parser.target_milestone_timestamp(), + AddressSwapSplitMap::default(), snapshot_parser.outputs(), object_snapshot_writer, )?; diff --git a/crates/iota-genesis-builder/src/main.rs b/crates/iota-genesis-builder/src/main.rs index b85c4358fee..066b0b0cbdf 100644 --- a/crates/iota-genesis-builder/src/main.rs +++ b/crates/iota-genesis-builder/src/main.rs @@ -14,7 +14,7 @@ use iota_genesis_builder::{ migration::{Migration, MigrationTargetNetwork}, parse::HornetSnapshotParser, process_outputs::scale_amount_for_iota, - types::address_swap_map::AddressSwapMap, + types::{address_swap_map::AddressSwapMap, address_swap_split_map::AddressSwapSplitMap}, }, }; use iota_types::stardust::coin_type::CoinType; @@ -41,6 +41,11 @@ enum Snapshot { 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: Option, + #[clap( + long, + help = "Path to the address swap split map file. This must be a CSV file with four columns, where an entry contains in the first column a (bech32) Address present in the Hornet full-snapshot, in the second column an IotaAddress that will be used for the swap, in the third column a target amount of iota tokens to be split from the origin address to the destination address and in the fourth column the amount of timelocked iota tokens used for the same scope." + )] + address_swap_split_map_path: Option, #[clap(long, value_parser = clap::value_parser!(MigrationTargetNetwork), help = "Target network for migration")] target_network: MigrationTargetNetwork, }, @@ -55,15 +60,23 @@ fn main() -> Result<()> { // Parse the CLI arguments let cli = Cli::parse(); - let (snapshot_path, address_swap_map_path, target_network, coin_type) = match cli.snapshot { + let ( + snapshot_path, + address_swap_map_path, + target_network, + address_swap_split_map_path, + coin_type, + ) = match cli.snapshot { Snapshot::Iota { snapshot_path, address_swap_map_path, + address_swap_split_map_path, target_network, } => ( snapshot_path, address_swap_map_path, target_network, + address_swap_split_map_path, CoinType::Iota, ), }; @@ -83,6 +96,13 @@ fn main() -> Result<()> { } else { AddressSwapMap::default() }; + + let address_swap_split_map = + if let Some(address_swap_split_map_path) = address_swap_split_map_path { + AddressSwapSplitMap::from_csv(&address_swap_split_map_path)? + } else { + AddressSwapSplitMap::default() + }; // Prepare the migration using the parser output stream let migration = Migration::new( snapshot_parser.target_milestone_timestamp(), @@ -100,6 +120,7 @@ fn main() -> Result<()> { CoinType::Iota => { migration.run_for_iota( snapshot_parser.target_milestone_timestamp(), + address_swap_split_map, snapshot_parser.outputs(), object_snapshot_writer, )?; diff --git a/crates/iota-genesis-builder/src/stardust/migration/migration.rs b/crates/iota-genesis-builder/src/stardust/migration/migration.rs index b7f8738adf7..d80096d99d5 100644 --- a/crates/iota-genesis-builder/src/stardust/migration/migration.rs +++ b/crates/iota-genesis-builder/src/stardust/migration/migration.rs @@ -33,7 +33,10 @@ use crate::stardust::{ }, native_token::package_data::NativeTokenPackageData, process_outputs::process_outputs_for_iota, - types::{address_swap_map::AddressSwapMap, output_header::OutputHeader}, + types::{ + address_swap_map::AddressSwapMap, address_swap_split_map::AddressSwapSplitMap, + output_header::OutputHeader, + }, }; /// We fix the protocol version used in the migration. @@ -169,11 +172,12 @@ impl Migration { pub fn run_for_iota<'a>( self, target_milestone_timestamp: u32, + swap_split_map: AddressSwapSplitMap, outputs: impl Iterator> + 'a, writer: impl Write, ) -> Result<()> { itertools::process_results( - process_outputs_for_iota(target_milestone_timestamp, outputs), + process_outputs_for_iota(target_milestone_timestamp, swap_split_map, outputs), |outputs| self.run(outputs, writer), )? } diff --git a/crates/iota-genesis-builder/src/stardust/migration/tests/basic.rs b/crates/iota-genesis-builder/src/stardust/migration/tests/basic.rs index b503be9862c..7dcddce5933 100644 --- a/crates/iota-genesis-builder/src/stardust/migration/tests/basic.rs +++ b/crates/iota-genesis-builder/src/stardust/migration/tests/basic.rs @@ -76,7 +76,7 @@ fn basic_simple_coin_id() { fn basic_simple_coin_id_with_expired_timelock() { for header in [ random_output_header(), - OutputHeader::new_testing( + OutputHeader::new( // A potential vesting reward output transaction ID. *TransactionId::from_str( "0xb191c4bc825ac6983789e50545d5ef07a1d293a98ad974fc9498cb1812345678", diff --git a/crates/iota-genesis-builder/src/stardust/migration/tests/mod.rs b/crates/iota-genesis-builder/src/stardust/migration/tests/mod.rs index f694f1b8758..97615e4f0f1 100644 --- a/crates/iota-genesis-builder/src/stardust/migration/tests/mod.rs +++ b/crates/iota-genesis-builder/src/stardust/migration/tests/mod.rs @@ -52,7 +52,7 @@ mod foundry; mod nft; fn random_output_header() -> OutputHeader { - OutputHeader::new_testing( + OutputHeader::new( random(), random_output_index(), random(), diff --git a/crates/iota-genesis-builder/src/stardust/types/address_swap_split_map.rs b/crates/iota-genesis-builder/src/stardust/types/address_swap_split_map.rs new file mode 100644 index 00000000000..5f78c872ff6 --- /dev/null +++ b/crates/iota-genesis-builder/src/stardust/types/address_swap_split_map.rs @@ -0,0 +1,110 @@ +use std::collections::HashMap; + +use iota_sdk::types::block::address::Address; +use iota_types::base_types::IotaAddress; + +type OriginAddress = Address; +type Destination = (IotaAddress, u64, u64); + +#[derive(Default)] +pub struct AddressSwapSplitMap { + addresses: HashMap, +} + +impl AddressSwapSplitMap { + /// If the `address` passed as input is present in the map, then return + /// a mutable reference to the destination, i.e., a tuple containing a + /// destination address, a tokens target and a timelocked tokens target. + pub fn get_destination_maybe_mut( + &mut self, + address: &OriginAddress, + ) -> Option<&mut (IotaAddress, u64, u64)> { + self.addresses.get_mut(address) + } + + /// Check whether the map has all targets set to 0. Return the first + /// occurrence of an entry where one or both the two targets are greater + /// than zero. If none is found, then return None. + pub fn validate_successfull_swap_split( + &self, + ) -> Option<(&OriginAddress, &IotaAddress, u64, u64)> { + for (origin, (destination, tokens_target, tokens_timelocked_target)) in + self.addresses.iter() + { + if *tokens_target > 0 || *tokens_timelocked_target > 0 { + return Some(( + origin, + destination, + *tokens_target, + *tokens_timelocked_target, + )); + } + } + None + } + + /// Initializes an [`AddressSwapSplitMap`] by reading address pairs from a + /// CSV file. + /// + /// The function expects the file to contain four columns: the origin + /// address (first column), the destination address (second column), the + /// tokens target (third column) and the timelocked tokens target + /// (fourth column). These are parsed into a [`HashMap`] that maps + /// origin addresses to tuples containing the destination address and + /// the two targets. + /// + /// # Example CSV File + /// ```csv + /// Origin,Destination,Tokens,TokensTimelocked + /// iota1qrukjnd6jhgwc0ls6dgt574sxuulcsmq5lnzhtv4jmlwkydhe2zvy69t7jj,0x1336d143de5eb55bcb069f55da5fc9f0c84e368022fd2bbe0125b1093b446313,107667149000,107667149000 + /// iota1qr4chj9jwhauvegqy40sdhj93mzmvc3mg9cmzlv2y6j8vpyxpvug2y6h5jd,0x83b5ed87bac715ecb09017a72d531ccc3c43bcb58edeb1ce383f1c46cfd79bec,388647312000,0 + + /// ``` + /// + /// # Parameters + /// - `file_path`: The relative path to the CSV file containing the address + /// mappings. + /// + /// # Returns + /// - An [`AddressSwapSplitMap`] containing the parsed mappings. + /// + /// # Errors + /// - Returns an error if the file cannot be found, read, or parsed + /// correctly. + /// - Returns an error if the origin, destination addresses, or targets + /// cannot be parsed into. + pub fn from_csv(file_path: &str) -> Result { + let current_dir = std::env::current_dir()?; + let file_path = current_dir.join(file_path); + let mut reader = csv::ReaderBuilder::new().from_path(file_path)?; + let mut addresses = HashMap::new(); + + verify_headers(reader.headers()?)?; + + for result in reader.records() { + let record = result?; + let origin = OriginAddress::try_from_bech32(&record[0])?; + let destination_address = record[1].parse()?; + let tokens_target = record[2].parse()?; + let tokens_timelocked_target = record[3].parse()?; + addresses.insert( + origin, + (destination_address, tokens_target, tokens_timelocked_target), + ); + } + + Ok(AddressSwapSplitMap { addresses }) + } +} + +fn verify_headers(headers: &csv::StringRecord) -> Result<(), anyhow::Error> { + if headers.len() != 4 + || &headers[0] != "Origin" + || &headers[1] != "Destination" + || &headers[2] != "Tokens" + || &headers[3] != "TokensTimelocked" + { + anyhow::bail!("Invalid CSV headers"); + } + Ok(()) +} diff --git a/crates/iota-genesis-builder/src/stardust/types/mod.rs b/crates/iota-genesis-builder/src/stardust/types/mod.rs index 49ea9368cf5..89cf2b9264a 100644 --- a/crates/iota-genesis-builder/src/stardust/types/mod.rs +++ b/crates/iota-genesis-builder/src/stardust/types/mod.rs @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 pub mod address_swap_map; +pub mod address_swap_split_map; pub mod output_header; pub mod output_index; pub mod snapshot; diff --git a/crates/iota-genesis-builder/src/stardust/types/output_header.rs b/crates/iota-genesis-builder/src/stardust/types/output_header.rs index ae6de6eb270..e1dd2fbd2f6 100644 --- a/crates/iota-genesis-builder/src/stardust/types/output_header.rs +++ b/crates/iota-genesis-builder/src/stardust/types/output_header.rs @@ -57,7 +57,7 @@ impl OutputHeader { } /// Creates a new OutputHeader for testing. - pub fn new_testing( + pub fn new( transaction_id_bytes: [u8; 32], output_index: OutputIndex, block_id_bytes: [u8; 32],