feat(staking): migration to rewards by timestamp

Author: noa-starkware

src/flow_test/flow_ideas.md

@@ -55,3 +55,9 @@ more ideas:
 - 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)
 - set block time config and test rewards after
+
+## rewards by timestamp - migration
+- set_consensus_rewards to future epoch, call update_rewards before consensus epoch and after, test rewards.
+- set_consensus_rewards to future epoch, call update_rewards only after consensus epoch, test rewards.
+- set_consensus_rewards to curr_epoch + 2. test rewards before and after. tets avg block time is update correctly.
+- set_consensus_rewards, update_rewards, then set_consensus_rewards to later epoch, update_rewards, then set_consensus_rewards to earlier epoch, update_rewards, test avg block time.

src/flow_test/test.cairo

@@ -18,11 +18,11 @@ use staking::staking::interface::{
 };
 use staking::staking::utils::{BTC_WEIGHT_FACTOR, STAKING_POWER_BASE_VALUE, STRK_WEIGHT_FACTOR};
 use staking::test_utils::constants::{
-    BTC_18D_CONFIG, BTC_5D_CONFIG, BTC_8D_CONFIG, PUBLIC_KEY, STRK_BASE_VALUE,
+    AVG_BLOCK_DURATION, BTC_18D_CONFIG, BTC_5D_CONFIG, BTC_8D_CONFIG, PUBLIC_KEY, STRK_BASE_VALUE,
     TEST_MIN_BTC_FOR_REWARDS,
 };
 use staking::test_utils::{
-    StakingInitConfig, calculate_staker_btc_pool_rewards_v2,
+    StakingInitConfig, advance_blocks, calculate_staker_btc_pool_rewards_v2,
     calculate_staker_strk_rewards_with_balances_v2, calculate_staker_strk_rewards_with_balances_v3,
     calculate_strk_pool_rewards_with_pool_balance_v2, compute_rewards_per_unit,
     custom_decimals_token, deploy_mock_erc20_decimals_contract,
@@ -2830,7 +2830,7 @@ fn update_rewards_disable_rewards_consensus_rewards_flow_test() {
     assert_panic_with_error(
         :result, expected_error: StakingError::REWARDS_ALREADY_UPDATED.describe(),
     );
-    advance_block_number_global(blocks: 1);
+    advance_blocks(blocks: 1, block_duration: AVG_BLOCK_DURATION);
 
     // Disable rewards = false before consensus epoch - no rewards
     system.update_rewards(:staker, disable_rewards: false);
@@ -2860,7 +2860,7 @@ fn update_rewards_disable_rewards_consensus_rewards_flow_test() {
     assert_panic_with_error(
         :result, expected_error: StakingError::REWARDS_ALREADY_UPDATED.describe(),
     );
-    advance_block_number_global(blocks: 1);
+    advance_blocks(blocks: 1, block_duration: AVG_BLOCK_DURATION);
 
     // Disable rewards = false with consensus on - rewards
     system.update_rewards(:staker, disable_rewards: false);

src/staking/staking.cairo

@@ -1278,8 +1278,22 @@ pub mod Staking {
 
         fn set_consensus_rewards_first_epoch(ref self: ContractState, epoch_id: Epoch) {
             self.roles.only_app_governor();
-            assert!(epoch_id >= self.get_current_epoch() + 2, "{}", Error::INVALID_EPOCH);
+            let curr_epoch = self.get_current_epoch();
+            assert!(epoch_id >= curr_epoch + 2, "{}", Error::INVALID_EPOCH);
             assert!(self.is_pre_consensus(), "{}", Error::CONSENSUS_REWARDS_IS_ACTIVE);
+            // If it's the first time setting the consensus rewards first epoch, initialize block
+            // rewards.
+            // We use `curr_epoch + 1` to ensure the average is computed only after completing
+            // at least one full epoch. This prevents calculating the average over too few blocks,
+            // since we are currently mid-epoch. The next calculation will therefore occur after
+            // at least one complete epoch (the next epoch) has passed.
+            if self.consensus_rewards_first_epoch.read().is_zero() {
+                self
+                    .calculate_block_rewards(
+                        reward_supplier_dispatcher: self.reward_supplier_dispatcher.read(),
+                        curr_epoch: curr_epoch + 1,
+                    );
+            }
             self.consensus_rewards_first_epoch.write(epoch_id);
             self
                 .emit(
@@ -1496,9 +1510,9 @@ 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.
+        /// storage reads. This precondition is only violated during the migration to consensus
+        /// rewards.
         // TODO: Consider view?
-        // TODO: Migration.
         fn calculate_block_rewards(
             ref self: ContractState,
             reward_supplier_dispatcher: IRewardSupplierDispatcher,

src/staking/tests/test.cairo

@@ -67,7 +67,7 @@ use staking::staking::staking::Staking::{
     DEFAULT_EXIT_WAIT_WINDOW, MAX_EXIT_WAIT_WINDOW, V3_PREV_CONTRACT_VERSION,
 };
 use staking::staking::utils::STRK_WEIGHT_FACTOR;
-use staking::types::{Epoch, InternalStakerInfoLatest, VecIndex};
+use staking::types::{Amount, Epoch, InternalStakerInfoLatest, VecIndex};
 use staking::{event_test_utils, test_utils};
 use starknet::class_hash::ClassHash;
 use starknet::{ContractAddress, Store, get_block_number};
@@ -3441,6 +3441,7 @@ fn test_update_rewards_from_attestation_contract_assertions() {
     let attestation_contract = cfg.test_info.attestation_contract;
     let staking_dispatcher = IStakingDispatcher { contract_address: staking_contract };
     let staking_config_dispatcher = IStakingConfigDispatcher { contract_address: staking_contract };
+    advance_epoch_global();
     stake_for_testing_using_dispatcher(:cfg);
 
     // Catch CALLER_IS_NOT_ATTESTATION_CONTRACT.
@@ -3492,6 +3493,7 @@ fn test_update_rewards_only_staker() {
     };
     let staking_config_dispatcher = IStakingConfigDispatcher { contract_address: staking_contract };
     let minting_curve_contract = cfg.reward_supplier.minting_curve_contract;
+    advance_epoch_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
     cheat_caller_address_once(
         contract_address: staking_contract, caller_address: cfg.test_info.app_governor,
@@ -3533,6 +3535,7 @@ fn test_update_rewards_miss_blocks() {
     };
     let staking_config_dispatcher = IStakingConfigDispatcher { contract_address: staking_contract };
     let minting_curve_contract = cfg.reward_supplier.minting_curve_contract;
+    advance_epoch_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
     cheat_caller_address_once(
         contract_address: staking_contract, caller_address: cfg.test_info.app_governor,
@@ -3570,6 +3573,7 @@ fn test_update_rewards_with_strk_pool() {
     let token = cfg.test_info.strk_token;
     let token_address = token.contract_address();
     let token_dispatcher = IERC20Dispatcher { contract_address: token_address };
+    advance_epoch_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
     cheat_caller_address_once(
         contract_address: staking_contract, caller_address: cfg.test_info.app_governor,
@@ -3658,6 +3662,7 @@ fn test_update_rewards_with_both_strk_and_btc() {
     let btc_token = cfg.test_info.btc_token;
     let btc_token_address = btc_token.contract_address();
     let token_dispatcher = IERC20Dispatcher { contract_address: token_address };
+    advance_epoch_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
     cheat_caller_address_once(
         contract_address: staking_contract, caller_address: cfg.test_info.app_governor,
@@ -3910,6 +3915,7 @@ fn test_update_rewards_assertions_already_consensus() {
     };
     let staking_config_dispatcher = IStakingConfigDispatcher { contract_address: staking_contract };
     let staker_address = cfg.test_info.staker_address;
+    advance_epoch_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
     cheat_caller_address_once(
         contract_address: staking_contract, caller_address: cfg.test_info.app_governor,
@@ -4566,6 +4572,9 @@ fn test_set_consensus_rewards_first_epoch() {
     let staking_safe_config_dispatcher = IStakingConfigSafeDispatcher {
         contract_address: staking_contract,
     };
+    let minting_curve_contract = cfg.reward_supplier.minting_curve_contract;
+    stake_for_testing_using_dispatcher(:cfg);
+    advance_k_epochs_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
     let mut spy = snforge_std::spy_events();
     cheat_caller_address_once(
@@ -4578,6 +4587,31 @@ fn test_set_consensus_rewards_first_epoch() {
         .try_into()
         .unwrap();
     assert!(consensus_rewards_first_epoch == current_epoch + 2);
+    // Test that block rewards are initialized.
+    let block_rewards = snforge_std::load(
+        target: staking_contract,
+        storage_address: selector!("block_rewards"),
+        size: Store::<(Amount, Amount)>::size().into(),
+    )
+        .span();
+    let strk_block_rewards = (*block_rewards.at(0)).try_into().unwrap();
+    let btc_block_rewards = (*block_rewards.at(1)).try_into().unwrap();
+    let (expected_strk_block_rewards, expected_btc_block_rewards) =
+        calculate_current_block_rewards_v3(
+        :minting_curve_contract,
+    );
+    assert!(expected_strk_block_rewards.is_non_zero());
+    assert!(expected_btc_block_rewards.is_non_zero());
+    assert!(strk_block_rewards == expected_strk_block_rewards);
+    assert!(btc_block_rewards == expected_btc_block_rewards);
+    let last_calculated_epoch = load_one_felt(
+        target: staking_contract, storage_address: selector!("last_calculated_epoch"),
+    )
+        .try_into()
+        .unwrap();
+    let expected_last_calculated_epoch = current_epoch + 1;
+    assert!(last_calculated_epoch == expected_last_calculated_epoch);
+    // Test that event is emitted.
     let events = spy.get_events().emitted_by(contract_address: staking_contract).events;
     assert_number_of_events(
         actual: events.len(), expected: 1, message: "set_consensus_rewards_first_epoch",
@@ -4597,6 +4631,23 @@ fn test_set_consensus_rewards_first_epoch() {
         .try_into()
         .unwrap();
     assert!(consensus_rewards_first_epoch == current_epoch + 4);
+    // Test that block rewards are initialized.
+    let block_rewards = snforge_std::load(
+        target: staking_contract,
+        storage_address: selector!("block_rewards"),
+        size: Store::<(Amount, Amount)>::size().into(),
+    )
+        .span();
+    let strk_block_rewards = (*block_rewards.at(0)).try_into().unwrap();
+    let btc_block_rewards = (*block_rewards.at(1)).try_into().unwrap();
+    assert!(strk_block_rewards == expected_strk_block_rewards);
+    assert!(btc_block_rewards == expected_btc_block_rewards);
+    let last_calculated_epoch = load_one_felt(
+        target: staking_contract, storage_address: selector!("last_calculated_epoch"),
+    )
+        .try_into()
+        .unwrap();
+    assert!(last_calculated_epoch == expected_last_calculated_epoch);
     // Set consensus rewards epoch again with earlier epoch.
     cheat_caller_address_once(
         contract_address: staking_contract, caller_address: cfg.test_info.app_governor,
@@ -4608,6 +4659,23 @@ fn test_set_consensus_rewards_first_epoch() {
         .try_into()
         .unwrap();
     assert!(consensus_rewards_first_epoch == current_epoch + 3);
+    // Test that block rewards are initialized.
+    let block_rewards = snforge_std::load(
+        target: staking_contract,
+        storage_address: selector!("block_rewards"),
+        size: Store::<(Amount, Amount)>::size().into(),
+    )
+        .span();
+    let strk_block_rewards = (*block_rewards.at(0)).try_into().unwrap();
+    let btc_block_rewards = (*block_rewards.at(1)).try_into().unwrap();
+    assert!(strk_block_rewards == expected_strk_block_rewards);
+    assert!(btc_block_rewards == expected_btc_block_rewards);
+    let last_calculated_epoch = load_one_felt(
+        target: staking_contract, storage_address: selector!("last_calculated_epoch"),
+    )
+        .try_into()
+        .unwrap();
+    assert!(last_calculated_epoch == expected_last_calculated_epoch);
     // Set consensus rewards epoch again when currently set to current epoch + 1.
     advance_k_epochs_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
@@ -4621,6 +4689,23 @@ fn test_set_consensus_rewards_first_epoch() {
         .try_into()
         .unwrap();
     assert!(consensus_rewards_first_epoch == current_epoch + 2);
+    // Test that block rewards are initialized.
+    let block_rewards = snforge_std::load(
+        target: staking_contract,
+        storage_address: selector!("block_rewards"),
+        size: Store::<(Amount, Amount)>::size().into(),
+    )
+        .span();
+    let strk_block_rewards = (*block_rewards.at(0)).try_into().unwrap();
+    let btc_block_rewards = (*block_rewards.at(1)).try_into().unwrap();
+    assert!(strk_block_rewards == expected_strk_block_rewards);
+    assert!(btc_block_rewards == expected_btc_block_rewards);
+    let last_calculated_epoch = load_one_felt(
+        target: staking_contract, storage_address: selector!("last_calculated_epoch"),
+    )
+        .try_into()
+        .unwrap();
+    assert!(last_calculated_epoch == expected_last_calculated_epoch);
     // Advance to consensus rewards and test that consensus rewards epoch cannot be changed.
     advance_k_epochs_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
@@ -4643,6 +4728,7 @@ fn test_set_consensus_rewards_first_epoch_assertions() {
     let staking_safe_config_dispatcher = IStakingConfigSafeDispatcher {
         contract_address: staking_contract,
     };
+    advance_epoch_global();
     let current_epoch = staking_dispatcher.get_current_epoch();
 
     // Catch ONLY_APP_GOVERNOR.

src/test_utils.cairo

@@ -938,6 +938,13 @@ pub(crate) fn advance_epoch_global_custom_time(block_time: TimeDelta) {
     advance_time_global(:time);
 }
 
+/// Advance the block number by the given `blocks` and the timestamp by the given `block_duration`
+/// in seconds * `blocks`.
+pub(crate) fn advance_blocks(blocks: u64, block_duration: Seconds) {
+    advance_block_number_global(:blocks);
+    advance_time_global(time: TimeDelta { seconds: block_duration * blocks });
+}
+
 // ---- Calculate Rewards - V0 (index based) -----
 
 /// Update rewards for STRK pool.