draft commit, not for review

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

@@ -115,6 +115,7 @@ pub struct Builder {
     migration_sources: Vec<SnapshotSource>,
     migration_tx_data: Option<MigrationTxData>,
     delegator: Option<String>,
+    delegator_map: Option<BTreeMap<String, String>>,
 }
 
 impl Default for Builder {
@@ -137,6 +138,7 @@ impl Builder {
             migration_sources: Default::default(),
             migration_tx_data: Default::default(),
             delegator: None,
+            delegator_map: None,
         }
     }
 
@@ -145,6 +147,11 @@ impl Builder {
         self
     }
 
+    pub fn with_delegator_map(mut self, delegator_map: BTreeMap<String, String>) -> Self {
+        self.delegator_map = Some(delegator_map);
+        self
+    }
+
     /// Checks if the genesis to be built has no migration or if it includes
     /// Stardust migration stakes
     pub fn contains_migrations(&self) -> bool {
@@ -259,10 +266,11 @@ impl Builder {
     fn create_and_cache_genesis_stake(&mut self) -> anyhow::Result<()> {
         if !self.migration_objects.is_empty() {
             if let Some(if_delegator_address) = &self.delegator {
+                // This means we have one delegator address supplied as argument
+                // with the iota genesis command. There should be no delegator_map.
                 let delegator = stardust_to_iota_address(
-                    Address::try_from_bech32(if_delegator_address).unwrap(),
-                )
-                .unwrap();
+                    Address::try_from_bech32(if_delegator_address)?,
+                )?;
                 // TODO: check whether we need to start with
                 // VALIDATOR_LOW_STAKE_THRESHOLD_NANOS
                 let minimum_stake = iota_types::governance::MIN_VALIDATOR_JOINING_STAKE_NANOS;
@@ -271,7 +279,50 @@ impl Builder {
                     delegator,
                     &self.migration_objects,
                     minimum_stake,
+                    None, // no genesis stake, we initialize it in the function.
                 )?;
+            } else if let Some(delegator_map) = &self.delegator_map {
+                // todo: probably better to merge the logic here with the above
+                // and modify the delegate_genesis_stake function.
+
+                // We can keep the delegate_genesis_stake function as is, and just
+                // run a for loop here with a new map, that maps a delegator to a list
+                // of validators. This way we can keep the logic in the function clean.
+
+                // todo: add a custom type for this map DelegatorMap for example.
+                let mut delegator_to_validators: BTreeMap<String, Vec<GenesisValidatorInfo>> = BTreeMap::new();
+                for validator in self.validators.values() {
+                    let validator_address = validator.info.account_address.to_string();
+                    // now check if the validator is in the delegator_map and add it to the delegator_to_validators map
+                    if let Some(delegator) = delegator_map.get(&validator_address) {
+                        delegator_to_validators
+                            .entry(delegator.to_string())
+                            .or_insert_with(Vec::new)
+                            .push(validator.clone());
+                    } else {
+                        bail!("A genesis with migrated state should have a delegator assigned");
+                    }
+                }
+
+                // todo: we can put delegate_genesis_stake in the implementation of GenesisStake
+                let mut genesis_stake = GenesisStake::default();
+                // And then we can do the for loop and use the same function as above
+                for (delegator, validators) in &delegator_to_validators {
+                    let delegator = stardust_to_iota_address(
+                        Address::try_from_bech32(delegator)?,
+                    )?;
+
+                    let minimum_stake = iota_types::governance::MIN_VALIDATOR_JOINING_STAKE_NANOS;
+                    genesis_stake = delegate_genesis_stake(
+                        validators.iter(),
+                        delegator,
+                        &self.migration_objects,
+                        minimum_stake,
+                        Some(genesis_stake)
+                    )?;
+                }
+
+                self.genesis_stake = genesis_stake;
             } else {
                 bail!("A genesis with migrated state should have a delegator assigned");
             }
@@ -828,7 +879,8 @@ impl Builder {
             genesis_stake: Default::default(),
             migration_sources,
             migration_tx_data,
-            delegator: None,
+            delegator: None, // todo: Probably need to load the delegator?
+            delegator_map: None, // todo: Probably need to load the delegator_map?
         };
 
         let unsigned_genesis_file = path.join(GENESIS_BUILDER_UNSIGNED_GENESIS_FILE);
@@ -916,6 +968,8 @@ impl Builder {
                 .save(file)?;
         }
 
+        // todo: probably need to save the delegator and delegator_map?
+
         Ok(())
     }
 }

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

@@ -179,6 +179,7 @@ pub fn delegate_genesis_stake<'info>(
     delegator: IotaAddress,
     migration_objects: &MigrationObjects,
     amount_nanos: u64,
+    genesis_stake: Option<GenesisStake>,
 ) -> anyhow::Result<GenesisStake> {
     let timelocks_pool = migration_objects.get_sorted_timelocks_and_expiration_by_owner(delegator);
     let gas_coins_pool = migration_objects.get_gas_coins_by_owner(delegator);
@@ -190,7 +191,9 @@ pub fn delegate_genesis_stake<'info>(
         .unwrap_or_default()
         .into_iter()
         .map(|object| (object, 0));
-    let mut genesis_stake = GenesisStake::default();
+
+    let mut genesis_stake = genesis_stake.unwrap_or_default();
+    // let mut genesis_stake = GenesisStake::default();
 
     // For each validator we try to fill their allocation up to
     // total_amount_to_stake_per_validator

crates/iota/src/genesis_ceremony.rs

@@ -3,7 +3,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::{fs::File, path::PathBuf};
-
+use std::collections::BTreeMap;
 use anyhow::Result;
 use camino::Utf8PathBuf;
 use clap::Parser;
@@ -118,6 +118,12 @@ pub enum CeremonyCommand {
         #[clap(long, name = "iota|<full-url>", help = "Remote migration snapshots.")]
         #[arg(num_args(0..))]
         remote_migration_snapshots: Vec<SnapshotUrl>,
+        #[clap(
+            long,
+            help = "Path to the delegator map file.",
+            name = "delegator_map.csv"
+        )]
+        delegator_map: Option<PathBuf>,
     },
     /// Examine the details of the built Genesis checkpoint.
     ExamineGenesisCheckpoint,
@@ -259,6 +265,7 @@ pub async fn run(cmd: Ceremony) -> Result<()> {
         CeremonyCommand::BuildUnsignedCheckpoint {
             local_migration_snapshots,
             remote_migration_snapshots,
+            delegator_map,
         } => {
             let local_snapshots = local_migration_snapshots
                 .into_iter()
@@ -271,6 +278,25 @@ pub async fn run(cmd: Ceremony) -> Result<()> {
             for source in local_snapshots.chain(remote_snapshots) {
                 builder = builder.add_migration_source(source);
             }
+
+            if let Some(delegator_map) = delegator_map {
+                // validator, delegator
+                let file = File::open(delegator_map)?;
+                let mut reader = csv::Reader::from_reader(file);
+                let mut map: BTreeMap<String, String> = BTreeMap::new();
+
+                // Add values to map.
+                for result in reader.records() {
+                    let record = result?;
+                    let address = record.get(0).ok_or_else(|| anyhow::anyhow!("Missing address"))?.to_string();
+                    let delegator = record.get(1).ok_or_else(|| anyhow::anyhow!("Missing delegator"))?.to_string();
+                    map.insert(address, delegator);
+                }
+
+                // Add map to builder.
+                builder = builder.with_delegator_map(map);
+            }
+
             tokio::task::spawn_blocking(move || {
                 let UnsignedGenesis { checkpoint, .. } = builder.get_or_build_unsigned_genesis();
                 println!(
@@ -472,6 +498,7 @@ mod test {
             command: CeremonyCommand::BuildUnsignedCheckpoint {
                 local_migration_snapshots: vec![],
                 remote_migration_snapshots: vec![],
+                delegator_map: PathBuf::new(),
             },
         };
         command.run().await?;