Refactor validate_block

pallets/parachain-system/proc-macro/src/lib.rs

@@ -115,14 +115,19 @@ pub fn register_validate_block(input: proc_macro::TokenStream) -> proc_macro::To
 
 				#[no_mangle]
 				unsafe fn validate_block(arguments: *mut u8, arguments_len: usize) -> u64 {
-					let params = #crate_::validate_block::polkadot_parachain::load_params(
-						arguments,
-						arguments_len,
+					// We convert the `arguments` into a boxed slice and then into `Bytes`.
+					let args = #crate_::validate_block::sp_std::boxed::Box::from_raw(
+						#crate_::validate_block::sp_std::slice::from_raw_parts_mut(
+							arguments,
+							arguments_len,
+						)
 					);
+					let args = #crate_::validate_block::bytes::Bytes::from(args);
 
-					// Memory is scarce resource, so let's free the `arguments` as we don't
-					// need them anymore.
-					#crate_::validate_block::sp_io::allocator::free(arguments);
+					// Then we decode from these bytes the `MemoryOptimizedValidationParams`.
+					let params = #crate_::validate_block::decode_from_bytes::<
+						#crate_::validate_block::MemoryOptimizedValidationParams
+					>(args).expect("Invalid arguments to `validate_block`.");
 
 					let res = #crate_::validate_block::implementation::validate_block::<
 						<#runtime as #crate_::validate_block::GetRuntimeBlockType>::RuntimeBlock,

pallets/parachain-system/src/validate_block/implementation.rs

@@ -16,9 +16,7 @@
 
 //! The actual implementation of the validate block functionality.
 
-use frame_support::traits::{ExecuteBlock, ExtrinsicCall, Get, IsSubType};
-use sp_runtime::traits::{Block as BlockT, Extrinsic, HashFor, Header as HeaderT};
-
+use super::MemoryOptimizedValidationParams;
 use cumulus_primitives_core::{
 	relay_chain::v2::Hash as RHash, ParachainBlockData, PersistedValidationData,
 };
@@ -30,9 +28,11 @@ use polkadot_parachain::primitives::{
 
 use codec::{Decode, Encode};
 
+use frame_support::traits::{ExecuteBlock, ExtrinsicCall, Get, IsSubType};
 use sp_core::storage::{ChildInfo, StateVersion};
 use sp_externalities::{set_and_run_with_externalities, Externalities};
 use sp_io::KillStorageResult;
+use sp_runtime::traits::{Block as BlockT, Extrinsic, HashFor, Header as HeaderT};
 use sp_std::{mem, prelude::*};
 use sp_trie::MemoryDB;
 
@@ -77,39 +77,41 @@ pub fn validate_block<
 	PSC: crate::Config,
 	CI: crate::CheckInherents<B>,
 >(
-	params: ValidationParams,
+	MemoryOptimizedValidationParams {
+		block_data,
+		parent_head,
+		relay_parent_number,
+		relay_parent_storage_root,
+	}: MemoryOptimizedValidationParams,
 ) -> ValidationResult
 where
 	B::Extrinsic: ExtrinsicCall,
 	<B::Extrinsic as Extrinsic>::Call: IsSubType<crate::Call<PSC>>,
 {
-	let block_data = ParachainBlockData::<B>::decode(&mut &params.block_data.0[..])
+	let block_data = codec::decode_from_bytes::<ParachainBlockData<B>>(block_data)
 		.expect("Invalid parachain block data");
 
-	let parent_head =
-		B::Header::decode(&mut &params.parent_head.0[..]).expect("Invalid parent head");
-
-	// We decoded the block data and don't need it anymore. So, we release it to free the memory.
-	mem::drop(params.block_data);
+	let parent_header =
+		codec::decode_from_bytes::<B::Header>(parent_head.clone()).expect("Invalid parent head");
 
 	let (header, extrinsics, storage_proof) = block_data.deconstruct();
 
 	let head_data = HeadData(header.encode());
 
 	let block = B::new(header, extrinsics);
-	assert!(parent_head.hash() == *block.header().parent_hash(), "Invalid parent hash");
+	assert!(parent_header.hash() == *block.header().parent_hash(), "Invalid parent hash");
 
 	let inherent_data = extract_parachain_inherent_data(&block);
 
 	validate_validation_data(
 		&inherent_data.validation_data,
-		params.relay_parent_number,
-		params.relay_parent_storage_root,
-		params.parent_head,
+		relay_parent_number,
+		relay_parent_storage_root,
+		parent_head,
 	);
 
 	// Create the db
-	let db = match storage_proof.to_memory_db(Some(parent_head.state_root())) {
+	let db = match storage_proof.to_memory_db(Some(parent_header.state_root())) {
 		Ok((db, _)) => db,
 		Err(_) => panic!("Compact proof decoding failure."),
 	};
@@ -118,7 +120,8 @@ where
 
 	// We use the storage root of the `parent_head` to ensure that it is the correct root.
 	// This is already being done above while creating the in-memory db, but let's be paranoid!!
-	let backend = sp_state_machine::TrieBackendBuilder::new(db, *parent_head.state_root()).build();
+	let backend =
+		sp_state_machine::TrieBackendBuilder::new(db, *parent_header.state_root()).build();
 
 	let _guard = (
 		// Replace storage calls with our own implementations
@@ -236,9 +239,9 @@ fn validate_validation_data(
 	validation_data: &PersistedValidationData,
 	relay_parent_number: RelayChainBlockNumber,
 	relay_parent_storage_root: RHash,
-	parent_head: HeadData,
+	parent_head: bytes::Bytes,
 ) {
-	assert_eq!(parent_head, validation_data.parent_head, "Parent head doesn't match");
+	assert_eq!(parent_head, validation_data.parent_head.0, "Parent head doesn't match");
 	assert_eq!(
 		relay_parent_number, validation_data.relay_parent_number,
 		"Relay parent number doesn't match",

pallets/parachain-system/src/validate_block/mod.rs

@@ -27,7 +27,32 @@ mod tests;
 pub use polkadot_parachain;
 #[cfg(not(feature = "std"))]
 #[doc(hidden)]
-pub use sp_io;
+pub use sp_std;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use bytes;
+#[cfg(not(feature = "std"))]
+#[doc(hidden)]
+pub use codec::decode_from_bytes;
 #[cfg(not(feature = "std"))]
 #[doc(hidden)]
 pub use sp_runtime::traits::GetRuntimeBlockType;
+
+/// Basically the same as [`ValidationParams`](polkadot_parachain::primitives::ValidationParams),
+/// but a little bit optimized for our use case here.
+///
+/// `block_data` and `head_data` are represented as [`bytes::Bytes`] to make them reuse
+/// the memory of the input parameter of the exported `validate_blocks` function.
+///
+/// The layout of this type must match exactly the layout of
+/// [`ValidationParams`](polkadot_parachain::primitives::ValidationParams) to have the same
+/// SCALE encoding.
+#[derive(codec::Decode)]
+#[cfg_attr(feature = "std", derive(codec::Encode))]
+#[doc(hidden)]
+pub struct MemoryOptimizedValidationParams {
+	pub parent_head: bytes::Bytes,
+	pub block_data: bytes::Bytes,
+	pub relay_parent_number: cumulus_primitives_core::relay_chain::v2::BlockNumber,
+	pub relay_parent_storage_root: cumulus_primitives_core::relay_chain::v2::Hash,
+}

pallets/parachain-system/src/validate_block/tests.rs

@@ -14,7 +14,7 @@
 // You should have received a copy of the GNU General Public License
 // along with Cumulus.  If not, see <http://www.gnu.org/licenses/>.
 
-use codec::{Decode, Encode};
+use codec::{Decode, Encode, DecodeAll};
 use cumulus_primitives_core::{ParachainBlockData, PersistedValidationData};
 use cumulus_test_client::{
 	generate_extrinsic,
@@ -27,6 +27,8 @@ use sp_keyring::AccountKeyring::*;
 use sp_runtime::traits::Header as HeaderT;
 use std::{env, process::Command};
 
+use crate::validate_block::MemoryOptimizedValidationParams;
+
 fn call_validate_block_encoded_header(
 	parent_head: Header,
 	block_data: ParachainBlockData<Block>,
@@ -289,3 +291,33 @@ fn check_inherents_are_unsigned_and_before_all_other_extrinsics() {
 			.contains("Could not find `set_validation_data` inherent"));
 	}
 }
+
+/// Test that ensures that `ValidationParams` and `MemoryOptimizedValidationParams`
+/// are encoding/decoding.
+#[test]
+fn validation_params_and_memory_optimized_validation_params_encode_and_decode() {
+	const BLOCK_DATA: &[u8] = &[1, 2, 3, 4, 5];
+	const PARENT_HEAD: &[u8] = &[1, 3, 4, 5, 6, 7, 9];
+
+
+	let validation_params =
+		ValidationParams {
+			block_data: BlockData(BLOCK_DATA.encode()),
+			parent_head: HeadData(PARENT_HEAD.encode()),
+			relay_parent_number: 1,
+			relay_parent_storage_root: Hash::random(),
+		};
+
+	let encoded = validation_params.encode();
+
+	let decoded = MemoryOptimizedValidationParams::decode_all(&mut &encoded[..]).unwrap();
+	assert_eq!(decoded.relay_parent_number, validation_params.relay_parent_number);
+	assert_eq!(decoded.relay_parent_storage_root, validation_params.relay_parent_storage_root);
+	assert_eq!(decoded.block_data, validation_params.block_data.0);
+	assert_eq!(decoded.parent_head, validation_params.parent_head.0);
+
+	let encoded = decoded.encode();
+
+	let decoded = ValidationParams::decode_all(&mut &encoded[..]).unwrap();
+	assert_eq!(decoded, validation_params);
+}