Merge pull request #996 from opentensor/mechanism-enum

Use an enum to represent mechanism ID

Author: keithtensor

pallets/admin-utils/tests/tests.rs

@@ -5,7 +5,7 @@ use frame_support::{
 };
 use frame_system::Config;
 use pallet_admin_utils::Error;
-use pallet_subtensor::{migrations, Error as SubtensorError, Event};
+use pallet_subtensor::{migrations, subnets, Error as SubtensorError, Event};
 use sp_core::U256;
 
 mod mock;
@@ -1203,7 +1203,7 @@ fn test_set_alpha_values_dispatch_info_ok() {
 fn test_sudo_get_set_alpha() {
     new_test_ext().execute_with(|| {
         let netuid: u16 = 1;
-        let mechid: u16 = 1;
+        let mechid: subnets::Mechanism = subnets::Mechanism::Dynamic;
         let alpha_low: u16 = 12_u16;
         let alpha_high: u16 = u16::MAX - 10;
 

pallets/subtensor/src/benchmarks.rs

@@ -2,8 +2,8 @@
 #![allow(clippy::arithmetic_side_effects, clippy::unwrap_used)]
 #![cfg(feature = "runtime-benchmarks")]
 
-use crate::Pallet as Subtensor;
 use crate::*;
+use crate::{subnets::Mechanism, Pallet as Subtensor};
 use frame_benchmarking::{account, benchmarks, whitelisted_caller};
 use frame_support::assert_ok;
 use frame_system::RawOrigin;
@@ -300,7 +300,7 @@ benchmarks! {
     let amount: u64 = 1;
     let amount_to_be_staked = 100_000_000_000_000u64;
     Subtensor::<T>::add_balance_to_coldkey_account(&coldkey.clone(), amount_to_be_staked);
-  }: register_network(RawOrigin::Signed(coldkey), hotkey, 1)
+  }: register_network(RawOrigin::Signed(coldkey), hotkey, Mechanism::Dynamic)
 
   // swap_hotkey {
   //   let seed: u32 = 1;

pallets/subtensor/src/coinbase/run_coinbase.rs

@@ -1,6 +1,7 @@
 use super::*;
 use crate::epoch::math::safe_modulo;
 use alloc::collections::BTreeMap;
+use subnets::Mechanism;
 use substrate_fixed::types::I96F32;
 
 impl<T: Config> Pallet<T> {
@@ -28,7 +29,7 @@ impl<T: Config> Pallet<T> {
 
         // --- 4. Sum all the SubnetTAO associated with the same mechanism
         let mut total_active_tao: I96F32 = I96F32::from_num(0);
-        let mut mechanism_tao: BTreeMap<u16, I96F32> = BTreeMap::new();
+        let mut mechanism_tao: BTreeMap<Mechanism, I96F32> = BTreeMap::new();
         for netuid in subnets.iter() {
             if *netuid == 0 {
                 continue;
@@ -63,7 +64,7 @@ impl<T: Config> Pallet<T> {
                 continue;
             }
             // 1. Get subnet mechanism ID
-            let mechid: u16 = SubnetMechanism::<T>::get(*netuid);
+            let mechid: Mechanism = SubnetMechanism::<T>::get(*netuid);
             // 2. Get subnet TAO (T_s)
             let subnet_tao: I96F32 = I96F32::from_num(SubnetTAO::<T>::get(*netuid));
             // 3. Get the denominator as the sum of all TAO associated with a specific mechanism (T_m)
@@ -94,7 +95,7 @@ impl<T: Config> Pallet<T> {
             // 11. Increase total issuance: I_new = I_old + E_s
             TotalIssuance::<T>::mutate(|total| *total = total.saturating_add(subnet_emission));
             // 12. Switch on dynamic or Stable.
-            if mechid == 1 {
+            if mechid.is_dynamic() {
                 // 12a Dynamic: Add the SubnetAlpha directly into the pool immediately: A_s_new = A_s_old + E_m
                 SubnetAlphaIn::<T>::mutate(*netuid, |total| {
                     *total = total.saturating_add(block_emission.to_num::<u64>())

pallets/subtensor/src/lib.rs

@@ -61,6 +61,7 @@ extern crate alloc;
 #[frame_support::pallet]
 pub mod pallet {
     use crate::migrations;
+    use crate::subnets::Mechanism;
     use frame_support::{
         dispatch::GetDispatchInfo,
         pallet_prelude::{DispatchResult, StorageMap, ValueQuery, *},
@@ -225,11 +226,6 @@ pub mod pallet {
         0
     }
     #[pallet::type_value]
-    /// Default value for zero.
-    pub fn DefaultZeroU16<T: Config>() -> u16 {
-        0
-    }
-    #[pallet::type_value]
     /// Default value for false.
     pub fn DefaultFalse<T: Config>() -> bool {
         false
@@ -1012,8 +1008,7 @@ pub mod pallet {
     /// ==== Subnet Parameters =====
     /// ============================
     #[pallet::storage] // --- MAP ( netuid ) --> subnet mechanism
-    pub type SubnetMechanism<T: Config> =
-        StorageMap<_, Identity, u16, u16, ValueQuery, DefaultZeroU16<T>>;
+    pub type SubnetMechanism<T: Config> = StorageMap<_, Identity, u16, Mechanism, ValueQuery>;
     #[pallet::storage]
     /// --- MAP ( netuid ) --> subnetwork_n (Number of UIDs in the network).
     pub type SubnetworkN<T: Config> = StorageMap<_, Identity, u16, u16, ValueQuery, DefaultN<T>>;

pallets/subtensor/src/macros/dispatches.rs

@@ -1065,7 +1065,7 @@ mod dispatches {
         pub fn register_network(
             origin: OriginFor<T>,
             hotkey: T::AccountId,
-            mechid: u16,
+            mechid: Mechanism,
         ) -> DispatchResult {
             Self::do_register_network(origin, &hotkey, mechid, None)
         }
@@ -1339,7 +1339,7 @@ mod dispatches {
         pub fn register_network_with_identity(
             origin: OriginFor<T>,
             hotkey: T::AccountId,
-            mechid: u16,
+            mechid: Mechanism,
             identity: Option<SubnetIdentityOf>,
         ) -> DispatchResult {
             Self::do_register_network(origin, &hotkey, mechid, identity)

pallets/subtensor/src/macros/genesis.rs

@@ -47,7 +47,7 @@ mod genesis {
             for net in 1..2 {
                 let netuid: u16 = net as u16;
                 let hotkey = DefaultAccount::<T>::get();
-                SubnetMechanism::<T>::insert(netuid, 1); // Make dynamic.
+                SubnetMechanism::<T>::insert(netuid, Mechanism::Dynamic); // Make dynamic.
                 Owner::<T>::insert(hotkey.clone(), hotkey.clone());
                 SubnetAlphaIn::<T>::insert(netuid, 1);
                 SubnetTAO::<T>::insert(netuid, 10_000_000_000);

pallets/subtensor/src/migrations/migrate_rao.rs

@@ -3,6 +3,7 @@ use alloc::string::String;
 use frame_support::IterableStorageMap;
 use frame_support::{traits::Get, weights::Weight};
 use log;
+use subnets::Mechanism;
 
 pub fn migrate_rao<T: Config>() -> Weight {
     let migration_name = b"migrate_rao".to_vec();
@@ -87,7 +88,7 @@ pub fn migrate_rao<T: Config>() -> Weight {
             *total = total.saturating_add(lock);
         }); // Increase the stake.
         TotalStake::<T>::put(TotalStake::<T>::get().saturating_add(lock)); // Increase the total stake.
-        SubnetMechanism::<T>::insert(netuid, 1); // Convert to dynamic immediately with initialization.
+        SubnetMechanism::<T>::insert(netuid, Mechanism::Dynamic); // Convert to dynamic immediately with initialization.
         SubnetLocked::<T>::insert(netuid, lock);
         // Update all tempos to default
         Tempo::<T>::insert(netuid, DefaultTempo::<T>::get());

pallets/subtensor/src/staking/stake_ops.rs

@@ -1,4 +1,5 @@
 use super::*;
+use subnets::Mechanism;
 use substrate_fixed::types::I96F32;
 
 impl<T: Config> Pallet<T> {
@@ -116,14 +117,14 @@ impl<T: Config> Pallet<T> {
         tao_staked: u64,
     ) -> u64 {
         // Step 1: Get the mechanism type for the subnet (0 for Stable, 1 for Dynamic)
-        let mechanism_id: u16 = SubnetMechanism::<T>::get(netuid);
+        let mechanism_id: Mechanism = SubnetMechanism::<T>::get(netuid);
 
         let alpha_staked: I96F32;
         let new_subnet_alpha: I96F32;
         // Step 2: Convert tao_staked to a fixed-point number for precise calculations
         let tao_staked_float: I96F32 = I96F32::from_num(tao_staked);
 
-        if mechanism_id == 1 {
+        if mechanism_id.is_dynamic() {
             // Step 3: Dynamic mechanism calculations
             // Step 3a: Get current TAO and Alpha in the subnet
             let subnet_tao: I96F32 = I96F32::from_num(SubnetTAO::<T>::get(netuid));
@@ -286,7 +287,7 @@ impl<T: Config> Pallet<T> {
         alpha_unstaked: u64,
     ) -> u64 {
         // Step 1: Get the mechanism type for this subnet
-        let mechid: u16 = SubnetMechanism::<T>::get(netuid);
+        let mechid: Mechanism = SubnetMechanism::<T>::get(netuid);
 
         // Step 2: Initialize variables for TAO unstaked and new subnet alpha
         let tao_unstaked: I96F32;
@@ -306,7 +307,7 @@ impl<T: Config> Pallet<T> {
             return 0;
         }
 
-        if mechid == 1 {
+        if mechid.is_dynamic() {
             // Step 4: Dynamic mechanism
             // Step 4a: Get current TAO in the subnet
             let subnet_tao: I96F32 = I96F32::from_num(SubnetTAO::<T>::get(netuid));

pallets/subtensor/src/subnets/mod.rs

@@ -1,8 +1,39 @@
 use super::*;
+use codec::{Decode, Encode, MaxEncodedLen};
 pub mod lock;
 pub mod registration;
 pub mod serving;
 pub mod subnet;
 pub mod tempo;
 pub mod uids;
 pub mod weights;
+
+#[derive(
+    Clone,
+    Copy,
+    Debug,
+    Decode,
+    Default,
+    Encode,
+    PartialEq,
+    Eq,
+    PartialOrd,
+    Ord,
+    MaxEncodedLen,
+    scale_info::TypeInfo,
+)]
+pub enum Mechanism {
+    #[default]
+    Stable = 0,
+    Dynamic = 1,
+}
+
+impl Mechanism {
+    pub fn is_stable(&self) -> bool {
+        matches!(self, Self::Stable)
+    }
+
+    pub fn is_dynamic(&self) -> bool {
+        matches!(self, Self::Dynamic)
+    }
+}

pallets/subtensor/src/subnets/subnet.rs

@@ -1,6 +1,7 @@
 use super::*;
 use frame_support::IterableStorageMap;
 use sp_core::Get;
+use subnets::Mechanism;
 
 impl<T: Config> Pallet<T> {
     /// Fetches the total count of subnets.
@@ -81,7 +82,7 @@ impl<T: Config> Pallet<T> {
     /// # Returns:
     /// * 'u16': The subnet mechanism
     ///
-    pub fn get_subnet_mechanism(netuid: u16) -> u16 {
+    pub fn get_subnet_mechanism(netuid: u16) -> Mechanism {
         SubnetMechanism::<T>::get(netuid)
     }
 
@@ -122,7 +123,7 @@ impl<T: Config> Pallet<T> {
     pub fn do_register_network(
         origin: T::RuntimeOrigin,
         hotkey: &T::AccountId,
-        mechid: u16,
+        mechid: Mechanism,
         identity: Option<SubnetIdentityOf>,
     ) -> DispatchResult {
         // --- 1. Ensure the caller is a signed user.
@@ -135,7 +136,7 @@ impl<T: Config> Pallet<T> {
         );
 
         // --- 3. Ensure the mechanism is Dynamic.
-        ensure!(mechid == 1, Error::<T>::MechanismDoesNotExist);
+        ensure!(mechid.is_dynamic(), Error::<T>::MechanismDoesNotExist);
 
         // --- 4. Rate limit for network registrations.
         let current_block = Self::get_current_block_as_u64();