apollo_staking: align the mock staking contract interface with the real contract

Author: dafnamatsry

crates/apollo_staking/src/contract_types.rs

@@ -4,6 +4,7 @@ use starknet_api::staking::StakingWeight;
 use starknet_types_core::felt::Felt;
 use thiserror::Error;
 
+#[cfg(test)]
 use crate::committee_provider::Staker;
 
 pub(crate) const GET_STAKERS_ENTRY_POINT: &str = "get_stakers";
@@ -14,46 +15,89 @@ pub(crate) const EPOCH_LENGTH: u64 = 100; // Number of heights in an epoch.
 #[derive(Debug, PartialEq, Eq)]
 struct ArrayRetdata<const N: usize, T>(Vec<T>);
 
+#[derive(Debug, PartialEq, Eq)]
+pub(crate) struct ContractStaker {
+    pub(crate) contract_address: ContractAddress,
+    pub(crate) staking_power: StakingWeight,
+    pub(crate) public_key: Option<Felt>,
+}
+
 #[derive(Debug, Error)]
 pub enum RetdataDeserializationError {
     #[error("Failed to convert Felt to ContractAddress: {address}")]
     ContractAddressConversionError { address: Felt },
     #[error("Failed to convert Felt to u128: {felt}")]
     U128ConversionError { felt: Felt },
+    #[error("Failed to convert Felt to usize: {felt}")]
+    USizeConversionError { felt: Felt },
     #[error(
         "Invalid retdata length: expected 1 Felt followed by {num_structs} (number of structs) *
          {struct_size} (number of Felts per struct), but received {length} Felts."
     )]
     InvalidArrayLength { length: usize, num_structs: usize, struct_size: usize },
+    #[error("Invalid retdata length: expected {expected} Felts, but received {received} Felts.")]
+    InvalidObjectLength { expected: usize, received: usize },
+    #[error("Unexpected enum variant: {variant}")]
+    UnexpectedEnumVariant { variant: usize },
 }
 
-impl Staker {
-    pub const CAIRO_OBJECT_NUM_FELTS: usize = 3;
+impl ContractStaker {
+    pub const CAIRO_OBJECT_NUM_FELTS: usize = 4;
 
     pub fn from_retdata_many(retdata: Retdata) -> Result<Vec<Self>, RetdataDeserializationError> {
-        Ok(ArrayRetdata::<{ Self::CAIRO_OBJECT_NUM_FELTS }, Staker>::try_from(retdata)?.0)
+        Ok(ArrayRetdata::<{ Self::CAIRO_OBJECT_NUM_FELTS }, ContractStaker>::try_from(retdata)?.0)
     }
 }
 
-impl TryFrom<[Felt; Self::CAIRO_OBJECT_NUM_FELTS]> for Staker {
+impl TryFrom<[Felt; Self::CAIRO_OBJECT_NUM_FELTS]> for ContractStaker {
     type Error = RetdataDeserializationError;
 
     fn try_from(felts: [Felt; Self::CAIRO_OBJECT_NUM_FELTS]) -> Result<Self, Self::Error> {
-        let [address, weight, public_key] = felts;
-        let address = ContractAddress::try_from(address)
-            .map_err(|_| RetdataDeserializationError::ContractAddressConversionError { address })?;
-        let weight = StakingWeight(
-            u128::try_from(weight)
-                .map_err(|_| RetdataDeserializationError::U128ConversionError { felt: weight })?,
-        );
-        Ok(Self { address, weight, public_key })
+        let [contract_address, staking_power, option_variant, public_key] = felts;
+        let contract_address = ContractAddress::try_from(contract_address).map_err(|_| {
+            RetdataDeserializationError::ContractAddressConversionError {
+                address: contract_address,
+            }
+        })?;
+        let staking_power = StakingWeight(u128::try_from(staking_power).map_err(|_| {
+            RetdataDeserializationError::U128ConversionError { felt: staking_power }
+        })?);
+        let option_variant = usize::try_from(option_variant).map_err(|_| {
+            RetdataDeserializationError::USizeConversionError { felt: option_variant }
+        })?;
+        let public_key = match option_variant {
+            0 => Some(public_key),
+            1 => None,
+            _ => {
+                return Err(RetdataDeserializationError::UnexpectedEnumVariant {
+                    variant: option_variant,
+                });
+            }
+        };
+        Ok(Self { contract_address, staking_power, public_key })
     }
 }
 
 #[cfg(test)]
-impl From<&Staker> for Vec<Felt> {
+impl From<&ContractStaker> for Vec<Felt> {
+    fn from(staker: &ContractStaker) -> Self {
+        vec![
+            Felt::from(staker.contract_address),
+            Felt::from(staker.staking_power.0),
+            Felt::from(if staker.public_key.is_some() { 0 } else { 1 }),
+            staker.public_key.unwrap_or_default(),
+        ]
+    }
+}
+
+#[cfg(test)]
+impl From<&Staker> for ContractStaker {
     fn from(staker: &Staker) -> Self {
-        vec![Felt::from(staker.address), Felt::from(staker.weight.0), staker.public_key]
+        Self {
+            contract_address: staker.address,
+            staking_power: staker.weight,
+            public_key: Some(staker.public_key),
+        }
     }
 }
 

crates/apollo_staking/src/staking_manager.rs

@@ -19,7 +19,7 @@ use crate::committee_provider::{
     ExecutionContext,
     Staker,
 };
-use crate::contract_types::{EPOCH_LENGTH, GET_STAKERS_ENTRY_POINT};
+use crate::contract_types::{ContractStaker, EPOCH_LENGTH, GET_STAKERS_ENTRY_POINT};
 use crate::utils::BlockRandomGenerator;
 
 pub type StakerSet = Vec<Staker>;
@@ -132,7 +132,17 @@ impl StakingManager {
             Calldata(vec![Felt::from(epoch)].into()),
         )?;
 
-        let stakers = Staker::from_retdata_many(call_info.execution.retdata)?;
+        let stakers: Vec<Staker> = ContractStaker::from_retdata_many(call_info.execution.retdata)?
+            .into_iter()
+            .filter_map(|staker| {
+                // Filter out stakers that don't have a public key.
+                staker.public_key.map(|public_key| Staker {
+                    address: staker.contract_address,
+                    weight: staker.staking_power,
+                    public_key,
+                })
+            })
+            .collect();
         let committee_members = self.select_committee(stakers);
 
         // Prepare the data needed for proposer selection.

crates/apollo_staking/src/staking_manager_test.rs

@@ -32,8 +32,8 @@ use crate::committee_provider::{
     ExecutionContext,
     Staker,
 };
-use crate::contract_types::RetdataDeserializationError;
-use crate::staking_manager::{StakerSet, StakingManager, StakingManagerConfig};
+use crate::contract_types::{ContractStaker, RetdataDeserializationError};
+use crate::staking_manager::{StakingManager, StakingManagerConfig};
 use crate::utils::MockBlockRandomGenerator;
 
 const STAKING_CONTRACT: FeatureContract =
@@ -93,7 +93,10 @@ fn default_config() -> StakingManagerConfig {
 }
 
 fn set_stakers(state: &mut State, block_context: &Context, stakers: &[Staker]) {
-    let mut stakers_as_felts: Vec<Felt> = stakers.iter().flat_map(<Vec<Felt>>::from).collect();
+    let mut stakers_as_felts: Vec<Felt> = stakers
+        .iter()
+        .flat_map(|staker| <Vec<Felt>>::from(&ContractStaker::from(staker)))
+        .collect();
     stakers_as_felts.insert(0, Felt::from(stakers.len()));
 
     // Invoke the set_stakers function on the mock staking contract.
@@ -123,7 +126,7 @@ fn get_committee_success(
     default_config: StakingManagerConfig,
     mut state: State,
     block_context: Context,
-    #[case] stakers: StakerSet,
+    #[case] stakers: Vec<Staker>,
     #[case] expected_committee: Committee,
 ) {
     set_stakers(&mut state, &block_context, &stakers);
@@ -276,25 +279,39 @@ async fn get_proposer_random_value_exceeds_total_weight(
     let _ = committee_manager.get_proposer(BlockNumber(1), 0, context).await;
 }
 
-// --- TryFrom tests for Staker and ArrayRetdata ---
+// --- TryFrom tests for ContractStaker and ArrayRetdata ---
 
 #[rstest]
 fn staker_try_from_valid() {
-    let staker = Staker::try_from([Felt::ONE, Felt::TWO, Felt::THREE]).unwrap();
-    assert_eq!(staker.address, contract_address!("0x1"));
-    assert_eq!(staker.weight, StakingWeight(2));
-    assert_eq!(staker.public_key, Felt::THREE);
+    let staker = ContractStaker::try_from([Felt::ONE, Felt::TWO, Felt::ZERO, Felt::THREE]).unwrap();
+    assert_eq!(staker.contract_address, contract_address!("0x1"));
+    assert_eq!(staker.staking_power, StakingWeight(2));
+    assert_eq!(staker.public_key, Some(Felt::THREE));
+
+    // A valid staker with no public key.
+    let staker = ContractStaker::try_from([Felt::ONE, Felt::TWO, Felt::ONE, Felt::THREE]).unwrap();
+    assert_eq!(staker.contract_address, contract_address!("0x1"));
+    assert_eq!(staker.staking_power, StakingWeight(2));
+    assert_eq!(staker.public_key, None);
 }
 
 #[rstest]
 fn staker_try_from_invalid_address() {
-    let err = Staker::try_from([CONTRACT_ADDRESS_DOMAIN_SIZE, Felt::ONE, Felt::ONE]).unwrap_err();
+    let err =
+        ContractStaker::try_from([CONTRACT_ADDRESS_DOMAIN_SIZE, Felt::ONE, Felt::ZERO, Felt::ONE])
+            .unwrap_err();
     assert_matches!(err, RetdataDeserializationError::ContractAddressConversionError { .. });
 }
 
+#[rstest]
+fn staker_try_from_invalid_public_key() {
+    let err = ContractStaker::try_from([Felt::ONE, Felt::TWO, Felt::TWO, Felt::THREE]).unwrap_err();
+    assert_matches!(err, RetdataDeserializationError::UnexpectedEnumVariant { .. });
+}
+
 #[rstest]
 fn staker_try_from_invalid_staked_amount() {
-    let err = Staker::try_from([Felt::ONE, Felt::MAX, Felt::ONE]).unwrap_err(); // Felt::MAX is too big for u128
+    let err = ContractStaker::try_from([Felt::ONE, Felt::MAX, Felt::ZERO, Felt::ONE]).unwrap_err(); // Felt::MAX is too big for u128
     assert_matches!(err, RetdataDeserializationError::U128ConversionError { .. });
 }
 
@@ -304,19 +321,19 @@ fn staker_try_from_invalid_staked_amount() {
 fn staker_array_retdata_try_from_valid(#[case] num_structs: usize) {
     let valid_retdata = [
         [Felt::from(num_structs)].as_slice(),
-        vec![Felt::ONE; Staker::CAIRO_OBJECT_NUM_FELTS * num_structs].as_slice(),
+        vec![Felt::ONE; ContractStaker::CAIRO_OBJECT_NUM_FELTS * num_structs].as_slice(),
     ]
     .concat();
 
-    let result = Staker::from_retdata_many(Retdata(valid_retdata)).unwrap();
+    let result = ContractStaker::from_retdata_many(Retdata(valid_retdata)).unwrap();
     assert_eq!(result.len(), num_structs);
 }
 
 #[rstest]
 #[case::empty_retdata(vec![])]
-#[case::missing_num_structs(vec![Felt::ONE; Staker::CAIRO_OBJECT_NUM_FELTS * 2])]
-#[case::invalid_staker_length(vec![Felt::ONE; Staker::CAIRO_OBJECT_NUM_FELTS - 1])]
+#[case::missing_num_structs(vec![Felt::ONE; ContractStaker::CAIRO_OBJECT_NUM_FELTS * 2])]
+#[case::invalid_staker_length(vec![Felt::ONE; ContractStaker::CAIRO_OBJECT_NUM_FELTS - 1])]
 fn staker_array_retdata_try_from_invalid_length(#[case] retdata: Vec<Felt>) {
-    let err = Staker::from_retdata_many(Retdata(retdata)).unwrap_err();
+    let err = ContractStaker::from_retdata_many(Retdata(retdata)).unwrap_err();
     assert_matches!(err, RetdataDeserializationError::InvalidArrayLength { .. });
 }

crates/blockifier_test_utils/resources/feature_contracts/cairo1/mock_staking.cairo

@@ -1,32 +1,44 @@
+use starknet::ContractAddress;
+
+pub type BlockNumber = u64;
+pub type Epoch = u64;
+pub type PublicKey = felt252;
+pub type StakingPower = u128;
+
 #[derive(Drop, Serde, starknet::Store)]
 pub struct Staker {
-    pub address: felt252,
-    pub staked_amount: u128,
-    pub pubkey: felt252,
+    pub contract_address: ContractAddress,
+    pub staking_power: StakingPower,
+    pub pub_key: Option<PublicKey>,
 }
 
 #[derive(Drop, Serde, starknet::Store)]
 pub struct EpochInfo {
-    pub epoch: u64,
-    pub start_block: u64,
-    pub end_block: u64,
+    pub epoch_id: Epoch,
+    pub start_block: BlockNumber,
+    pub epoch_length: u32,
 }
 
 #[starknet::interface]
 pub trait IStaking<TContractState> {
     fn add_staker(ref self: TContractState, staker: Staker);
     fn set_stakers(ref self: TContractState, stakers: Array<Staker>);
-    fn get_stakers(self: @TContractState, epoch: u64) -> Array<Staker>;
     fn set_current_epoch(ref self: TContractState, epoch: EpochInfo);
-    fn get_current_epoch(self: @TContractState) -> EpochInfo;
+
+    // The following functions have exactly the same interface as the real Staking contract.
+    fn get_stakers(
+        self: @TContractState, epoch_id: Epoch,
+    ) -> Span<(ContractAddress, StakingPower, Option<PublicKey>)>;
+    fn get_current_epoch_data(self: @TContractState) -> (Epoch, BlockNumber, u32);
 }
 
 #[starknet::contract]
 mod Staking {
+    use starknet::ContractAddress;
     use starknet::storage::{
         MutableVecTrait, StoragePointerReadAccess, StoragePointerWriteAccess, Vec, VecTrait,
     };
-    use super::{EpochInfo, Staker};
+    use super::{BlockNumber, Epoch, EpochInfo, PublicKey, Staker, StakingPower};
 
     #[storage]
     struct Storage {
@@ -50,21 +62,25 @@ mod Staking {
             }
         }
 
-        // epoch is not used in this mock, but should be part of the interface.
-        fn get_stakers(self: @ContractState, epoch: u64) -> Array<Staker> {
+        fn set_current_epoch(ref self: ContractState, epoch: EpochInfo) {
+            self.current_epoch.write(epoch);
+        }
+
+        // epoch_id is not used in this mock, but should be part of the interface.
+        fn get_stakers(
+            self: @ContractState, epoch_id: Epoch,
+        ) -> Span<(ContractAddress, StakingPower, Option<PublicKey>)> {
             let mut stakers = array![];
             for i in 0..self.stakers.len() {
-                stakers.append(self.stakers.at(i).read());
+                let staker = self.stakers.at(i).read();
+                stakers.append((staker.contract_address, staker.staking_power, staker.pub_key));
             }
-            stakers
-        }
-
-        fn set_current_epoch(ref self: ContractState, epoch: EpochInfo) {
-            self.current_epoch.write(epoch);
+            stakers.span()
         }
 
-        fn get_current_epoch(self: @ContractState) -> EpochInfo {
-            self.current_epoch.read()
+        fn get_current_epoch_data(self: @ContractState) -> (Epoch, BlockNumber, u32) {
+            let epoch_info = self.current_epoch.read();
+            (epoch_info.epoch_id, epoch_info.start_block, epoch_info.epoch_length)
         }
     }
 }