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,35 @@ 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, vec![hotkey]);
+
+                match call_result {
+                    Ok(outcome) => {
+                        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

@@ -1358,6 +1358,67 @@ 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_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 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());
+    });
+}
+
 #[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);