refactor: optimize read state handler

runtime/devnet/src/config/api.rs

@@ -1,6 +1,7 @@
 use crate::{config::assets::TrustBackedAssetsInstance, fungibles, Runtime, RuntimeCall};
 use codec::{Decode, Encode, MaxEncodedLen};
 use frame_support::traits::Contains;
+use pop_runtime_extension::ReadStateHandler;
 
 /// A query of runtime state.
 #[derive(Encode, Decode, Debug, MaxEncodedLen)]
@@ -11,6 +12,14 @@ pub enum RuntimeRead {
 	Fungibles(fungibles::Read<Runtime>),
 }
 
+impl ReadStateHandler<Runtime> for RuntimeRead {
+	fn handle_read(read: RuntimeRead) -> sp_std::vec::Vec<u8> {
+		match read {
+			RuntimeRead::Fungibles(key) => fungibles::Pallet::read_state(key),
+		}
+	}
+}
+
 /// A type to identify allowed calls to the Runtime from the API.
 pub struct AllowedApiCalls;
 

runtime/devnet/src/config/contracts.rs

@@ -8,6 +8,8 @@ use frame_support::{
 };
 use frame_system::pallet_prelude::BlockNumberFor;
 
+use super::api::{AllowedApiCalls, RuntimeRead};
+
 pub enum AllowBalancesCall {}
 
 impl frame_support::traits::Contains<RuntimeCall> for AllowBalancesCall {
@@ -44,6 +46,11 @@ parameter_types! {
 	pub const CodeHashLockupDepositPercent: Perbill = Perbill::from_percent(0);
 }
 
+impl pop_runtime_extension::Config for Runtime {
+	type RuntimeRead = RuntimeRead;
+	type AllowedApiCalls = AllowedApiCalls;
+}
+
 impl pallet_contracts::Config for Runtime {
 	type Time = Timestamp;
 	type Randomness = DummyRandomness<Self>;

runtime/devnet/src/config/extension.rs

@@ -1,39 +1,5 @@
-use crate::{
-	config::api::AllowedApiCalls,
-	fungibles::{self},
-	Runtime,
-};
-use frame_support::{ensure, traits::Contains};
-use pop_runtime_extension::{constants::UNKNOWN_CALL_ERROR, decode_checked, StateReadHandler};
-
-use sp_runtime::DispatchError;
-use sp_std::vec::Vec;
-
-use super::api::RuntimeRead;
-
-pub struct ContractExecutionContext;
-
-impl StateReadHandler for ContractExecutionContext {
-	fn handle_params(params: &[u8]) -> Result<Vec<u8>, DispatchError> {
-		let read = decode_checked::<RuntimeRead>(&mut &params[..])?;
-		ensure!(AllowedApiCalls::contains(&read), UNKNOWN_CALL_ERROR);
-		let result = match read {
-			RuntimeRead::Fungibles(key) => fungibles::Pallet::read_state(key),
-		};
-		Ok(result)
-	}
-}
-
-impl pop_runtime_extension::Config for Runtime {
-	type StateReadHandler = ContractExecutionContext;
-	type AllowedDispatchCalls = AllowedApiCalls;
-}
-
 #[cfg(test)]
 mod tests {
-
-	use super::*;
-
 	use crate::{config::assets::TrustBackedAssetsInstance, Assets, Runtime, System};
 	use codec::{Decode, Encode};
 	use sp_runtime::{

runtime/extension/src/lib.rs

@@ -23,17 +23,17 @@ type ContractSchedule<T> = <T as pallet_contracts::Config>::Schedule;
 pub trait Config:
 	frame_system::Config<RuntimeCall: GetDispatchInfo + Dispatchable<PostInfo = PostDispatchInfo>>
 {
-	/// Receives parameter byte array from the environment buffer and handles the logic.
-	type StateReadHandler: StateReadHandler;
-	/// Whitelisting list of runtime calls that can be dispatched.
-	type AllowedDispatchCalls: Contains<Self::RuntimeCall>;
+	/// A query of runtime state.
+	type RuntimeRead: Decode + ReadStateHandler<Self>;
+	/// Whitelisting list of runtime calls and read state calls.
+	type AllowedApiCalls: Contains<Self::RuntimeCall> + Contains<Self::RuntimeRead>;
 }
 
 /// Trait for handling parameters from the chain extension environment during state read operations.
-pub trait StateReadHandler {
+pub trait ReadStateHandler<T: Config> {
 	/// Processes the parameters needed to execute a call within the runtime environment.
 	/// Layout of `params`: [version, pallet_index, call_index, ...Vec<u8>].
-	fn handle_params(params: &[u8]) -> Result<Vec<u8>, DispatchError>;
+	fn handle_read(params: T::RuntimeRead) -> Vec<u8>;
 }
 
 #[derive(Default)]
@@ -130,13 +130,17 @@ where
 	params.insert(0, version);
 	params.insert(1, pallet_index);
 	params.insert(2, call_index);
-	let (mut encoded_version, encoded_read) = (&params[..1], &params[1..]);
+	let (mut encoded_version, mut encoded_read) = (&params[..1], &params[1..]);
 	let version = decode_checked::<VersionedStateRead>(&mut encoded_version)?;
 
 	// Charge weight for doing one storage read.
 	env.charge_weight(T::DbWeight::get().reads(1_u64))?;
 	let result = match version {
-		VersionedStateRead::V0 => T::StateReadHandler::handle_params(encoded_read)?,
+		VersionedStateRead::V0 => {
+			let read = decode_checked::<T::RuntimeRead>(&mut encoded_read)?;
+			ensure!(T::AllowedApiCalls::contains(&read), UNKNOWN_CALL_ERROR);
+			T::RuntimeRead::handle_read(read)
+		},
 	};
 	log::trace!(
 		target:LOG_TARGET,
@@ -182,7 +186,7 @@ where
 {
 	let charged_dispatch_weight = env.charge_weight(call.get_dispatch_info().weight)?;
 	log::debug!(target:LOG_TARGET, "{} Inputted RuntimeCall: {:?}", log_prefix, call);
-	origin.add_filter(T::AllowedDispatchCalls::contains);
+	origin.add_filter(T::AllowedApiCalls::contains);
 	match call.dispatch(origin) {
 		Ok(info) => {
 			log::debug!(target:LOG_TARGET, "{} success, actual weight: {:?}", log_prefix, info.actual_weight);