feat(chain-extensions): add claim-root-with-hotkey for contract root-basket claims

chain-extensions/src/lib.rs

@@ -19,6 +19,7 @@ use pallet_subtensor_proxy as pallet_proxy;
 use pallet_subtensor_proxy::WeightInfo;
 use sp_runtime::{DispatchError, Weight, traits::StaticLookup};
 use sp_std::marker::PhantomData;
+use sp_std::vec;
 use substrate_fixed::types::U64F64;
 use subtensor_runtime_common::{AlphaBalance, NetUid, ProxyType, TaoBalance};
 use subtensor_swap_interface::SwapHandler;
@@ -964,6 +965,36 @@ where
                     }
                 }
             }
+            FunctionId::ClaimRootWithHotkeyV1 => {
+                let hotkey: T::AccountId = env
+                    .read_as()
+                    .map_err(|_| DispatchError::Other("Failed to decode input parameters"))?;
+
+                let rows = pallet_subtensor::Pallet::<T>::get_basket_holdings(&hotkey).len();
+                if rows > pallet_subtensor::MAX_ROOT_CLAIM_WORK as usize {
+                    return Ok(RetVal::Converging(Output::RuntimeError as u32));
+                }
+
+                let weight = <<T as pallet_subtensor::Config>::WeightInfo as SubtensorWeightInfo>::claim_root(pallet_subtensor::MAX_ROOT_CLAIM_WORK);
+                env.charge_weight(weight)?;
+
+                let caller = env.caller();
+                let call_result =
+                    pallet_subtensor::Pallet::<T>::do_root_claim(caller.clone(), vec![hotkey]);
+
+                match call_result {
+                    Ok(outcome) => {
+                        pallet_subtensor::Pallet::<T>::maybe_add_coldkey_index(&caller);
+                        env.write_output(&outcome.tao.encode())
+                            .map_err(|_| DispatchError::Other("Failed to write output"))?;
+                        Ok(RetVal::Converging(Output::Success as u32))
+                    }
+                    Err(e) => {
+                        let error_code = Output::from(e) as u32;
+                        Ok(RetVal::Converging(error_code))
+                    }
+                }
+            }
         }
     }
 }

chain-extensions/src/tests.rs

@@ -12,7 +12,7 @@ use pallet_subtensor::weights::WeightInfo as SubtensorWeightInfo;
 use sp_core::Get;
 use sp_core::U256;
 use sp_runtime::DispatchError;
-use substrate_fixed::types::U64F64;
+use substrate_fixed::types::{I96F32, U64F64};
 use subtensor_runtime_common::{AlphaBalance, NetUid, TaoBalance, Token};
 use subtensor_swap_interface::SwapHandler;
 
@@ -1358,6 +1358,199 @@ fn assert_success(ret: RetVal) {
     }
 }
 
+#[test]
+fn claim_root_with_hotkey_noop_returns_zero() {
+    mock::new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(61001);
+        let hotkey = U256::from(61002);
+
+        let expected_weight = <<mock::Test as pallet_subtensor::Config>::WeightInfo as SubtensorWeightInfo>::claim_root(
+            pallet_subtensor::MAX_ROOT_CLAIM_WORK,
+        );
+
+        let mut env = MockEnv::new(
+            FunctionId::ClaimRootWithHotkeyV1,
+            coldkey,
+            hotkey.encode(),
+        )
+        .with_expected_weight(expected_weight);
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert_eq!(env.charged_weight(), Some(expected_weight));
+
+        let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(tao, 0);
+    });
+}
+
+#[test]
+fn claim_root_with_hotkey_noop_indexes_fresh_coldkey() {
+    mock::new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(61211);
+        let hotkey = U256::from(61212);
+
+        assert!(!pallet_subtensor::StakingColdkeys::<mock::Test>::contains_key(coldkey));
+        let num_before = pallet_subtensor::NumStakingColdkeys::<mock::Test>::get();
+
+        let expected_weight = <<mock::Test as pallet_subtensor::Config>::WeightInfo as SubtensorWeightInfo>::claim_root(
+            pallet_subtensor::MAX_ROOT_CLAIM_WORK,
+        );
+
+        let mut env = MockEnv::new(
+            FunctionId::ClaimRootWithHotkeyV1,
+            coldkey,
+            hotkey.encode(),
+        )
+        .with_expected_weight(expected_weight);
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+        assert_eq!(env.charged_weight(), Some(expected_weight));
+
+        let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(tao, 0);
+
+        assert!(pallet_subtensor::StakingColdkeys::<mock::Test>::contains_key(coldkey));
+        assert_eq!(
+            pallet_subtensor::NumStakingColdkeys::<mock::Test>::get(),
+            num_before.saturating_add(1)
+        );
+        let idx = pallet_subtensor::StakingColdkeys::<mock::Test>::get(coldkey).unwrap();
+        assert_eq!(
+            pallet_subtensor::StakingColdkeysByIndex::<mock::Test>::get(idx),
+            Some(coldkey)
+        );
+    });
+}
+
+#[test]
+fn claim_root_with_hotkey_repeat_claim_preserves_index() {
+    mock::new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(61221);
+        let hotkey = U256::from(61222);
+
+        pallet_subtensor::Pallet::<mock::Test>::maybe_add_coldkey_index(&coldkey);
+        let num_before = pallet_subtensor::NumStakingColdkeys::<mock::Test>::get();
+
+        let expected_weight = <<mock::Test as pallet_subtensor::Config>::WeightInfo as SubtensorWeightInfo>::claim_root(
+            pallet_subtensor::MAX_ROOT_CLAIM_WORK,
+        );
+
+        let mut env = MockEnv::new(
+            FunctionId::ClaimRootWithHotkeyV1,
+            coldkey,
+            hotkey.encode(),
+        )
+        .with_expected_weight(expected_weight);
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+
+        let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(tao, 0);
+
+        assert_eq!(pallet_subtensor::NumStakingColdkeys::<mock::Test>::get(), num_before);
+        let idx = pallet_subtensor::StakingColdkeys::<mock::Test>::get(coldkey).unwrap();
+        assert_eq!(
+            pallet_subtensor::StakingColdkeysByIndex::<mock::Test>::get(idx),
+            Some(coldkey)
+        );
+    });
+}
+
+#[test]
+fn claim_root_with_hotkey_rejects_basket_above_envelope() {
+    mock::new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(61101);
+        let hotkey = U256::from(61102);
+
+        // Seed the validator's basket with one escrow holding per netuid, 257 rows
+        // > MAX_ROOT_CLAIM_WORK (256), using the same helper the claim engine uses
+        // to build escrow holdings.
+        let escrow = pallet_subtensor::Pallet::<mock::Test>::get_beta_escrow_account_id();
+        for i in 0..=pallet_subtensor::MAX_ROOT_CLAIM_WORK {
+            pallet_subtensor::Pallet::<mock::Test>::increase_stake_for_hotkey_and_coldkey_on_subnet(
+                &hotkey,
+                &escrow,
+                NetUid::from(i as u16),
+                1u64.into(),
+            );
+        }
+        assert!(
+            pallet_subtensor::Pallet::<mock::Test>::get_basket_holdings(&hotkey).len()
+                > pallet_subtensor::MAX_ROOT_CLAIM_WORK as usize
+        );
+        let num_before = pallet_subtensor::NumStakingColdkeys::<mock::Test>::get();
+
+        let mut env = MockEnv::new(FunctionId::ClaimRootWithHotkeyV1, coldkey, hotkey.encode());
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        match ret {
+            RetVal::Converging(code) => assert_eq!(code, Output::RuntimeError as u32),
+            _ => panic!("expected converging error code"),
+        }
+
+        assert!(env.charged_weight().is_none());
+
+        // Rejection must not mutate the index.
+        assert!(!pallet_subtensor::StakingColdkeys::<mock::Test>::contains_key(coldkey));
+        assert_eq!(
+            pallet_subtensor::NumStakingColdkeys::<mock::Test>::get(),
+            num_before
+        );
+    });
+}
+
+#[test]
+fn claim_root_with_hotkey_payout_indexes_coldkey() {
+    mock::new_test_ext(1).execute_with(|| {
+        let coldkey = U256::from(61201);
+        let hotkey = U256::from(61202);
+
+        pallet_subtensor::Pallet::<mock::Test>::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &coldkey,
+            NetUid::ROOT,
+            1u64.into(),
+        );
+        let escrow = pallet_subtensor::Pallet::<mock::Test>::get_beta_escrow_account_id();
+        pallet_subtensor::Pallet::<mock::Test>::increase_stake_for_hotkey_and_coldkey_on_subnet(
+            &hotkey,
+            &escrow,
+            NetUid::ROOT,
+            1_000_000u64.into(),
+        );
+        pallet_subtensor::BasketShares::<mock::Test>::insert(hotkey, 1u64);
+        pallet_subtensor::BasketRate::<mock::Test>::insert(hotkey, I96F32::from_num(1));
+        pallet_subtensor::RootClaimableThreshold::<mock::Test>::insert(
+            NetUid::ROOT,
+            I96F32::from_num(0),
+        );
+
+        assert!(!pallet_subtensor::StakingColdkeys::<mock::Test>::contains_key(coldkey));
+
+        let expected_weight = <<mock::Test as pallet_subtensor::Config>::WeightInfo as SubtensorWeightInfo>::claim_root(
+            pallet_subtensor::MAX_ROOT_CLAIM_WORK,
+        );
+
+        let mut env = MockEnv::new(
+            FunctionId::ClaimRootWithHotkeyV1,
+            coldkey,
+            hotkey.encode(),
+        )
+        .with_expected_weight(expected_weight);
+
+        let ret = SubtensorChainExtension::<mock::Test>::dispatch(&mut env).unwrap();
+        assert_success(ret);
+
+        let tao: u64 = Decode::decode(&mut &env.output()[..]).unwrap();
+        assert_eq!(tao, 1_000_000);
+
+        assert!(pallet_subtensor::StakingColdkeys::<mock::Test>::contains_key(coldkey));
+    });
+}
+
 #[test]
 fn add_stake_recycle_rollback_on_recycle_failure() {
     mock::new_test_ext(1).execute_with(|| {

chain-extensions/src/types.rs

@@ -46,6 +46,7 @@ pub enum FunctionId {
     GetStakeAvailabilityV1 = 36,
     MoveStakeLimitV1 = 37,
     CallerMoveStakeLimitV1 = 38,
+    ClaimRootWithHotkeyV1 = 39,
 }
 
 #[freeze_struct("5dc33d60abed5c08")]
@@ -192,11 +193,12 @@ mod function_id_tests {
         assert_eq!(FunctionId::GetStakeAvailabilityV1 as u16, 36);
         assert_eq!(FunctionId::MoveStakeLimitV1 as u16, 37);
         assert_eq!(FunctionId::CallerMoveStakeLimitV1 as u16, 38);
+        assert_eq!(FunctionId::ClaimRootWithHotkeyV1 as u16, 39);
     }
 
     #[test]
     fn caller_ids_roundtrip_try_from_primitive() {
-        for id in 16u16..=38u16 {
+        for id in 16u16..=39u16 {
             let v = FunctionId::try_from_primitive(id)
                 .unwrap_or_else(|_| panic!("try_from_primitive failed for {id}"));
             assert_eq!(v as u16, id);