feat(genesis): add delegator map type to genesis builder

crates/iota-config/src/genesis.rs

@@ -516,12 +516,7 @@ impl TokenDistributionSchedule {
     /// Helper to read a TokenDistributionSchedule from a csv file.
     ///
     /// The file is encoded such that the final entry in the CSV file is used to
-    /// denote the allocation to the stake subsidy fund. It must be in the
-    /// following format:
-    /// `0x0000000000000000000000000000000000000000000000000000000000000000,
-    /// <pre>minted supply</pre>,`
-    ///
-    /// All entries in a token distribution schedule must add up to 10B Iota.
+    /// denote the allocation to the stake subsidy fund.
     pub fn from_csv<R: std::io::Read>(reader: R) -> Result<Self> {
         let mut reader = csv::Reader::from_reader(reader);
         let mut allocations: Vec<TokenAllocation> =
@@ -628,3 +623,110 @@ impl TokenDistributionScheduleBuilder {
         schedule
     }
 }
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub struct ValidatorAllocation {
+    /// The validator address receiving the stake and/or gas payment
+    pub address: IotaAddress,
+    // The amount of nanos to stake to the validator
+    pub amount_nanos_to_stake: u64,
+    /// The amount of nanos to transfer as gas payment to the validator
+    pub amount_nanos_to_pay_gas: u64,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub struct DelegatorDistribution {
+    /// The address from which take the nanos for staking/gas
+    pub delegator: IotaAddress,
+    /// The allocation to a validator receiving a stake and/or a gas payment
+    pub validator_allocation: ValidatorAllocation,
+}
+
+#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
+#[serde(rename_all = "kebab-case")]
+pub struct DelegatorMap {
+    // Maps a delegator address to a tuple containing the address of a validator (1nd element of
+    // the tuple) that will receive an amount of nanos to stake (2nd element) and an amount as gas
+    // payment (3rd element).
+    pub allocations: HashMap<IotaAddress, Vec<ValidatorAllocation>>,
+}
+
+impl DelegatorMap {
+    pub fn new_for_validators_with_default_allocation<I: IntoIterator<Item = IotaAddress>>(
+        validators: I,
+        delegator: IotaAddress,
+    ) -> Self {
+        let default_allocation = iota_types::governance::MIN_VALIDATOR_JOINING_STAKE_NANOS;
+
+        let allocations = validators.into_iter().fold(
+            HashMap::new(),
+            |mut allocations: HashMap<IotaAddress, Vec<_>>, address| {
+                allocations
+                    .entry(delegator)
+                    .or_default()
+                    .push(ValidatorAllocation {
+                        address,
+                        amount_nanos_to_stake: default_allocation,
+                        amount_nanos_to_pay_gas: 0,
+                    });
+                allocations
+            },
+        );
+
+        Self { allocations }
+    }
+
+    /// Helper to read a DelegatorMap from a csv file.
+    ///
+    /// The file is encoded such that the final entry in the CSV file is used to
+    /// denote the allocation coming from a delegator. It must be in the
+    /// following format:
+    /// `delegator,validator,amount-nanos-to-stake,amount-nanos-to-pay-gas
+    /// <delegator1-address>,<validator-1-address>,2000000000000000,5000000000
+    /// <delegator1-address>,<validator-2-address>,3000000000000000,5000000000
+    /// <delegator2-address>,<validator-3-address>,4500000000000000,5000000000`
+    pub fn from_csv<R: std::io::Read>(reader: R) -> Result<Self> {
+        let mut reader = csv::Reader::from_reader(reader);
+
+        let allocations = reader
+            .deserialize::<DelegatorDistribution>()
+            .collect::<Result<Vec<_>, _>>()?
+            .into_iter()
+            .fold(
+                HashMap::new(),
+                |mut allocations: HashMap<IotaAddress, Vec<_>>, allocation| {
+                    allocations
+                        .entry(allocation.delegator)
+                        .or_default()
+                        .push(allocation.validator_allocation);
+                    allocations
+                },
+            );
+
+        Ok(Self { allocations })
+    }
+
+    /// Helper to write a DelegatorMap into a csv file.
+    ///
+    /// It writes in the following format:
+    /// `delegator,validator,amount-nanos-to-stake,amount-nanos-to-pay-gas
+    /// <delegator1-address>,<validator-1-address>,2000000000000000,5000000000
+    /// <delegator1-address>,<validator-2-address>,3000000000000000,5000000000
+    /// <delegator2-address>,<validator-3-address>,4500000000000000,5000000000`
+    pub fn to_csv<W: std::io::Write>(&self, writer: W) -> Result<()> {
+        let mut writer = csv::Writer::from_writer(writer);
+
+        for (&delegator, validators_allocations) in &self.allocations {
+            for &validator_allocation in validators_allocations {
+                writer.serialize(DelegatorDistribution {
+                    delegator,
+                    validator_allocation,
+                })?;
+            }
+        }
+
+        Ok(())
+    }
+}

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

@@ -19,8 +19,8 @@ use genesis_build_effects::GenesisBuildEffects;
 use iota_config::{
     IOTA_GENESIS_MIGRATION_TX_DATA_FILENAME,
     genesis::{
-        Genesis, GenesisCeremonyParameters, GenesisChainParameters, TokenDistributionSchedule,
-        UnsignedGenesis,
+        DelegatorMap, Genesis, GenesisCeremonyParameters, GenesisChainParameters,
+        TokenDistributionSchedule, UnsignedGenesis,
     },
     migration_tx_data::{MigrationTxData, TransactionsData},
 };
@@ -78,7 +78,7 @@ use move_binary_format::CompiledModule;
 use move_core_types::ident_str;
 use serde::{Deserialize, Serialize};
 use shared_crypto::intent::{Intent, IntentMessage, IntentScope};
-use stake::{GenesisStake, delegate_genesis_stake};
+use stake::GenesisStake;
 use stardust::migration::MigrationObjects;
 use tracing::trace;
 use validator_info::{GenesisValidatorInfo, GenesisValidatorMetadata, ValidatorInfo};
@@ -115,6 +115,7 @@ pub struct Builder {
     migration_sources: Vec<SnapshotSource>,
     migration_tx_data: Option<MigrationTxData>,
     delegator: Option<String>,
+    delegator_map: Option<DelegatorMap>,
 }
 
 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: DelegatorMap) -> 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 {
@@ -258,23 +265,27 @@ impl Builder {
     /// contains migrated objects.
     fn create_and_cache_genesis_stake(&mut self) -> anyhow::Result<()> {
         if !self.migration_objects.is_empty() {
-            if let Some(if_delegator_address) = &self.delegator {
-                let delegator = stardust_to_iota_address(
-                    Address::try_from_bech32(if_delegator_address).unwrap(),
+            let mut genesis_stake = GenesisStake::default();
+
+            let delegator_map = if let Some(delegator_map) = &self.delegator_map {
+                delegator_map.clone()
+            } else if let Some(delegator) = &self.delegator {
+                DelegatorMap::new_for_validators_with_default_allocation(
+                    self.validators.values().map(|v| v.info.iota_address()),
+                    stardust_to_iota_address(Address::try_from_bech32(delegator)?)?,
                 )
-                .unwrap();
-                // TODO: check whether we need to start with
-                // VALIDATOR_LOW_STAKE_THRESHOLD_NANOS
-                let minimum_stake = iota_types::governance::MIN_VALIDATOR_JOINING_STAKE_NANOS;
-                self.genesis_stake = delegate_genesis_stake(
-                    self.validators.values(),
+            } else {
+                bail!("no delegator/s assigned with a migration");
+            };
+
+            for (delegator, validators_allocations) in delegator_map.allocations {
+                genesis_stake.delegate_genesis_stake(
+                    &validators_allocations,
                     delegator,
                     &self.migration_objects,
-                    minimum_stake,
                 )?;
-            } else {
-                bail!("A genesis with migrated state should have a delegator assigned");
             }
+            self.genesis_stake = genesis_stake;
         }
         Ok(())
     }
@@ -828,7 +839,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 +928,8 @@ impl Builder {
                 .save(file)?;
         }
 
+        // todo: probably need to save the delegator and delegator_map?
+
         Ok(())
     }
 }