Add try_execute method to PrecompileExt

Author: ales-otf

precompiles/src/extensions.rs

@@ -4,9 +4,10 @@ use alloc::format;
 
 use frame_support::dispatch::{GetDispatchInfo, Pays, PostDispatchInfo};
 use frame_system::RawOrigin;
+use pallet_admin_utils::{PrecompileEnable, PrecompileEnum};
 use pallet_evm::{
-    AddressMapping, BalanceConverter, ExitError, GasWeightMapping, PrecompileFailure,
-    PrecompileHandle,
+    AddressMapping, BalanceConverter, ExitError, GasWeightMapping, Precompile, PrecompileFailure,
+    PrecompileHandle, PrecompileResult,
 };
 use precompile_utils::EvmResult;
 use sp_core::{H160, U256, blake2_256};
@@ -109,7 +110,7 @@ pub(crate) trait PrecompileHandleExt: PrecompileHandle {
 
 impl<T> PrecompileHandleExt for T where T: PrecompileHandle {}
 
-pub(crate) trait PrecompileExt<AccountId: From<[u8; 32]>> {
+pub(crate) trait PrecompileExt<AccountId: From<[u8; 32]>>: Precompile {
     const INDEX: u64;
 
     // ss58 public key i.e., the contract sends funds it received to the destination address from
@@ -127,6 +128,24 @@ pub(crate) trait PrecompileExt<AccountId: From<[u8; 32]>> {
 
         hash.into()
     }
+
+    fn try_execute<R>(
+        handle: &mut impl PrecompileHandle,
+        precompile_enum: PrecompileEnum,
+    ) -> Option<PrecompileResult>
+    where
+        R: frame_system::Config + pallet_admin_utils::Config,
+    {
+        if PrecompileEnable::<R>::get(&precompile_enum) {
+            Some(Self::execute(handle))
+        } else {
+            Some(Err(PrecompileFailure::Error {
+                exit_status: ExitError::Other(
+                    format!("Precompile {:?} is disabled", precompile_enum).into(),
+                ),
+            }))
+        }
+    }
 }
 
 #[cfg(test)]
@@ -152,4 +171,8 @@ mod test {
     impl PrecompileExt<AccountId32> for TestPrecompile {
         const INDEX: u64 = 2051;
     }
+
+    #[allow(unreachable_code)]
+    #[precompile_utils::precompile]
+    impl TestPrecompile {}
 }

precompiles/src/lib.rs

@@ -138,61 +138,25 @@ where
             }
             // Subtensor specific precompiles :
             a if a == hash(BalanceTransferPrecompile::<R>::INDEX) => {
-                if PrecompileEnable::<R>::get(PrecompileEnum::BalanceTransfer) {
-                    Some(BalanceTransferPrecompile::<R>::execute(handle))
-                } else {
-                    Some(Err(PrecompileFailure::Error {
-                        exit_status: ExitError::Other(
-                            "Precompile Balance Transfer is disabled".into(),
-                        ),
-                    }))
-                }
+                BalanceTransferPrecompile::<R>::try_execute::<R>(
+                    handle,
+                    PrecompileEnum::BalanceTransfer,
+                )
             }
             a if a == hash(StakingPrecompile::<R>::INDEX) => {
-                if PrecompileEnable::<R>::get(PrecompileEnum::Staking) {
-                    Some(StakingPrecompile::<R>::execute(handle))
-                } else {
-                    Some(Err(PrecompileFailure::Error {
-                        exit_status: ExitError::Other("Precompile Staking is disabled".into()),
-                    }))
-                }
+                StakingPrecompile::<R>::try_execute::<R>(handle, PrecompileEnum::Staking)
             }
             a if a == hash(StakingPrecompileV2::<R>::INDEX) => {
-                if PrecompileEnable::<R>::get(PrecompileEnum::Staking) {
-                    Some(StakingPrecompileV2::<R>::execute(handle))
-                } else {
-                    Some(Err(PrecompileFailure::Error {
-                        exit_status: ExitError::Other("Precompile Staking is disabled".into()),
-                    }))
-                }
+                StakingPrecompile::<R>::try_execute::<R>(handle, PrecompileEnum::Staking)
             }
-
             a if a == hash(SubnetPrecompile::<R>::INDEX) => {
-                if PrecompileEnable::<R>::get(PrecompileEnum::Subnet) {
-                    Some(SubnetPrecompile::<R>::execute(handle))
-                } else {
-                    Some(Err(PrecompileFailure::Error {
-                        exit_status: ExitError::Other("Precompile Subnet is disabled".into()),
-                    }))
-                }
+                SubnetPrecompile::<R>::try_execute::<R>(handle, PrecompileEnum::Subnet)
             }
             a if a == hash(MetagraphPrecompile::<R>::INDEX) => {
-                if PrecompileEnable::<R>::get(PrecompileEnum::Metagraph) {
-                    Some(MetagraphPrecompile::<R>::execute(handle))
-                } else {
-                    Some(Err(PrecompileFailure::Error {
-                        exit_status: ExitError::Other("Precompile Metagrah is disabled".into()),
-                    }))
-                }
+                MetagraphPrecompile::<R>::try_execute::<R>(handle, PrecompileEnum::Metagraph)
             }
             a if a == hash(NeuronPrecompile::<R>::INDEX) => {
-                if PrecompileEnable::<R>::get(PrecompileEnum::Neuron) {
-                    Some(NeuronPrecompile::<R>::execute(handle))
-                } else {
-                    Some(Err(PrecompileFailure::Error {
-                        exit_status: ExitError::Other("Precompile Neuron is disabled".into()),
-                    }))
-                }
+                NeuronPrecompile::<R>::try_execute::<R>(handle, PrecompileEnum::Neuron)
             }
             _ => None,
         }