refactor: expand StakeRewardDelta with pool field

Signed-off-by: William Hankins <[email protected]>

Author: whankinsiv

common/src/types.rs

@@ -214,6 +214,7 @@ pub struct StakeRewardDelta {
     pub stake_address: StakeAddress,
     pub delta: u64,
     pub reward_type: RewardType,
+    pub pool: PoolId,
 }
 
 /// Type of reward being given

modules/accounts_state/src/rewards.rs

@@ -1,11 +1,11 @@
 //! Acropolis AccountsState: rewards calculations
 
 use crate::snapshot::{Snapshot, SnapshotSPO};
-use acropolis_common::RewardType;
 use acropolis_common::{
     protocol_params::ShelleyParams, rational_number::RationalNumber, KeyHash, Lovelace, SPORewards,
     StakeAddress,
 };
+use acropolis_common::{PoolId, RewardType};
 use anyhow::{bail, Result};
 use bigdecimal::{BigDecimal, One, ToPrimitive, Zero};
 use std::cmp::min;
@@ -25,6 +25,9 @@ pub struct RewardDetail {
 
     /// Reward amount
     pub amount: Lovelace,
+
+    // Pool that reward came from
+    pub pool: PoolId,
 }
 
 /// Result of a rewards calculation
@@ -386,6 +389,7 @@ fn calculate_spo_rewards(
                     account: delegator_stake_address.clone(),
                     rtype: RewardType::Member,
                     amount: to_pay,
+                    pool: operator_id.to_vec(),
                 });
                 total_paid += to_pay;
                 delegators_paid += 1;
@@ -403,6 +407,7 @@ fn calculate_spo_rewards(
             account: spo.reward_account.clone(),
             rtype: RewardType::Leader,
             amount: spo_benefit,
+            pool: operator_id.to_vec(),
         });
     } else {
         info!(

modules/accounts_state/src/state.rs

@@ -15,8 +15,9 @@ use acropolis_common::{
     protocol_params::ProtocolParams,
     stake_addresses::{StakeAddressMap, StakeAddressState},
     BlockInfo, DRepChoice, DRepCredential, DelegatedStake, InstantaneousRewardSource,
-    InstantaneousRewardTarget, KeyHash, Lovelace, MoveInstantaneousReward, PoolLiveStakeInfo,
-    PoolRegistration, Pot, SPORewards, StakeAddress, StakeRewardDelta, TxCertificate,
+    InstantaneousRewardTarget, KeyHash, Lovelace, MoveInstantaneousReward, PoolId,
+    PoolLiveStakeInfo, PoolRegistration, Pot, SPORewards, StakeAddress, StakeRewardDelta,
+    TxCertificate,
 };
 use anyhow::Result;
 use imbl::OrdMap;
@@ -115,7 +116,7 @@ pub struct State {
     previous_protocol_parameters: Option<ProtocolParams>,
 
     /// Pool refunds to apply next epoch (list of reward accounts to refund to)
-    pool_refunds: Vec<StakeAddress>,
+    pool_refunds: Vec<(PoolId, StakeAddress)>,
 
     /// Addresses registration changes in current epoch
     current_epoch_registration_changes: Arc<Mutex<Vec<RegistrationChange>>>,
@@ -466,14 +467,15 @@ impl State {
         }
 
         // Send them their deposits back
-        for stake_address in refunds {
+        for (pool, stake_address) in refunds {
             // If their reward account has been deregistered, it goes to Treasury
             let mut stake_addresses = self.stake_addresses.lock().unwrap();
             if stake_addresses.is_registered(&stake_address) {
                 reward_deltas.push(StakeRewardDelta {
                     stake_address: stake_address.clone(),
                     delta: deposit,
                     reward_type: RewardType::PoolRefund,
+                    pool,
                 });
                 stake_addresses.add_to_reward(&stake_address, deposit);
             } else {
@@ -622,6 +624,7 @@ impl State {
                                 stake_address: reward.account.clone(),
                                 delta: reward.amount,
                                 reward_type: reward.rtype.clone(),
+                                pool: reward.pool.clone(),
                             })
                             .collect::<Vec<_>>(),
                     );
@@ -736,7 +739,11 @@ impl State {
                     hex::encode(id),
                     retired_spo.reward_account
                 );
-                self.pool_refunds.push(retired_spo.reward_account.clone()); // Store full StakeAddress
+                self.pool_refunds.push((
+                    retired_spo.operator.clone(),
+                    retired_spo.reward_account.clone(),
+                ));
+                // Store full StakeAddress
             }
 
             // Schedule to retire - we need them to still be in place when we count

modules/accounts_state/src/verifier.rs

@@ -166,10 +166,11 @@ impl Verifier {
                     continue;
                 };
 
-                expected_rewards.entry(spo).or_default().push(RewardDetail {
+                expected_rewards.entry(spo.clone()).or_default().push(RewardDetail {
                     account: stake_address,
                     rtype,
                     amount,
+                    pool: spo,
                 });
             }
 

modules/historical_accounts_state/src/state.rs

@@ -111,7 +111,7 @@ impl State {
             let update = RewardHistory {
                 epoch,
                 amount: reward.delta,
-                pool: Vec::new(),
+                pool: reward.pool.clone(),
                 reward_type: reward.reward_type.clone(),
             };
             entry.reward_history.get_or_insert_with(Vec::new).push(update);