feat(chain-extensions): add read-only WASM queries for subnet registration, coldkey lock, and stake availability

chain-extensions/Cargo.toml

@@ -21,6 +21,7 @@ sp-std.workspace = true
 codec = { workspace = true, features = ["derive"] }
 scale-info = { workspace = true, features = ["derive"] }
 subtensor-runtime-common.workspace = true
+subtensor-macros.workspace = true
 pallet-contracts.workspace = true
 pallet-subtensor.workspace = true
 pallet-subtensor-swap.workspace = true

chain-extensions/src/lib.rs

@@ -7,7 +7,7 @@ mod tests;
 
 pub mod types;
 
-use crate::types::{FunctionId, Output};
+use crate::types::{ColdkeyLock, FunctionId, Output, StakeAvailability, SubnetRegistrationState};
 use codec::{Decode, Encode, MaxEncodedLen};
 use frame_support::{DebugNoBound, traits::Get};
 use frame_system::RawOrigin;
@@ -595,6 +595,61 @@ where
 
                 Ok(RetVal::Converging(Output::Success as u32))
             }
+            FunctionId::GetSubnetRegistrationStateV1 => {
+                let netuid: NetUid = env
+                    .read_as()
+                    .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?;
+
+                let state = SubnetRegistrationState {
+                    netuid,
+                    exists: pallet_subtensor::Pallet::<T>::if_subnet_exist(netuid),
+                    registered_subnet_counter:
+                        pallet_subtensor::Pallet::<T>::get_registered_subnet_counter(netuid),
+                };
+
+                env.write_output(&state.encode())
+                    .map_err(|_| DispatchError::Other("Failed to write output"))?;
+
+                Ok(RetVal::Converging(Output::Success as u32))
+            }
+            FunctionId::GetColdkeyLockV1 => {
+                let (coldkey, netuid): (T::AccountId, NetUid) = env
+                    .read_as()
+                    .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?;
+
+                let lock =
+                    pallet_subtensor::Pallet::<T>::get_coldkey_lock(&coldkey, netuid).map(|lock| {
+                        ColdkeyLock {
+                            locked_mass: lock.locked_mass,
+                            conviction_bits: lock.conviction.to_bits(),
+                            last_update: lock.last_update,
+                        }
+                    });
+
+                env.write_output(&lock.encode())
+                    .map_err(|_| DispatchError::Other("Failed to write output"))?;
+
+                Ok(RetVal::Converging(Output::Success as u32))
+            }
+            FunctionId::GetStakeAvailabilityV1 => {
+                let (coldkey, netuid): (T::AccountId, NetUid) = env
+                    .read_as()
+                    .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?;
+
+                let (total, locked, available) =
+                    pallet_subtensor::Pallet::<T>::stake_availability(&coldkey, netuid);
+                let availability = StakeAvailability {
+                    netuid,
+                    total,
+                    locked,
+                    available,
+                };
+
+                env.write_output(&availability.encode())
+                    .map_err(|_| DispatchError::Other("Failed to write output"))?;
+
+                Ok(RetVal::Converging(Output::Success as u32))
+            }
             FunctionId::AddStakeV1 => {
                 let origin = RawOrigin::Signed(env.caller());
                 Self::dispatch_add_stake_v1(env, origin)

chain-extensions/src/tests.rs

@@ -1,7 +1,7 @@
 #![allow(clippy::unwrap_used)]
 
 use super::{SubtensorChainExtension, SubtensorExtensionEnv, mock};
-use crate::types::{FunctionId, Output};
+use crate::types::{ColdkeyLock, FunctionId, Output, StakeAvailability, SubnetRegistrationState};
 use codec::{Decode, Encode};
 use frame_support::pallet_prelude::Zero;
 use frame_support::{assert_ok, weights::Weight};
@@ -1625,6 +1625,300 @@ fn get_alpha_price_returns_encoded_price() {
     });
 }
 
+#[test]
+fn get_subnet_registration_state_returns_existing_subnet_counter() {
+    mock::new_test_ext(1).execute_with(|| {
+        let owner_hotkey = U256::from(8101);
+        let owner_coldkey = U256::from(8102);
+        let caller = U256::from(8103);
+
+        let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
+        let expected_counter =
+            pallet_subtensor::Pallet::<mock::Test>::get_registered_subnet_counter(netuid);
+
+        let mut env = MockEnv::new(
+            FunctionId::GetSubnetRegistrationStateV1,
+            caller,
+            netuid.encode(),
+        );
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert!(env.charged_weight().is_none());
+
+        let state = SubnetRegistrationState::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(
+            state,
+            SubnetRegistrationState {
+                netuid,
+                exists: true,
+                registered_subnet_counter: expected_counter,
+            }
+        );
+    });
+}
+
+#[test]
+fn get_subnet_registration_state_preserves_counter_after_dissolve() {
+    mock::new_test_ext(1).execute_with(|| {
+        let owner_hotkey = U256::from(8201);
+        let owner_coldkey = U256::from(8202);
+        let caller = U256::from(8203);
+
+        let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
+        let registered_counter =
+            pallet_subtensor::Pallet::<mock::Test>::get_registered_subnet_counter(netuid);
+
+        assert_ok!(pallet_subtensor::Pallet::<mock::Test>::do_dissolve_network(
+            netuid
+        ));
+
+        let mut env = MockEnv::new(
+            FunctionId::GetSubnetRegistrationStateV1,
+            caller,
+            netuid.encode(),
+        );
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+
+        let state = SubnetRegistrationState::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(state.netuid, netuid);
+        assert!(!state.exists);
+        assert_eq!(state.registered_subnet_counter, registered_counter);
+    });
+}
+
+#[test]
+fn get_subnet_registration_state_detects_reused_netuid_generation() {
+    mock::new_test_ext(1).execute_with(|| {
+        let first_hotkey = U256::from(8301);
+        let first_coldkey = U256::from(8302);
+        let second_hotkey = U256::from(8303);
+        let second_coldkey = U256::from(8304);
+        let caller = U256::from(8305);
+
+        let netuid = mock::add_dynamic_network(&first_hotkey, &first_coldkey);
+        let first_counter =
+            pallet_subtensor::Pallet::<mock::Test>::get_registered_subnet_counter(netuid);
+
+        assert_ok!(pallet_subtensor::Pallet::<mock::Test>::do_dissolve_network(
+            netuid
+        ));
+        let reused_netuid = mock::add_dynamic_network(&second_hotkey, &second_coldkey);
+        assert_eq!(reused_netuid, netuid);
+
+        let mut env = MockEnv::new(
+            FunctionId::GetSubnetRegistrationStateV1,
+            caller,
+            netuid.encode(),
+        );
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+
+        let state = SubnetRegistrationState::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(state.netuid, netuid);
+        assert!(state.exists);
+        assert!(state.registered_subnet_counter > first_counter);
+    });
+}
+
+#[test]
+fn get_coldkey_lock_returns_none_without_lock() {
+    mock::new_test_ext(1).execute_with(|| {
+        let caller = U256::from(8401);
+        let coldkey = U256::from(8402);
+        let netuid = NetUid::from(1);
+
+        let mut env = MockEnv::new(
+            FunctionId::GetColdkeyLockV1,
+            caller,
+            (coldkey, netuid).encode(),
+        );
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert!(env.charged_weight().is_none());
+
+        let lock = Option::<ColdkeyLock>::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(lock, None);
+    });
+}
+
+#[test]
+fn get_coldkey_lock_returns_rolled_lock() {
+    mock::new_test_ext(1).execute_with(|| {
+        let owner_hotkey = U256::from(8501);
+        let owner_coldkey = U256::from(8502);
+        let coldkey = U256::from(8503);
+        let hotkey = U256::from(8504);
+        let stake = AlphaBalance::from(10_000u64);
+        let lock_amount = AlphaBalance::from(5_000u64);
+
+        let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey, &coldkey, netuid, stake,
+        );
+
+        assert_ok!(pallet_subtensor::Pallet::<mock::Test>::do_lock_stake(
+            &coldkey,
+            netuid,
+            &hotkey,
+            lock_amount,
+        ));
+
+        frame_system::Pallet::<mock::Test>::set_block_number(1_001);
+        let expected =
+            pallet_subtensor::Pallet::<mock::Test>::get_coldkey_lock(&coldkey, netuid).unwrap();
+
+        let mut env = MockEnv::new(
+            FunctionId::GetColdkeyLockV1,
+            coldkey,
+            (coldkey, netuid).encode(),
+        );
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert!(env.charged_weight().is_none());
+
+        let lock = Option::<ColdkeyLock>::decode(&mut &env.output()[..])
+            .unwrap()
+            .unwrap();
+        assert_eq!(lock.locked_mass, expected.locked_mass);
+        assert_eq!(lock.conviction_bits, expected.conviction.to_bits());
+        assert!(lock.conviction_bits > 0);
+        assert_eq!(
+            lock.last_update,
+            pallet_subtensor::Pallet::<mock::Test>::get_current_block_as_u64()
+        );
+    });
+}
+
+#[test]
+fn get_stake_availability_returns_zeroes_without_stake_or_lock() {
+    mock::new_test_ext(1).execute_with(|| {
+        let caller = U256::from(8601);
+        let coldkey = U256::from(8602);
+        let netuid = NetUid::from(1);
+
+        let mut env = MockEnv::new(
+            FunctionId::GetStakeAvailabilityV1,
+            caller,
+            (coldkey, netuid).encode(),
+        );
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert!(env.charged_weight().is_none());
+
+        let availability = StakeAvailability::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(
+            availability,
+            StakeAvailability {
+                netuid,
+                total: AlphaBalance::ZERO,
+                locked: AlphaBalance::ZERO,
+                available: AlphaBalance::ZERO,
+            }
+        );
+    });
+}
+
+#[test]
+fn get_stake_availability_returns_partial_lock_breakdown() {
+    mock::new_test_ext(1).execute_with(|| {
+        let owner_hotkey = U256::from(8701);
+        let owner_coldkey = U256::from(8702);
+        let coldkey = U256::from(8703);
+        let hotkey = U256::from(8704);
+        let stake = AlphaBalance::from(10_000u64);
+        let lock_amount = AlphaBalance::from(4_000u64);
+
+        let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey, &coldkey, netuid, stake,
+        );
+        assert_ok!(pallet_subtensor::Pallet::<mock::Test>::do_lock_stake(
+            &coldkey,
+            netuid,
+            &hotkey,
+            lock_amount,
+        ));
+
+        let mut env = MockEnv::new(
+            FunctionId::GetStakeAvailabilityV1,
+            coldkey,
+            (coldkey, netuid).encode(),
+        );
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert!(env.charged_weight().is_none());
+
+        let availability = StakeAvailability::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(
+            availability,
+            StakeAvailability {
+                netuid,
+                total: stake,
+                locked: lock_amount,
+                available: stake.saturating_sub(lock_amount),
+            }
+        );
+    });
+}
+
+#[test]
+fn get_stake_availability_uses_rolled_forward_lock() {
+    mock::new_test_ext(1).execute_with(|| {
+        let owner_hotkey = U256::from(8801);
+        let owner_coldkey = U256::from(8802);
+        let coldkey = U256::from(8803);
+        let hotkey = U256::from(8804);
+        let stake = AlphaBalance::from(10_000u64);
+        let lock_amount = AlphaBalance::from(5_000u64);
+
+        let netuid = mock::add_dynamic_network(&owner_hotkey, &owner_coldkey);
+        mock::register_ok_neuron(netuid, hotkey, coldkey, 0);
+        pallet_subtensor::Pallet::<mock::Test>::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey, &coldkey, netuid, stake,
+        );
+        assert_ok!(pallet_subtensor::Pallet::<mock::Test>::do_lock_stake(
+            &coldkey,
+            netuid,
+            &hotkey,
+            lock_amount,
+        ));
+
+        frame_system::Pallet::<mock::Test>::set_block_number(1_001);
+        let expected_locked =
+            pallet_subtensor::Pallet::<mock::Test>::get_current_locked(&coldkey, netuid);
+        assert!(expected_locked < lock_amount);
+
+        let mut env = MockEnv::new(
+            FunctionId::GetStakeAvailabilityV1,
+            coldkey,
+            (coldkey, netuid).encode(),
+        );
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert!(env.charged_weight().is_none());
+
+        let availability = StakeAvailability::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(availability.netuid, netuid);
+        assert_eq!(availability.total, stake);
+        assert_eq!(availability.locked, expected_locked);
+        assert_eq!(
+            availability.available,
+            stake.saturating_sub(expected_locked)
+        );
+    });
+}
+
 /// `Caller*` dispatch uses `env.origin()` via `convert_origin`; with [`MockEnv`] both match
 /// `Signed(caller)`, so outcomes align with non-`Caller` arms. Weight expectations match the shared
 /// `dispatch_*_v1` helpers used by each pair.