feat(iota-genesis-builder): Allow Shimmer test outputs generation

crates/iota-genesis-builder/examples/snapshot_add_test_outputs.rs

@@ -5,13 +5,38 @@
 
 use std::{fs::File, path::Path};
 
+use clap::{Parser, Subcommand};
 use iota_genesis_builder::stardust::{
     parse::HornetSnapshotParser,
-    test_outputs::{add_snapshot_test_outputs, to_micros},
+    test_outputs::{add_snapshot_test_outputs, to_micros, STARDUST_TOTAL_SUPPLY_SHIMMER_MICRO},
 };
-use iota_types::gas_coin::STARDUST_TOTAL_SUPPLY_IOTA;
+use iota_types::{gas_coin::STARDUST_TOTAL_SUPPLY_IOTA, stardust::coin_type::CoinType};
 
-fn parse_snapshot<const VERIFY: bool>(path: impl AsRef<Path>) -> anyhow::Result<()> {
+#[derive(Parser, Debug)]
+#[clap(about = "Tool for adding test data to Iota and Shimmer Hornet full-snapshots")]
+struct Cli {
+    #[clap(subcommand)]
+    snapshot: Snapshot,
+}
+
+#[derive(Subcommand, Debug)]
+enum Snapshot {
+    #[clap(about = "Parse an Iota Hornet full-snapshot file")]
+    Iota {
+        #[clap(long, help = "Path to the Iota Hornet full-snapshot file")]
+        snapshot_path: String,
+    },
+    #[clap(about = "Parse a Shimmer Hornet full-snapshot file")]
+    Shimmer {
+        #[clap(long, help = "Path to the Shimmer Hornet full-snapshot file")]
+        snapshot_path: String,
+    },
+}
+
+fn parse_snapshot<const VERIFY: bool>(
+    path: impl AsRef<Path>,
+    coin_type: CoinType,
+) -> anyhow::Result<()> {
     let file = File::open(path)?;
     let mut parser = HornetSnapshotParser::new::<VERIFY>(file)?;
 
@@ -21,8 +46,12 @@ fn parse_snapshot<const VERIFY: bool>(path: impl AsRef<Path>) -> anyhow::Result<
         Ok::<_, anyhow::Error>(acc + output?.1.amount())
     })?;
 
-    // Total supply is in IOTA, snapshot supply is Micros
-    assert_eq!(total_supply, to_micros(STARDUST_TOTAL_SUPPLY_IOTA));
+    let expected_total_supply = match coin_type {
+        CoinType::Iota => to_micros(STARDUST_TOTAL_SUPPLY_IOTA),
+        CoinType::Shimmer => STARDUST_TOTAL_SUPPLY_SHIMMER_MICRO,
+    };
+
+    assert_eq!(total_supply, expected_total_supply);
 
     println!("Total supply: {total_supply}");
 
@@ -31,8 +60,10 @@ fn parse_snapshot<const VERIFY: bool>(path: impl AsRef<Path>) -> anyhow::Result<
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let Some(current_path) = std::env::args().nth(1) else {
-        anyhow::bail!("please provide path to the full-snapshot file");
+    let cli = Cli::parse();
+    let (current_path, coin_type) = match cli.snapshot {
+        Snapshot::Iota { snapshot_path } => (snapshot_path, CoinType::Iota),
+        Snapshot::Shimmer { snapshot_path } => (snapshot_path, CoinType::Shimmer),
     };
     let mut new_path = String::from("test-");
     // prepend "test-" before the file name
@@ -44,13 +75,23 @@ async fn main() -> anyhow::Result<()> {
         new_path.push_str(&current_path);
     }
 
-    parse_snapshot::<false>(&current_path)?;
+    parse_snapshot::<false>(&current_path, coin_type)?;
 
-    let randomness_seed = 0;
-    add_snapshot_test_outputs::<false>(&current_path, &new_path, randomness_seed, None, false)
-        .await?;
+    let randomness_seed = match coin_type {
+        CoinType::Iota => 0,
+        CoinType::Shimmer => 1,
+    };
+    add_snapshot_test_outputs::<false>(
+        &current_path,
+        &new_path,
+        coin_type,
+        randomness_seed,
+        None,
+        false,
+    )
+    .await?;
 
-    parse_snapshot::<false>(&new_path)?;
+    parse_snapshot::<false>(&current_path, coin_type)?;
 
     Ok(())
 }

crates/iota-genesis-builder/examples/snapshot_only_test_outputs.rs

@@ -6,19 +6,47 @@
 
 use std::{fs::File, path::Path};
 
+use clap::{Parser, Subcommand};
 use iota_genesis_builder::{
     stardust::{
         parse::HornetSnapshotParser,
-        test_outputs::{add_snapshot_test_outputs, to_micros},
+        test_outputs::{add_snapshot_test_outputs, to_micros, STARDUST_TOTAL_SUPPLY_SHIMMER_MICRO},
     },
     IF_STARDUST_ADDRESS,
 };
 use iota_sdk::types::block::address::Address;
-use iota_types::gas_coin::STARDUST_TOTAL_SUPPLY_IOTA;
+use iota_types::{gas_coin::STARDUST_TOTAL_SUPPLY_IOTA, stardust::coin_type::CoinType};
+
+pub const IF_SHIMMER_STARDUST_ADDRESS: &str =
+    "smr1qqzhsp9x3m22l55wlclawlw25536e3rgghep77awcfrgh60uxhxuq6vlak7";
 
 const WITH_SAMPLING: bool = false;
 
-fn parse_snapshot<const VERIFY: bool>(path: impl AsRef<Path>) -> anyhow::Result<()> {
+#[derive(Parser, Debug)]
+#[clap(about = "Tool for generating Iota and Shimmer Hornet full-snapshot files with test data")]
+struct Cli {
+    #[clap(subcommand)]
+    snapshot: Snapshot,
+}
+
+#[derive(Subcommand, Debug)]
+enum Snapshot {
+    #[clap(about = "Parse an Iota Hornet full-snapshot file")]
+    Iota {
+        #[clap(long, help = "Path to the Iota Hornet full-snapshot file")]
+        snapshot_path: String,
+    },
+    #[clap(about = "Parse a Shimmer Hornet full-snapshot file")]
+    Shimmer {
+        #[clap(long, help = "Path to the Shimmer Hornet full-snapshot file")]
+        snapshot_path: String,
+    },
+}
+
+fn parse_snapshot<const VERIFY: bool>(
+    path: impl AsRef<Path>,
+    coin_type: CoinType,
+) -> anyhow::Result<()> {
     let file = File::open(path)?;
     let mut parser = HornetSnapshotParser::new::<VERIFY>(file)?;
 
@@ -28,8 +56,12 @@ fn parse_snapshot<const VERIFY: bool>(path: impl AsRef<Path>) -> anyhow::Result<
         Ok::<_, anyhow::Error>(acc + output?.1.amount())
     })?;
 
-    // Total supply is in IOTA, snapshot supply is Micros
-    assert_eq!(total_supply, to_micros(STARDUST_TOTAL_SUPPLY_IOTA));
+    let expected_total_supply = match coin_type {
+        CoinType::Iota => to_micros(STARDUST_TOTAL_SUPPLY_IOTA),
+        CoinType::Shimmer => STARDUST_TOTAL_SUPPLY_SHIMMER_MICRO,
+    };
+
+    assert_eq!(total_supply, expected_total_supply);
 
     println!("Total supply: {total_supply}");
 
@@ -38,8 +70,10 @@ fn parse_snapshot<const VERIFY: bool>(path: impl AsRef<Path>) -> anyhow::Result<
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
-    let Some(current_path) = std::env::args().nth(1) else {
-        anyhow::bail!("please provide path to the full-snapshot file");
+    let cli = Cli::parse();
+    let (current_path, coin_type) = match cli.snapshot {
+        Snapshot::Iota { snapshot_path } => (snapshot_path, CoinType::Iota),
+        Snapshot::Shimmer { snapshot_path } => (snapshot_path, CoinType::Shimmer),
     };
     let mut new_path = String::from("test-");
     // prepend "test-" before the file name
@@ -51,19 +85,30 @@ async fn main() -> anyhow::Result<()> {
         new_path.push_str(&current_path);
     }
 
-    parse_snapshot::<false>(&current_path)?;
+    parse_snapshot::<false>(&current_path, coin_type)?;
+
+    let (randomness_seed, delegator_address) = match coin_type {
+        CoinType::Iota => {
+            // IOTA coin type values
+            (0, IF_STARDUST_ADDRESS)
+        }
+        CoinType::Shimmer => {
+            // Shimmer coin type values
+            (1, IF_SHIMMER_STARDUST_ADDRESS)
+        }
+    };
 
-    let randomness_seed = 0;
     add_snapshot_test_outputs::<false>(
         &current_path,
         &new_path,
+        coin_type,
         randomness_seed,
-        *Address::try_from_bech32(IF_STARDUST_ADDRESS)?.as_ed25519(),
+        *Address::try_from_bech32(delegator_address)?.as_ed25519(),
         WITH_SAMPLING,
     )
     .await?;
 
-    parse_snapshot::<false>(&new_path)?;
+    parse_snapshot::<false>(&new_path, coin_type)?;
 
     Ok(())
 }

crates/iota-genesis-builder/src/stardust/test_outputs/alias_ownership.rs

@@ -27,15 +27,17 @@ use crate::stardust::{
 };
 
 const 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";
-const COIN_TYPE: u32 = 4218;
 const OWNING_ALIAS_COUNT: u32 = 10;
 
-pub(crate) async fn outputs(rng: &mut StdRng) -> anyhow::Result<Vec<(OutputHeader, Output)>> {
+pub(crate) async fn outputs(
+    rng: &mut StdRng,
+    coin_type: u32,
+) -> anyhow::Result<Vec<(OutputHeader, Output)>> {
     let mut outputs = Vec::new();
     let secret_manager = MnemonicSecretManager::try_from_mnemonic(MNEMONIC)?;
 
     let alias_owners = secret_manager
-        .generate_ed25519_addresses(COIN_TYPE, 0, 0..OWNING_ALIAS_COUNT, None)
+        .generate_ed25519_addresses(coin_type, 0, 0..OWNING_ALIAS_COUNT, None)
         .await?;
 
     // create 10 different alias outputs with each owning various other assets

crates/iota-genesis-builder/src/stardust/test_outputs/delegator_outputs.rs

@@ -12,12 +12,11 @@ use iota_types::timelock::timelock::VESTED_REWARD_ID_PREFIX;
 use rand::{rngs::StdRng, Rng};
 
 use super::to_micros;
-use crate::stardust::types::{
-    output_header::OutputHeader, output_index::random_output_index_with_rng,
+use crate::stardust::{
+    test_outputs::{MERGE_MILESTONE_INDEX, MERGE_TIMESTAMP_SECS},
+    types::{output_header::OutputHeader, output_index::random_output_index_with_rng},
 };
 
-const MERGE_MILESTONE_INDEX: u32 = 7669900;
-const MERGE_TIMESTAMP_SECS: u32 = 1696406475;
 const A_WEEK_IN_SECONDS: u32 = 604_800;
 const TIMELOCK_MAX_ENDING_TIME: u32 = A_WEEK_IN_SECONDS * 208;
 

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

@@ -23,6 +23,7 @@ use iota_sdk::types::block::{
 };
 use iota_types::{
     gas_coin::STARDUST_TOTAL_SUPPLY_IOTA,
+    stardust::coin_type::CoinType,
     timelock::timelock::{self},
 };
 use packable::{
@@ -33,10 +34,17 @@ use rand::{rngs::StdRng, Rng, SeedableRng};
 
 use crate::stardust::{parse::HornetSnapshotParser, types::output_header::OutputHeader};
 
-const OUTPUT_TO_DECREASE_AMOUNT_FROM: &str =
+pub const IOTA_COIN_TYPE: u32 = 4218;
+const IOTA_OUTPUT_TO_DECREASE_AMOUNT_FROM: &str =
     "0xb462c8b2595d40d3ff19924e3731f501aab13e215613ce3e248d0ed9f212db160000";
-const MERGE_MILESTONE_INDEX: u32 = 7669900;
-const MERGE_TIMESTAMP_SECS: u32 = 1696406475;
+
+pub const SHIMMER_COIN_TYPE: u32 = 4219;
+pub const STARDUST_TOTAL_SUPPLY_SHIMMER_MICRO: u64 = 1_813_620_509_061_365;
+const SHIMMER_OUTPUT_TO_DECREASE_AMOUNT_FROM: &str =
+    "0x4c337ea67697cb8dd0267cced8d9b51c479eb61dea04842138dcef31218d63810000";
+
+pub const MERGE_MILESTONE_INDEX: u32 = 7669900;
+pub const MERGE_TIMESTAMP_SECS: u32 = 1696406475;
 
 pub const fn to_micros(n: u64) -> u64 {
     match n.checked_mul(1_000_000) {
@@ -60,6 +68,7 @@ const PROBABILITY_OF_PICKING_A_BASIC_OUTPUT: f64 = 0.1;
 pub async fn add_snapshot_test_outputs<const VERIFY: bool>(
     current_path: impl AsRef<Path> + core::fmt::Debug,
     new_path: impl AsRef<Path> + core::fmt::Debug,
+    coin_type: CoinType,
     randomness_seed: u64,
     delegator: impl Into<Option<Ed25519Address>>,
     with_sampling: bool,
@@ -72,13 +81,29 @@ pub async fn add_snapshot_test_outputs<const VERIFY: bool>(
     let mut new_header = parser.header.clone();
     let mut vested_index = u32::MAX;
 
+    let address_derivation_coin_type = match coin_type {
+        CoinType::Iota => IOTA_COIN_TYPE,
+        CoinType::Shimmer => SHIMMER_COIN_TYPE,
+    };
+
     let mut rng = StdRng::seed_from_u64(randomness_seed);
     let mut new_outputs = [
-        alias_ownership::outputs(&mut rng).await?,
-        stardust_mix::outputs(&mut rng, &mut vested_index).await?,
-        vesting_schedule_entity::outputs(&mut rng, &mut vested_index).await?,
-        vesting_schedule_iota_airdrop::outputs(&mut rng, &mut vested_index).await?,
-        vesting_schedule_portfolio_mix::outputs(&mut rng, &mut vested_index).await?,
+        alias_ownership::outputs(&mut rng, address_derivation_coin_type).await?,
+        stardust_mix::outputs(&mut rng, &mut vested_index, address_derivation_coin_type).await?,
+        vesting_schedule_entity::outputs(&mut rng, &mut vested_index, address_derivation_coin_type)
+            .await?,
+        vesting_schedule_iota_airdrop::outputs(
+            &mut rng,
+            &mut vested_index,
+            address_derivation_coin_type,
+        )
+        .await?,
+        vesting_schedule_portfolio_mix::outputs(
+            &mut rng,
+            &mut vested_index,
+            address_derivation_coin_type,
+        )
+        .await?,
     ]
     .concat();
 
@@ -91,9 +116,10 @@ pub async fn add_snapshot_test_outputs<const VERIFY: bool>(
             delegator,
             with_sampling.then_some(&mut parser),
             new_temp_amount,
+            coin_type,
         )?
     } else {
-        add_all_previous_outputs_and_test_outputs(&mut parser, new_temp_amount)?
+        add_all_previous_outputs_and_test_outputs(&mut parser, new_temp_amount, coin_type)?
     });
 
     // Adjust the output count according to newly generated outputs.
@@ -117,12 +143,18 @@ pub async fn add_snapshot_test_outputs<const VERIFY: bool>(
 fn add_all_previous_outputs_and_test_outputs<R: Read>(
     parser: &mut HornetSnapshotParser<R>,
     new_amount: u64,
+    coin_type: CoinType,
 ) -> anyhow::Result<Vec<(OutputHeader, Output)>> {
     let mut new_outputs = Vec::new();
 
+    let target_output = match coin_type {
+        CoinType::Iota => IOTA_OUTPUT_TO_DECREASE_AMOUNT_FROM,
+        CoinType::Shimmer => SHIMMER_OUTPUT_TO_DECREASE_AMOUNT_FROM,
+    };
+
     // Writes previous outputs.
     for (output_header, output) in parser.outputs().filter_map(|o| o.ok()) {
-        if output_header.output_id() == OutputId::from_str(OUTPUT_TO_DECREASE_AMOUNT_FROM)? {
+        if output_header.output_id() == OutputId::from_str(target_output)? {
             let basic = output.as_basic();
             let amount = basic
                 .amount()
@@ -154,6 +186,7 @@ fn add_only_test_outputs<R: Read>(
     delegator: Ed25519Address,
     parser: Option<&mut HornetSnapshotParser<R>>,
     new_temp_amount: u64,
+    coin_type: CoinType,
 ) -> anyhow::Result<Vec<(OutputHeader, Output)>> {
     // Needed outputs for delegator
     let mut new_outputs = delegator_outputs::outputs(rng, vested_index, delegator)?;
@@ -165,7 +198,11 @@ fn add_only_test_outputs<R: Read>(
 
     // Add all the remainder tokens to the zero address
     let zero_address = Ed25519Address::new([0; 32]);
-    let remainder = to_micros(STARDUST_TOTAL_SUPPLY_IOTA)
+    let network_total_supply = match coin_type {
+        CoinType::Iota => to_micros(STARDUST_TOTAL_SUPPLY_IOTA),
+        CoinType::Shimmer => STARDUST_TOTAL_SUPPLY_SHIMMER_MICRO,
+    };
+    let remainder = network_total_supply
         .checked_sub(new_temp_amount + new_outputs.iter().map(|o| o.1.amount()).sum::<u64>())
         .ok_or_else(|| anyhow!("new amount should not be higher than total supply"))?;
     let remainder_per_output = remainder / 4;