Ignore trie nodes while recording a proof

cumulus/client/consensus/proposer/Cargo.toml

@@ -17,6 +17,11 @@ async-trait = { workspace = true }
 thiserror = { workspace = true }
 
 # Substrate
+sc-basic-authorship = { workspace = true }
+sc-block-builder = { workspace = true }
+sc-transaction-pool-api = { workspace = true }
+sp-api = { workspace = true, default-features = true }
+sp-blockchain = { workspace = true, default-features = true }
 sp-consensus = { workspace = true, default-features = true }
 sp-inherents = { workspace = true, default-features = true }
 sp-runtime = { workspace = true, default-features = true }

cumulus/client/consensus/proposer/src/lib.rs

@@ -15,19 +15,22 @@
 // You should have received a copy of the GNU General Public License
 // along with Cumulus. If not, see <https://www.gnu.org/licenses/>.
 
-//! The Cumulus [`Proposer`] is a wrapper around a Substrate [`sp_consensus::Environment`]
+//! The Cumulus [`ProposerInterface`] is an extension of the Substrate [`ProposerFactory`]
 //! for creating new parachain blocks.
 //!
 //! This utility is designed to be composed within any collator consensus algorithm.
 
 use async_trait::async_trait;
-
-use sp_consensus::{EnableProofRecording, Environment, Proposal, Proposer as SubstrateProposer};
+use cumulus_primitives_parachain_inherent::ParachainInherentData;
+use sc_basic_authorship::{ProposeArgs, ProposerFactory};
+use sc_block_builder::BlockBuilderApi;
+use sc_transaction_pool_api::TransactionPool;
+use sp_api::{ApiExt, CallApiAt, ProvideRuntimeApi};
+use sp_blockchain::HeaderBackend;
+use sp_consensus::{EnableProofRecording, Environment, Proposal};
 use sp_inherents::{InherentData, InherentDataProvider};
 use sp_runtime::{traits::Block as BlockT, Digest};
 use sp_state_machine::StorageProof;
-
-use cumulus_primitives_parachain_inherent::ParachainInherentData;
 use std::{fmt::Debug, time::Duration};
 
 /// Errors that can occur when proposing a parachain block.
@@ -50,7 +53,7 @@ impl Error {
 }
 
 /// A type alias for easily referring to the type of a proposal produced by a specific
-/// [`Proposer`].
+/// [`ProposerInterface`].
 pub type ProposalOf<B> = Proposal<B, StorageProof>;
 
 /// An interface for proposers.
@@ -80,39 +83,24 @@ pub trait ProposerInterface<Block: BlockT> {
 	) -> Result<Option<Proposal<Block, StorageProof>>, Error>;
 }
 
-/// A simple wrapper around a Substrate proposer for creating collations.
-pub struct Proposer<B, T> {
-	inner: T,
-	_marker: std::marker::PhantomData<B>,
-}
-
-impl<B, T> Proposer<B, T> {
-	/// Create a new Cumulus [`Proposer`].
-	pub fn new(inner: T) -> Self {
-		Proposer { inner, _marker: std::marker::PhantomData }
-	}
-}
-
 #[async_trait]
-impl<B, T> ProposerInterface<B> for Proposer<B, T>
+impl<Block, A, C> ProposerInterface<Block> for ProposerFactory<A, C, EnableProofRecording>
 where
-	B: sp_runtime::traits::Block,
-	T: Environment<B> + Send,
-	T::Error: Send + Sync + 'static,
-	T::Proposer: SubstrateProposer<B, ProofRecording = EnableProofRecording, Proof = StorageProof>,
-	<T::Proposer as SubstrateProposer<B>>::Error: Send + Sync + 'static,
+	A: TransactionPool<Block = Block> + 'static,
+	C: HeaderBackend<Block> + ProvideRuntimeApi<Block> + CallApiAt<Block> + Send + Sync + 'static,
+	C::Api: ApiExt<Block> + BlockBuilderApi<Block>,
+	Block: sp_runtime::traits::Block,
 {
 	async fn propose(
 		&mut self,
-		parent_header: &B::Header,
+		parent_header: &Block::Header,
 		paras_inherent_data: &ParachainInherentData,
 		other_inherent_data: InherentData,
 		inherent_digests: Digest,
 		max_duration: Duration,
 		block_size_limit: Option<usize>,
-	) -> Result<Option<Proposal<B, StorageProof>>, Error> {
+	) -> Result<Option<Proposal<Block, StorageProof>>, Error> {
 		let proposer = self
-			.inner
 			.init(parent_header)
 			.await
 			.map_err(|e| Error::proposer_creation(anyhow::Error::new(e)))?;
@@ -124,7 +112,13 @@ where
 			.map_err(|e| Error::proposing(anyhow::Error::new(e)))?;
 
 		proposer
-			.propose(inherent_data, inherent_digests, max_duration, block_size_limit)
+			.propose_block(ProposeArgs {
+				inherent_data,
+				inherent_digests,
+				max_duration,
+				block_size_limit,
+				ignored_nodes_by_proof_recording: None,
+			})
 			.await
 			.map(Some)
 			.map_err(|e| Error::proposing(anyhow::Error::new(e)).into())

cumulus/pallets/parachain-system/Cargo.toml

@@ -59,12 +59,13 @@ trie-standardmap = { workspace = true }
 
 
 # Substrate
+sc-consensus = { workspace = true }
+sp-api = { workspace = true, default-features = true }
 sp-consensus-slots = { workspace = true, default-features = true }
 sp-crypto-hashing = { workspace = true, default-features = true }
 sp-keyring = { workspace = true, default-features = true }
 sp-tracing = { workspace = true, default-features = true }
 sp-version = { workspace = true, default-features = true }
-
 # Cumulus
 cumulus-test-client = { workspace = true }
 cumulus-test-relay-sproof-builder = { workspace = true, default-features = true }

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

@@ -17,33 +17,32 @@
 //! The actual implementation of the validate block functionality.
 
 use super::{trie_cache, trie_recorder, MemoryOptimizedValidationParams};
-use crate::parachain_inherent::BasicParachainInherentData;
-use cumulus_primitives_core::{
-	relay_chain::Hash as RHash, ParachainBlockData, PersistedValidationData,
-};
-
-use polkadot_parachain_primitives::primitives::{
-	HeadData, RelayChainBlockNumber, ValidationResult,
-};
-
+use crate::{parachain_inherent::BasicParachainInherentData, ClaimQueueOffset, CoreSelector};
 use alloc::vec::Vec;
 use codec::{Decode, Encode};
-
-use cumulus_primitives_core::relay_chain::vstaging::{UMPSignal, UMP_SEPARATOR};
+use cumulus_primitives_core::{
+	relay_chain::{
+		vstaging::{UMPSignal, UMP_SEPARATOR},
+		Hash as RHash,
+	},
+	ParachainBlockData, PersistedValidationData,
+};
 use frame_support::{
 	traits::{ExecuteBlock, Get, IsSubType},
 	BoundedVec,
 };
+use polkadot_parachain_primitives::primitives::{
+	HeadData, RelayChainBlockNumber, ValidationResult,
+};
 use sp_core::storage::{ChildInfo, StateVersion};
 use sp_externalities::{set_and_run_with_externalities, Externalities};
 use sp_io::{hashing::blake2_128, KillStorageResult};
 use sp_runtime::traits::{
-	Block as BlockT, ExtrinsicCall, ExtrinsicLike, HashingFor, Header as HeaderT,
+	Block as BlockT, ExtrinsicCall, ExtrinsicLike, Hash as HashT, HashingFor, Header as HeaderT,
 };
-
 use sp_state_machine::OverlayedChanges;
-use sp_trie::ProofSizeProvider;
-use trie_recorder::SizeOnlyRecorderProvider;
+use sp_trie::{HashDBT, ProofSizeProvider, EMPTY_PREFIX};
+use trie_recorder::{SeenNodes, SizeOnlyRecorderProvider};
 
 type Ext<'a, Block, Backend> = sp_state_machine::Ext<'a, HashingFor<Block>, Backend>;
 
@@ -174,37 +173,40 @@ where
 	let num_blocks = blocks.len();
 
 	// Create the db
-	let db = match proof.to_memory_db(Some(parent_header.state_root())) {
+	let mut db = match proof.to_memory_db(Some(parent_header.state_root())) {
 		Ok((db, _)) => db,
 		Err(_) => panic!("Compact proof decoding failure."),
 	};
 
 	core::mem::drop(proof);
 
 	let cache_provider = trie_cache::CacheProvider::new();
-	// 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_with_cache(
-		db,
-		*parent_header.state_root(),
-		cache_provider,
-	)
-	.build();
-
-	// We use the same recorder when executing all blocks. So, each node only contributes once to
-	// the total size of the storage proof. This recorder should only be used for `execute_block`.
-	let mut execute_recorder = SizeOnlyRecorderProvider::default();
-	// `backend` with the `execute_recorder`. As the `execute_recorder`, this should only be used
-	// for `execute_block`.
-	let execute_backend = sp_state_machine::TrieBackendBuilder::wrap(&backend)
-		.with_recorder(execute_recorder.clone())
+	let seen_nodes = SeenNodes::<HashingFor<B>>::default();
+
+	for (block_index, block) in blocks.into_iter().enumerate() {
+		// 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_with_cache(
+			&db,
+			*parent_header.state_root(),
+			&cache_provider,
+		)
 		.build();
 
-	// We let all blocks contribute to the same overlay. Data written by a previous block will be
-	// directly accessible without going to the db.
-	let mut overlay = OverlayedChanges::default();
+		// We use the same recorder when executing all blocks. So, each node only contributes once
+		// to the total size of the storage proof. This recorder should only be used for
+		// `execute_block`.
+		let mut execute_recorder = SizeOnlyRecorderProvider::with_seen_nodes(seen_nodes.clone());
+		// `backend` with the `execute_recorder`. As the `execute_recorder`, this should only be
+		// used for `execute_block`.
+		let execute_backend = sp_state_machine::TrieBackendBuilder::wrap(&backend)
+			.with_recorder(execute_recorder.clone())
+			.build();
+
+		// We let all blocks contribute to the same overlay. Data written by a previous block will
+		// be directly accessible without going to the db.
+		let mut overlay = OverlayedChanges::default();
 
-	for (block_index, block) in blocks.into_iter().enumerate() {
 		parent_header = block.header().clone();
 		let inherent_data = extract_parachain_inherent_data(&block);
 
@@ -308,11 +310,33 @@ where
 					);
 				}
 			},
-		)
+		);
+
+		if block_index + 1 != num_blocks {
+			let mut changes = overlay
+				.drain_storage_changes(
+					&backend,
+					<PSC as frame_system::Config>::Version::get().state_version(),
+				)
+				.expect("Failed to get drain storage changes from the overlay.");
+
+			drop(backend);
+
+			// We just forward the changes directly to our db.
+			changes.transaction.drain().into_iter().for_each(|(_, (value, count))| {
+				// We only care about inserts and not deletes.
+				if count > 0 {
+					db.insert(EMPTY_PREFIX, &value);
+
+					let hash = HashingFor::<B>::hash(&value);
+					seen_nodes.borrow_mut().insert(hash);
+				}
+			});
+		}
 	}
 
 	if !upward_message_signals.is_empty() {
-		let mut selected_core = None;
+		let mut selected_core: Option<(CoreSelector, ClaimQueueOffset)> = None;
 		let mut approved_peer = None;
 
 		upward_message_signals.iter().for_each(|s| {