Refactoring, added Axon and Prometheus info queries

cybercongress · Dec 4, 2023 · c171c99 · c171c99
1 parent da415c2
commit c171c99
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 26 deletions.
diff --git a/src/contract.rs b/src/contract.rs
@@ -262,17 +262,16 @@ pub fn execute(
             weights,
             version_key,
         } => do_set_weights(deps, env, info, netuid, dests, weights, version_key),
-        ExecuteMsg::BecomeDelegate { hotkey, take } => {
-            do_become_delegate(deps, env, info, hotkey, take)
+        ExecuteMsg::BecomeDelegate { hotkey } => {
+            do_become_delegate(deps, env, info, hotkey)
         }
         ExecuteMsg::AddStake {
             hotkey,
-            amount_staked,
-        } => do_add_stake(deps, env, info, hotkey, amount_staked),
+        } => do_add_stake(deps, env, info, hotkey),
         ExecuteMsg::RemoveStake {
             hotkey,
-            amount_unstaked,
-        } => do_remove_stake(deps, env, info, hotkey, amount_unstaked),
+            amount,
+        } => do_remove_stake(deps, env, info, hotkey, amount),
         ExecuteMsg::ServeAxon {
             netuid,
             version,
@@ -495,11 +494,11 @@ pub fn sudo(deps: DepsMut, env: Env, msg: SudoMsg) -> Result<Response, ContractE
 pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
     match msg {
         QueryMsg::GetDelegates {} => to_json_binary(&get_delegates(deps)?),
-        QueryMsg::GetDelegate { delegate_account } => {
-            to_json_binary(&get_delegate(deps, delegate_account)?)
+        QueryMsg::GetDelegate { delegate } => {
+            to_json_binary(&get_delegate(deps, delegate)?)
         }
-        QueryMsg::GetDelegated { delegatee_account } => {
-            to_json_binary(&get_delegated(deps, delegatee_account)?)
+        QueryMsg::GetDelegated { delegatee } => {
+            to_json_binary(&get_delegated(deps, delegatee)?)
         }
         QueryMsg::GetNeuronsLite { netuid } => {
             to_json_binary(&get_neurons_lite(deps.storage, netuid)?)
@@ -516,11 +515,11 @@ pub fn query(deps: Deps, env: Env, msg: QueryMsg) -> StdResult<Binary> {
         QueryMsg::GetSubnetHyperparams { netuid } => {
             to_json_binary(&get_subnet_hyperparams(deps, netuid)?)
         }
-        QueryMsg::GetStakeInfoForColdkey { coldkey_account } => {
-            to_json_binary(&get_stake_info_for_coldkey(deps, coldkey_account)?)
+        QueryMsg::GetStakeInfoForColdkey { coldkey } => {
+            to_json_binary(&get_stake_info_for_coldkey(deps, coldkey)?)
         }
-        QueryMsg::GetStakeInfoForColdkeys { coldkey_accounts } => {
-            to_json_binary(&get_stake_info_for_coldkeys(deps, coldkey_accounts)?)
+        QueryMsg::GetStakeInfoForColdkeys { coldkeys } => {
+            to_json_binary(&get_stake_info_for_coldkeys(deps, coldkeys)?)
         }
         QueryMsg::GetNetworkRegistrationCost {} => to_json_binary(&get_network_lock_cost(
             deps.storage,

diff --git a/src/msg.rs b/src/msg.rs
@@ -6,7 +6,7 @@ pub struct InstantiateMsg {}
 #[cw_serde]
 pub enum ExecuteMsg {
     Activate {},
-    // TODO move to sudo and call as module from sdk
+    // TODO remove later, use for manual block_step
     BlockStep {},
 
     SetWeights {
@@ -17,15 +17,14 @@ pub enum ExecuteMsg {
     },
     BecomeDelegate {
         hotkey: String,
-        take: u16,
+        // take: u16,
     },
     AddStake {
         hotkey: String,
-        amount_staked: u64,
     },
     RemoveStake {
         hotkey: String,
-        amount_unstaked: u64,
+        amount: u64,
     },
     ServeAxon {
         netuid: u16,
@@ -226,9 +225,9 @@ pub enum QueryMsg {
     #[returns(Vec<crate::delegate_info::DelegateInfo>)]
     GetDelegates {},
     #[returns(Option<crate::delegate_info::DelegateInfo>)]
-    GetDelegate { delegate_account: String },
+    GetDelegate { delegate: String },
     #[returns(Vec<(crate::delegate_info::DelegateInfo, u64)>)]
-    GetDelegated { delegatee_account: String },
+    GetDelegated { delegatee: String },
 
     #[returns(Vec<crate::neuron_info::NeuronInfoLite>)]
     GetNeuronsLite { netuid: u16 },
@@ -247,9 +246,9 @@ pub enum QueryMsg {
     GetSubnetHyperparams { netuid: u16 },
 
     #[returns(crate::stake_info::StakeInfo)]
-    GetStakeInfoForColdkey { coldkey_account: String },
+    GetStakeInfoForColdkey { coldkey: String },
     #[returns(Vec<crate::stake_info::StakeInfo>)]
-    GetStakeInfoForColdkeys { coldkey_accounts: Vec<String> },
+    GetStakeInfoForColdkeys { coldkeys: Vec<String> },
 
     #[returns(u64)]
     GetNetworkRegistrationCost {},
@@ -312,10 +311,10 @@ pub enum QueryMsg {
     #[returns(u64)]
     GetTxRateLimit {},
 
-    #[returns(Option<AxonInfo>)]
+    #[returns(Option<crate::state::AxonInfoOf>)]
     GetAxonInfo { netuid: u16, hotkey: String },
 
-    #[returns(Option<PrometheusInfo>)]
+    #[returns(Option<crate::state::PrometheusInfoOf>)]
     GetPrometheusInfo { netuid: u16, hotkey: String },
 
     #[returns(Option<u64>)]

diff --git a/src/staking.rs b/src/staking.rs
@@ -45,9 +45,10 @@ pub fn do_become_delegate(
     env: Env,
     info: MessageInfo,
     hotkey_address: String,
-    take: u16,
+    // take: u16,
 ) -> Result<Response, ContractError> {
     // TODO set get_default_take() of custom take
+    let take = 11_796;
 
     // --- 1. We check the coldkey signuture.
     let coldkey = info.sender;
@@ -144,7 +145,6 @@ pub fn do_add_stake(
     env: Env,
     info: MessageInfo,
     hotkey_address: String,
-    _stake_to_be_added: u64,
 ) -> Result<Response, ContractError> {
     // --- 1. We check that the transaction is signed by the caller and retrieve the T::AccountId coldkey information.
     let coldkey = info.clone().sender;