feat: add peer ids to each staker

docs/spec.md

@@ -59,6 +59,8 @@
     - [get\_total\_stake\_for\_token](#get_total_stake_for_token)
     - [set\_public\_key](#set_public_key)
     - [get\_current\_public\_key](#get_current_public_key)
+    - [set\_peer\_id](#set_peer_id)
+    - [get\_current\_peer\_id](#get_current_peer_id)
     - [get\_current\_epoch\_data](#get_current_epoch_data)
     - [update\_rewards](#update_rewards)
     - [get\_stakers](#get_stakers)
@@ -91,6 +93,7 @@
     - [Token Enabled](#token-enabled)
     - [Token Disabled](#token-disabled)
     - [Public Key Set](#public-key-set)
+    - [Peer Id Set](#peer-id-set)
     - [Consensus Rewards First Epoch Set](#consensus-rewards-first-epoch-set)
 - [Pool contract](#pool-contract)
   - [Functions](#functions-1)
@@ -1571,6 +1574,42 @@ Returns the current public key for the specified staker.
 Any address.
 #### logic <!-- omit from toc -->
 
+### set_peer_id
+```rust
+fn set_peer_id(ref self: ContractState, peer_id: PeerId)
+```
+#### description <!-- omit from toc -->
+Sets the peer ID for the caller. Changes take effect after K epochs.
+#### emits <!-- omit from toc -->
+1. [Peer Id Set](#peer-id-set)
+#### errors <!-- omit from toc -->
+1. [CONTRACT\_IS\_PAUSED](#contract_is_paused)
+2. [CALLER\_IS\_ZERO\_ADDRESS](#caller_is_zero_address)
+3. [INVALID\_PEER\_ID](#invalid_peer_id)
+4. [STAKER\_NOT\_EXISTS](#staker_not_exists)
+5. [UNSTAKE\_IN\_PROGRESS](#unstake_in_progress)
+6. [PEER\_ID\_SET\_IN\_PROGRESS](#peer_id_set_in_progress)
+7. [PEER\_ID\_MUST\_DIFFER](#peer_id_must_differ)
+#### pre-condition <!-- omit from toc -->
+#### access control <!-- omit from toc -->
+Only staker address.
+#### logic <!-- omit from toc -->
+
+### get_current_peer_id
+```rust
+fn get_current_peer_id(self: @ContractState, staker_address: ContractAddress) -> PeerId
+```
+#### description <!-- omit from toc -->
+Returns the current peer ID for the specified staker.
+#### emits <!-- omit from toc -->
+#### errors <!-- omit from toc -->
+1. [STAKER\_NOT\_EXISTS](#staker_not_exists)
+2. [PEER\_ID\_NOT\_SET](#peer_id_not_set)
+#### pre-condition <!-- omit from toc -->
+#### access control <!-- omit from toc -->
+Any address.
+#### logic <!-- omit from toc -->
+
 ### get_current_epoch_data
 ```rust
 fn get_current_epoch_data(self: @ContractState) -> (Epoch, BlockNumber, u32)
@@ -1614,10 +1653,10 @@ Only starkware sequencer.
 
 ### get_stakers
 ```rust
-fn get_stakers(self: @TContractState, epoch_id: Epoch) -> Span<(ContractAddress, StakingPower, Option<PublicKey>)>
+fn get_stakers(self: @TContractState, epoch_id: Epoch) -> Span<(ContractAddress, StakingPower, Option<PublicKey>, Option<PeerId>)>
 ```
 #### description <!-- omit from toc -->
-Returns a span of (staker_address, staking_power, Option<public_key>) for all stakers
+Returns a span of (staker_address, staking_power, Option<public_key>, Option<peer_id>) for all stakers
 for the given `epoch_id`.
 **Note**: The staking power is the relative weight of the staker's stake
 out of the total stake, including pooled stake (STRK and BTC), multiplied by
@@ -1822,6 +1861,12 @@ Any address.
 | staker_address | address                  | ✅    |
 | public_key     | [PublicKey](#publickey)  | ❌    |
 
+### Peer Id Set
+| data           | type                 | keyed |
+| -------------- | -------------------- | ----- |
+| staker_address | address              | ✅    |
+| peer_id        | [PeerId](#peerid)    | ❌    |
+
 ### Consensus Rewards First Epoch Set
 | data                   | type                     | keyed |
 | ---------------------- | ------------------------ | ----- |
@@ -2819,6 +2864,18 @@ Only token admin.
 ### PUBLIC_KEY_NOT_SET
 "Public key is not set"
 
+### PEER_ID_SET_IN_PROGRESS
+"Peer ID set is in progress"
+
+### PEER_ID_MUST_DIFFER
+"Peer ID is already set to provided value"
+
+### INVALID_PEER_ID
+"Peer ID is invalid"
+
+### PEER_ID_NOT_SET
+"Peer ID is not set"
+
 ### ATTEST_WITH_ZERO_BALANCE
 "Cannot attest with zero balance"
 
@@ -2980,5 +3037,8 @@ Epoch: u64
 ### PublicKey
 PublicKey: felt252
 
+### PeerId
+PeerId: u256
+
 ### BlockNumber
 BlockNumber: u64

src/event_test_utils.cairo

@@ -7,7 +7,7 @@ use staking::staking::interface::{
     ConfigEvents as StakingConfigEvents, Events as StakingEvents, PauseEvents as StakingPauseEvents,
     TokenManagerEvents as StakingTokenManagerEvents,
 };
-use staking::types::{Amount, Commission, Epoch, Inflation, PublicKey};
+use staking::types::{Amount, Commission, Epoch, Inflation, PeerId, PublicKey};
 use starknet::ContractAddress;
 use starkware_utils::time::time::{TimeDelta, Timestamp};
 use starkware_utils_testing::test_utils::assert_expected_event_emitted;
@@ -637,6 +637,18 @@ pub(crate) fn assert_public_key_set_event(
     );
 }
 
+pub(crate) fn assert_peer_id_set_event(
+    spied_event: @(ContractAddress, Event), staker_address: ContractAddress, peer_id: PeerId,
+) {
+    let expected_event = StakingEvents::PeerIdSet { staker_address, peer_id };
+    assert_expected_event_emitted(
+        :spied_event,
+        :expected_event,
+        expected_event_selector: @selector!("PeerIdSet"),
+        expected_event_name: "PeerIdSet",
+    );
+}
+
 pub(crate) fn assert_consensus_rewards_first_epoch_set_event(
     spied_event: @(ContractAddress, Event), consensus_rewards_first_epoch: Epoch,
 ) {

src/flow_test/flows.cairo

@@ -32,8 +32,8 @@ use staking::staking::staking::Staking::V3_PREV_CONTRACT_VERSION;
 use staking::staking::utils::STRK_WEIGHT_FACTOR;
 use staking::test_utils::constants::{
     AVG_BLOCK_DURATION, BTC_18D_CONFIG, BTC_8D_CONFIG, BTC_DECIMALS_18, BTC_DECIMALS_8,
-    EPOCH_DURATION, PUBLIC_KEY, STRK_BASE_VALUE, TEST_BTC_DECIMALS, TEST_MIN_BTC_FOR_REWARDS,
-    TEST_ONE_BTC,
+    EPOCH_DURATION, PEER_ID, PUBLIC_KEY, STRK_BASE_VALUE, TEST_BTC_DECIMALS,
+    TEST_MIN_BTC_FOR_REWARDS, TEST_ONE_BTC,
 };
 use staking::test_utils::{
     advance_blocks, calculate_pool_member_rewards, calculate_staker_btc_pool_rewards_v2,
@@ -7124,7 +7124,9 @@ pub(crate) impl SetPublicKeySameEpochAsUpgradeFlowImpl of FlowTrait<
         let result = staking_safe.get_current_public_key(staker_address: staker.staker.address);
         assert_panic_with_error(result, StakingError::PUBLIC_KEY_NOT_SET.describe());
         let stakers = staking_consensus.get_stakers(epoch_id: staking.get_current_epoch());
-        let expected_stakers = array![(staker.staker.address, STRK_WEIGHT_FACTOR, Option::None)]
+        let expected_stakers = array![
+            (staker.staker.address, STRK_WEIGHT_FACTOR, Option::None, Option::None),
+        ]
             .span();
         assert!(stakers == expected_stakers);
 
@@ -7135,7 +7137,177 @@ pub(crate) impl SetPublicKeySameEpochAsUpgradeFlowImpl of FlowTrait<
         assert!(actual_public_key == public_key);
         let stakers = staking_consensus.get_stakers(epoch_id: staking.get_current_epoch());
         let expected_stakers = array![
-            (staker.staker.address, STRK_WEIGHT_FACTOR, Option::Some(public_key)),
+            (staker.staker.address, STRK_WEIGHT_FACTOR, Option::Some(public_key), Option::None),
+        ]
+            .span();
+        assert!(stakers == expected_stakers);
+    }
+}
+
+/// Flow:
+/// Stake
+/// Upgrade
+/// Attempt to get peer id
+/// Catch PEER_ID_NOT_SET
+#[derive(Drop, Copy)]
+pub(crate) struct GetPeerIdAfterUpgradeFlow {
+    pub(crate) staker: Option<Staker>,
+}
+pub(crate) impl GetPeerIdAfterUpgradeFlowImpl of FlowTrait<GetPeerIdAfterUpgradeFlow> {
+    fn setup(ref self: GetPeerIdAfterUpgradeFlow, ref system: SystemState) {
+        let amount = system.staking.get_min_stake();
+        let staker = system.new_staker(:amount);
+        let commission = 200;
+        system.stake(:staker, :amount, pool_enabled: false, :commission);
+
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+        self.staker = Option::Some(staker);
+    }
+
+    #[feature("safe_dispatcher")]
+    fn test(self: GetPeerIdAfterUpgradeFlow, ref system: SystemState) {
+        let staker = self.staker.unwrap();
+        let staker_address = staker.staker.address;
+        let result = system.staking.safe_dispatcher().get_current_peer_id(:staker_address);
+        assert_panic_with_error(result, StakingError::PEER_ID_NOT_SET.describe());
+    }
+}
+
+/// Flow:
+/// Stake
+/// Exit
+/// Upgrade
+/// Attempt to set peer id
+#[derive(Drop, Copy)]
+pub(crate) struct ExitUpgradeSetPeerIdFlow {
+    pub(crate) staker: Option<Staker>,
+}
+pub(crate) impl ExitUpgradeSetPeerIdFlowImpl of MultiVersionFlowTrait<ExitUpgradeSetPeerIdFlow> {
+    fn versions(self: ExitUpgradeSetPeerIdFlow) -> Span<ReleaseVersion> {
+        [V0, V1, V2].span()
+    }
+
+    fn setup(ref self: ExitUpgradeSetPeerIdFlow, ref system: SystemState, version: ReleaseVersion) {
+        let amount = system.staking.get_min_stake();
+        let staker = system.new_staker(:amount);
+        let commission = 200;
+        system.stake(:staker, :amount, pool_enabled: true, :commission);
+        system.staker_exit_intent(:staker);
+        system.advance_exit_wait_window();
+        system.staker_exit_action(:staker);
+
+        self.staker = Option::Some(staker);
+    }
+
+    #[feature("safe_dispatcher")]
+    fn test(self: ExitUpgradeSetPeerIdFlow, ref system: SystemState, version: ReleaseVersion) {
+        let staker = self.staker.unwrap();
+        let peer_id = PEER_ID;
+
+        cheat_caller_address_once(
+            contract_address: system.staking.address, caller_address: staker.staker.address,
+        );
+        let result = system.staking.safe_dispatcher().set_peer_id(:peer_id);
+        assert_panic_with_error(result, StakingError::STAKER_NOT_EXISTS.describe());
+    }
+}
+
+/// Flow:
+/// Stake
+/// Exit intent
+/// Upgrade
+/// Attempt to set peer id
+#[derive(Drop, Copy)]
+pub(crate) struct IntentUpgradeSetPeerIdFlow {
+    pub(crate) staker: Option<Staker>,
+}
+pub(crate) impl IntentUpgradeSetPeerIdFlowImpl of MultiVersionFlowTrait<
+    IntentUpgradeSetPeerIdFlow,
+> {
+    fn versions(self: IntentUpgradeSetPeerIdFlow) -> Span<ReleaseVersion> {
+        [V0, V1, V2].span()
+    }
+
+    fn setup(
+        ref self: IntentUpgradeSetPeerIdFlow, ref system: SystemState, version: ReleaseVersion,
+    ) {
+        let amount = system.staking.get_min_stake();
+        let staker = system.new_staker(:amount);
+        let commission = 200;
+        system.stake(:staker, :amount, pool_enabled: false, :commission);
+        system.staker_exit_intent(:staker);
+
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+        self.staker = Option::Some(staker);
+    }
+
+    #[feature("safe_dispatcher")]
+    fn test(self: IntentUpgradeSetPeerIdFlow, ref system: SystemState, version: ReleaseVersion) {
+        let staker = self.staker.unwrap();
+        let peer_id = PEER_ID;
+
+        cheat_caller_address_once(
+            contract_address: system.staking.address, caller_address: staker.staker.address,
+        );
+        let result = system.staking.safe_dispatcher().set_peer_id(:peer_id);
+        assert_panic_with_error(result, StakingError::UNSTAKE_IN_PROGRESS.describe());
+    }
+}
+
+/// Flow:
+/// Stake
+/// Upgrade
+/// Set peer id (same epoch as upgrade)
+/// Test get_peer_id and get_stakers (still same epoch)
+/// Advance K epochs
+/// Test get_peer_id and get_stakers
+#[derive(Drop, Copy)]
+pub(crate) struct SetPeerIdSameEpochAsUpgradeFlow {
+    pub(crate) staker: Option<Staker>,
+}
+pub(crate) impl SetPeerIdSameEpochAsUpgradeFlowImpl of FlowTrait<SetPeerIdSameEpochAsUpgradeFlow> {
+    fn setup_v2(ref self: SetPeerIdSameEpochAsUpgradeFlow, ref system: SystemState) {
+        let amount = system.staking.get_min_stake();
+        let staker = system.new_staker(:amount);
+        let commission = 200;
+        system.stake(:staker, :amount, pool_enabled: false, :commission);
+        system.advance_epoch();
+
+        system.set_staker_for_migration(staker_address: staker.staker.address);
+        self.staker = Option::Some(staker);
+    }
+
+    #[feature("safe_dispatcher")]
+    fn test(self: SetPeerIdSameEpochAsUpgradeFlow, ref system: SystemState) {
+        let staker = self.staker.unwrap();
+        let staking = system.staking.dispatcher();
+        let staking_safe = system.staking.safe_dispatcher();
+        let staking_consensus = system.staking.consensus_dispatcher();
+        let peer_id = PEER_ID;
+
+        // Set peer id.
+        cheat_caller_address_once(
+            contract_address: system.staking.address, caller_address: staker.staker.address,
+        );
+        staking.set_peer_id(:peer_id);
+
+        // Test get_peer_id and get_stakers same epoch.
+        let result = staking_safe.get_current_peer_id(staker_address: staker.staker.address);
+        assert_panic_with_error(result, StakingError::PEER_ID_NOT_SET.describe());
+        let stakers = staking_consensus.get_stakers(epoch_id: staking.get_current_epoch());
+        let expected_stakers = array![
+            (staker.staker.address, STRK_WEIGHT_FACTOR, Option::None, Option::None),
+        ]
+            .span();
+        assert!(stakers == expected_stakers);
+
+        // Test get_peer_id and get_stakers after K epochs.
+        system.advance_k_epochs();
+        let actual_peer_id = staking.get_current_peer_id(staker_address: staker.staker.address);
+        assert!(actual_peer_id == peer_id);
+        let stakers = staking_consensus.get_stakers(epoch_id: staking.get_current_epoch());
+        let expected_stakers = array![
+            (staker.staker.address, STRK_WEIGHT_FACTOR, Option::None, Option::Some(peer_id)),
         ]
             .span();
         assert!(stakers == expected_stakers);
@@ -7179,7 +7351,9 @@ pub(crate) impl GetStakersAfterUpgradeFlowImpl of FlowTrait<GetStakersAfterUpgra
 
         // Test next epoch.
         let stakers = staking_consensus.get_stakers(epoch_id: epoch_id + 1);
-        let expected_stakers = array![(staker.staker.address, STRK_WEIGHT_FACTOR, Option::None)]
+        let expected_stakers = array![
+            (staker.staker.address, STRK_WEIGHT_FACTOR, Option::None, Option::None),
+        ]
             .span();
         assert!(stakers == expected_stakers);
         system.advance_epoch();

src/flow_test/fork_test.cairo

@@ -784,3 +784,17 @@ fn accrue_rewards_different_balances_across_versions_flow_test() {
     };
     test_flow_mainnet(ref :flow);
 }
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn get_peer_id_after_upgrade_flow_test() {
+    let mut flow = flows::GetPeerIdAfterUpgradeFlow { staker: Option::None };
+    test_flow_mainnet(ref :flow);
+}
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn set_peer_id_same_epoch_as_upgrade_flow_test() {
+    let mut flow = flows::SetPeerIdSameEpochAsUpgradeFlow { staker: Option::None };
+    test_flow_mainnet(ref :flow);
+}

src/flow_test/multi_version_tests.cairo

@@ -258,3 +258,17 @@ fn member_change_balance_claim_rewards_one_rewards_flow_test() {
     };
     test_multi_version_flow_mainnet(ref :flow);
 }
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn exit_upgrade_set_peer_id_flow_test() {
+    let mut flow = flows::ExitUpgradeSetPeerIdFlow { staker: Option::None };
+    test_multi_version_flow_mainnet(ref :flow);
+}
+
+#[test]
+#[fork("MAINNET_LATEST")]
+fn intent_upgrade_set_peer_id_flow_test() {
+    let mut flow = flows::IntentUpgradeSetPeerIdFlow { staker: Option::None };
+    test_multi_version_flow_mainnet(ref :flow);
+}