feat(staking): calculate block rewards by timestamp

Author: noa-starkware

src/flow_test/flow_ideas.md

@@ -49,3 +49,8 @@ more ideas:
 - enable token, update rewards, advance epoch, update rewards, advance epoch, update rewards - token does not get rewards until after 2 epochs
 - same as above with disable (can be implemented together as one test)
 - enable token A and disable token B, next epoch upgrade, test views and rewards.
+
+## block rewards by timestamp
+- advance blocks with different block times and check the avg is calculated correctly
+- update_rewards for blocks in same epoch - same rewards, then advance epoch, different rewards, update rewards for blocks in same epoch - same rewards.
+- update rewards is not called every block, still rewards is updated correctly (miss block, miss first block in epoch, miss epoch)

src/flow_test/test.cairo

@@ -2720,11 +2720,7 @@ fn update_rewards_transition_from_attestation_to_consensus_flow_test() {
     system.update_rewards(:staker, disable_rewards: false);
     let rewards = system.staker_claim_rewards(:staker);
     let (expected_rewards_v3, _) = calculate_staker_strk_rewards_with_balances_v3(
-        amount_own: stake_amount,
-        pool_amount: Zero::zero(),
-        :commission,
-        :staking_contract,
-        :minting_curve_contract,
+        amount_own: stake_amount, pool_amount: Zero::zero(), :commission, :minting_curve_contract,
     );
     assert!(expected_rewards_v3.is_non_zero());
     assert!(rewards == expected_rewards_v3);
@@ -2873,7 +2869,6 @@ fn update_rewards_disable_rewards_consensus_rewards_flow_test() {
         amount_own: stake_amount,
         pool_amount: Zero::zero(),
         :commission,
-        staking_contract: system.staking.address,
         minting_curve_contract: system.minting_curve.address,
     );
     assert!(expected_rewards.is_non_zero());
@@ -2920,7 +2915,6 @@ fn update_rewards_strk_pool_flow_test() {
         amount_own: stake_amount,
         pool_amount: delegation_amount,
         :commission,
-        staking_contract: system.staking.address,
         minting_curve_contract: system.minting_curve.address,
     );
     let staker_rewards = system.staker_claim_rewards(:staker);

src/staking/staking.cairo

@@ -188,6 +188,10 @@ pub mod Staking {
         pool_eic_class_hash: ClassHash,
         /// Map staker address to its version.
         staker_version: Map<ContractAddress, StakerVersion>,
+        /// Last epoch for which block rewards were calculated.
+        last_calculated_epoch: Epoch,
+        /// Block rewards (STRK, BTC) for the current epoch.
+        block_rewards: (Amount, Amount),
     }
 
     #[event]
@@ -1431,7 +1435,7 @@ pub mod Staking {
             // Get current block data and update rewards.
             let reward_supplier_dispatcher = self.reward_supplier_dispatcher.read();
             let (strk_block_rewards, btc_block_rewards) = self
-                .calculate_block_rewards(:reward_supplier_dispatcher);
+                .calculate_block_rewards(:reward_supplier_dispatcher, :curr_epoch);
             self
                 ._update_rewards(
                     :staker_address,
@@ -1490,13 +1494,24 @@ pub mod Staking {
         }
 
         /// Calculates the rewards for a block in the current epoch (for STRK and BTC).
+        ///
+        /// Precondition: `curr_epoch` must be `get_current_epoch()`, it's passed as a param to save
+        /// storage reads.
+        // TODO: Consider view?
+        // TODO: Migration.
         fn calculate_block_rewards(
-            self: @ContractState, reward_supplier_dispatcher: IRewardSupplierDispatcher,
+            ref self: ContractState,
+            reward_supplier_dispatcher: IRewardSupplierDispatcher,
+            curr_epoch: Epoch,
         ) -> (Amount, Amount) {
-            let (strk_rewards, btc_rewards) = reward_supplier_dispatcher
-                .calculate_current_epoch_rewards();
-            let epoch_len_in_blocks = self.get_epoch_info().epoch_len_in_blocks();
-            (strk_rewards / epoch_len_in_blocks.into(), btc_rewards / epoch_len_in_blocks.into())
+            if curr_epoch > self.last_calculated_epoch.read() {
+                self.last_calculated_epoch.write(curr_epoch);
+                let block_rewards = reward_supplier_dispatcher.update_current_epoch_block_rewards();
+                self.block_rewards.write(block_rewards);
+                block_rewards
+            } else {
+                self.block_rewards.read()
+            }
         }
 
         fn update_commission(

src/staking/tests/test.cairo

@@ -86,16 +86,17 @@ use starkware_utils_testing::test_utils::{
 };
 use test_utils::{
     StakingInitConfig, advance_block_into_attestation_window, advance_epoch_global,
-    advance_k_epochs_global, approve, calculate_staker_btc_pool_rewards_v2,
-    calculate_staker_btc_pool_rewards_v3, calculate_staker_strk_rewards_v2,
-    calculate_staker_strk_rewards_with_balances_v3, cheat_target_attestation_block_hash, constants,
-    custom_decimals_token, declare_pool_contract, declare_pool_eic_contract,
-    declare_staking_contract, declare_staking_eic_contract, deploy_mock_erc20_decimals_contract,
-    deploy_staking_contract, enter_delegation_pool_for_testing_using_dispatcher, fund,
-    general_contract_system_deployment, load_from_simple_map, load_one_felt, pause_staking_contract,
-    setup_btc_token, stake_for_testing_using_dispatcher, stake_from_zero_address,
-    stake_with_strk_pool_enabled, store_internal_staker_info_v0_to_map, store_to_simple_map,
-    to_amount_18_decimals, upgrade_implementation,
+    advance_k_epochs_global, approve, calculate_current_block_rewards_v3,
+    calculate_staker_btc_pool_rewards_v2, calculate_staker_btc_pool_rewards_v3,
+    calculate_staker_strk_rewards_v2, calculate_staker_strk_rewards_with_balances_v3,
+    cheat_target_attestation_block_hash, constants, custom_decimals_token, declare_pool_contract,
+    declare_pool_eic_contract, declare_staking_contract, declare_staking_eic_contract,
+    deploy_mock_erc20_decimals_contract, deploy_staking_contract,
+    enter_delegation_pool_for_testing_using_dispatcher, fund, general_contract_system_deployment,
+    load_from_simple_map, load_one_felt, pause_staking_contract, setup_btc_token,
+    stake_for_testing_using_dispatcher, stake_from_zero_address, stake_with_strk_pool_enabled,
+    store_internal_staker_info_v0_to_map, store_to_simple_map, to_amount_18_decimals,
+    upgrade_implementation,
 };
 
 #[test]
@@ -3490,10 +3491,7 @@ fn test_update_rewards_only_staker() {
         contract_address: staking_contract,
     };
     let staking_config_dispatcher = IStakingConfigDispatcher { contract_address: staking_contract };
-    let reward_supplier = cfg.staking_contract_info.reward_supplier;
-    let reward_supplier_dispatcher = IRewardSupplierDispatcher {
-        contract_address: reward_supplier,
-    };
+    let minting_curve_contract = cfg.reward_supplier.minting_curve_contract;
     let current_epoch = staking_dispatcher.get_current_epoch();
     cheat_caller_address_once(
         contract_address: staking_contract, caller_address: cfg.test_info.app_governor,
@@ -3505,9 +3503,7 @@ fn test_update_rewards_only_staker() {
     advance_k_epochs_global();
     let staker_address = cfg.test_info.staker_address;
     let staker_info_before = staking_dispatcher.staker_info_v1(:staker_address);
-    let (strk_epoch_rewards, _) = reward_supplier_dispatcher.calculate_current_epoch_rewards();
-    let epoch_len_in_blocks = cfg.staking_contract_info.epoch_info.epoch_len_in_blocks();
-    let strk_block_rewards = strk_epoch_rewards / epoch_len_in_blocks.into();
+    let (strk_block_rewards, _) = calculate_current_block_rewards_v3(:minting_curve_contract);
     let staker_info_expected = StakerInfoV1 {
         unclaimed_rewards_own: strk_block_rewards, ..staker_info_before,
     };
@@ -3536,10 +3532,7 @@ fn test_update_rewards_miss_blocks() {
         contract_address: staking_contract,
     };
     let staking_config_dispatcher = IStakingConfigDispatcher { contract_address: staking_contract };
-    let reward_supplier = cfg.staking_contract_info.reward_supplier;
-    let reward_supplier_dispatcher = IRewardSupplierDispatcher {
-        contract_address: reward_supplier,
-    };
+    let minting_curve_contract = cfg.reward_supplier.minting_curve_contract;
     let current_epoch = staking_dispatcher.get_current_epoch();
     cheat_caller_address_once(
         contract_address: staking_contract, caller_address: cfg.test_info.app_governor,
@@ -3551,9 +3544,7 @@ fn test_update_rewards_miss_blocks() {
     advance_k_epochs_global();
     let staker_address = cfg.test_info.staker_address;
     let staker_info_before = staking_dispatcher.staker_info_v1(:staker_address);
-    let (strk_epoch_rewards, _) = reward_supplier_dispatcher.calculate_current_epoch_rewards();
-    let epoch_len_in_blocks = cfg.staking_contract_info.epoch_info.epoch_len_in_blocks();
-    let strk_block_rewards = strk_epoch_rewards / epoch_len_in_blocks.into();
+    let (strk_block_rewards, _) = calculate_current_block_rewards_v3(:minting_curve_contract);
     let staker_info_expected = StakerInfoV1 {
         unclaimed_rewards_own: strk_block_rewards * 2, ..staker_info_before,
     };
@@ -3603,7 +3594,6 @@ fn test_update_rewards_with_strk_pool() {
         amount_own: stake_amount,
         pool_amount: delegated_amount,
         :commission,
-        :staking_contract,
         :minting_curve_contract,
     );
     // Assert staker rewards, delegator rewards, and pool balance before update.
@@ -3733,7 +3723,6 @@ fn test_update_rewards_with_both_strk_and_btc() {
         amount_own: stake_amount,
         pool_amount: strk_delegated_amount,
         :commission,
-        :staking_contract,
         :minting_curve_contract,
     );
     // Same calculation for both BTC pools (both have the same decimals).
@@ -3748,7 +3737,6 @@ fn test_update_rewards_with_both_strk_and_btc() {
         :normalized_pool_balance,
         :normalized_staker_total_btc_balance,
         :commission,
-        :staking_contract,
         :minting_curve_contract,
         token_address: btc_token_address,
     );

src/test_utils.cairo

@@ -27,7 +27,7 @@ use snforge_std::{
 };
 use staking::attestation::attestation::Attestation::MIN_ATTESTATION_WINDOW;
 use staking::attestation::interface::{IAttestationDispatcher, IAttestationDispatcherTrait};
-use staking::constants::{K, STARTING_EPOCH, STRK_IN_FRIS, STRK_TOKEN_ADDRESS};
+use staking::constants::{K, SECONDS_IN_YEAR, STARTING_EPOCH, STRK_IN_FRIS, STRK_TOKEN_ADDRESS};
 use staking::errors::InternalError;
 use staking::minting_curve::interface::{
     IMintingCurveConfigDispatcher, IMintingCurveConfigDispatcherTrait, IMintingCurveDispatcher,
@@ -40,6 +40,7 @@ use staking::pool::pool::Pool;
 use staking::pool::pool_member_balance_trace::trace::PoolMemberCheckpointTrait;
 use staking::pool::utils::compute_rewards_rounded_down;
 use staking::reward_supplier::reward_supplier::RewardSupplier;
+use staking::reward_supplier::reward_supplier::RewardSupplier::BLOCK_DURATION_SCALE;
 use staking::reward_supplier::utils::calculate_btc_rewards;
 use staking::staking::interface::{
     IStakingDispatcher, IStakingDispatcherTrait, IStakingPauseDispatcher,
@@ -1121,12 +1122,9 @@ pub(crate) fn calculate_staker_strk_rewards_with_balances_v3(
     amount_own: Amount,
     pool_amount: Amount,
     commission: Commission,
-    staking_contract: ContractAddress,
     minting_curve_contract: ContractAddress,
 ) -> (Amount, Amount) {
-    let (strk_block_rewards, _) = _calculate_current_block_rewards_v3(
-        :staking_contract, :minting_curve_contract,
-    );
+    let (strk_block_rewards, _) = calculate_current_block_rewards_v3(:minting_curve_contract);
     let total_stake = amount_own + pool_amount;
     // Calculate staker own rewards.
     let mut staker_rewards = mul_wide_and_div(
@@ -1166,7 +1164,6 @@ pub(crate) fn calculate_strk_pool_rewards_v3(
         amount_own: staker_info.amount_own,
         pool_amount: pool_info.amount,
         commission: pool_info.commission,
-        :staking_contract,
         :minting_curve_contract,
     );
     pool_rewards
@@ -1180,13 +1177,10 @@ pub(crate) fn calculate_staker_btc_pool_rewards_v3(
     normalized_pool_balance: NormalizedAmount,
     normalized_staker_total_btc_balance: NormalizedAmount,
     commission: Commission,
-    staking_contract: ContractAddress,
     minting_curve_contract: ContractAddress,
     token_address: ContractAddress,
 ) -> (Amount, Amount) {
-    let (_, btc_block_rewards) = _calculate_current_block_rewards_v3(
-        :staking_contract, :minting_curve_contract,
-    );
+    let (_, btc_block_rewards) = calculate_current_block_rewards_v3(:minting_curve_contract);
     // Calculate pool rewards including commission.
     let pool_rewards_including_commission = mul_wide_and_div(
         lhs: btc_block_rewards,
@@ -1203,17 +1197,33 @@ pub(crate) fn calculate_staker_btc_pool_rewards_v3(
 }
 
 
-fn _calculate_current_block_rewards_v3(
-    staking_contract: ContractAddress, minting_curve_contract: ContractAddress,
+/// Calculate block rewards for the current epoch. Use `AVG_BLOCK_TIME` (default value for testing)
+/// for `avg_block_time`.
+pub(crate) fn calculate_current_block_rewards_v3(
+    minting_curve_contract: ContractAddress,
 ) -> (Amount, Amount) {
-    let (strk_epoch_rewards, btc_epoch_rewards) = _calculate_current_epoch_rewards_v2(
-        :staking_contract, :minting_curve_contract,
-    );
-    let staking_dispatcher = IStakingDispatcher { contract_address: staking_contract };
-    let epoch_len_in_blocks = staking_dispatcher.get_epoch_info().epoch_len_in_blocks();
-    let strk_block_rewards = strk_epoch_rewards / epoch_len_in_blocks.into();
-    let btc_block_rewards = btc_epoch_rewards / epoch_len_in_blocks.into();
-    (strk_block_rewards, btc_block_rewards)
+    calculate_current_block_rewards_with_avg_block_time_v3(
+        :minting_curve_contract, avg_block_time: AVG_BLOCK_DURATION * BLOCK_DURATION_SCALE,
+    )
+}
+
+/// Calculate block rewards for the current epoch based on the given `avg_block_time`.
+pub(crate) fn calculate_current_block_rewards_with_avg_block_time_v3(
+    minting_curve_contract: ContractAddress, avg_block_time: u64,
+) -> (Amount, Amount) {
+    let minting_curve_dispatcher = IMintingCurveDispatcher {
+        contract_address: minting_curve_contract,
+    };
+    let yearly_mint = minting_curve_dispatcher.yearly_mint();
+    let total_rewards = mul_wide_and_div(
+        lhs: yearly_mint,
+        rhs: avg_block_time.into(),
+        div: BLOCK_DURATION_SCALE.into() * SECONDS_IN_YEAR.into(),
+    )
+        .expect_with_err(err: InternalError::REWARDS_COMPUTATION_OVERFLOW);
+    let btc_rewards = calculate_btc_rewards(:total_rewards);
+    let strk_rewards = total_rewards - btc_rewards;
+    (strk_rewards, btc_rewards)
 }
 
 // ---- Calculate Rewards - Helpers -----