Remove requirement on Hash = H256, make Proposer return StorageChanges and Proof (#3860)

* Extend `Proposer` to optionally generate a proof of the proposal * Something * Refactor sr-api to not depend on client anymore * Fix benches * Apply suggestions from code review Co-Authored-By: Tomasz Drwięga <tomusdrw@users.noreply.github.com> * Apply suggestions from code review * Introduce new `into_storage_changes` function * Switch to runtime api for `execute_block` and don't require `H256` anywhere in the code * Put the `StorageChanges` into the `Proposal` * Move the runtime api error to its own trait * Adds `StorageTransactionCache` to the runtime api This requires that we add `type NodeBlock = ` to the `impl_runtime_apis!` macro to work around some bugs in rustc :( * Remove `type NodeBlock` and switch to a "better" hack * Start using the transaction cache from the runtime api * Make it compile * Move `InMemory` to its own file * Make all tests work again * Return block, storage_changes and proof from Blockbuilder::bake() * Make sure that we use/set `storage_changes` when possible * Add test * Fix deadlock * Remove accidentally added folders * Introduce `RecordProof` as argument type to be more explicit * Update client/src/client.rs Co-Authored-By: Tomasz Drwięga <tomusdrw@users.noreply.github.com> * Update primitives/state-machine/src/ext.rs Co-Authored-By: Tomasz Drwięga <tomusdrw@users.noreply.github.com> * Integrates review feedback * Remove `unsafe` usage * Update client/block-builder/src/lib.rs Co-Authored-By: Benjamin Kampmann <ben@gnunicorn.org> * Update client/src/call_executor.rs * Bump versions Co-authored-by: Tomasz Drwięga <tomusdrw@users.noreply.github.com> Co-authored-by: Benjamin Kampmann <ben.kampmann@googlemail.com>
2026-06-14 05:11:09 +00:00 · 2020-01-10 10:48:32 +01:00
parent 74d6e660c6
commit fd6b29dd2c
140 changed files with 4860 additions and 3339 deletions
@@ -22,35 +22,13 @@ use sp_consensus::{
 	BlockImportParams, BlockImport, BlockOrigin, Error as ConsensusError,
 	ForkChoiceStrategy,
 };
-use hash_db::Hasher;
 use sp_runtime::Justification;
 use sp_runtime::traits::{Block as BlockT};
 use sp_runtime::generic::BlockId;
-use sp_core::Blake2Hasher;
 use codec::alloc::collections::hash_map::HashMap;

 /// Extension trait for a test client.
 pub trait ClientExt<Block: BlockT>: Sized {
-	/// Import block to the chain. No finality.
-	fn import(&self, origin: BlockOrigin, block: Block)
-		-> Result<(), ConsensusError>;
-
-	/// Import a block and make it our best block if possible.
-	fn import_as_best(&self, origin: BlockOrigin, block: Block)
-		-> Result<(), ConsensusError>;
-
-	/// Import a block and finalize it.
-	fn import_as_final(&self, origin: BlockOrigin, block: Block)
-		-> Result<(), ConsensusError>;
-
-	/// Import block with justification, finalizes block.
-	fn import_justified(
-		&self,
-		origin: BlockOrigin,
-		block: Block,
-		justification: Justification
-	) -> Result<(), ConsensusError>;
-
 	/// Finalize a block.
 	fn finalize_block(
 		&self,
@@ -62,96 +40,34 @@ pub trait ClientExt<Block: BlockT>: Sized {
 	fn genesis_hash(&self) -> <Block as BlockT>::Hash;
 }

-impl<B, E, RA, Block> ClientExt<Block> for Client<B, E, Block, RA>
-	where
-		B: sc_client_api::backend::Backend<Block, Blake2Hasher>,
-		E: sc_client::CallExecutor<Block, Blake2Hasher>,
-		for<'r> &'r Self: BlockImport<Block, Error=ConsensusError>,
-		Block: BlockT<Hash=<Blake2Hasher as Hasher>::Out>,
-{
-	fn import(&self, origin: BlockOrigin, block: Block)
-		-> Result<(), ConsensusError>
-	{
-		let (header, extrinsics) = block.deconstruct();
-		let import = BlockImportParams {
-			origin,
-			header,
-			justification: None,
-			post_digests: vec![],
-			body: Some(extrinsics),
-			finalized: false,
-			auxiliary: Vec::new(),
-			fork_choice: ForkChoiceStrategy::LongestChain,
-			allow_missing_state: false,
-			import_existing: false,
-		};
+/// Extension trait for a test client around block importing.
+pub trait ClientBlockImportExt<Block: BlockT>: Sized {
+	/// Import block to the chain. No finality.
+	fn import(&mut self, origin: BlockOrigin, block: Block) -> Result<(), ConsensusError>;

-		BlockImport::import_block(&mut (&*self), import, HashMap::new()).map(|_| ())
-	}
+	/// Import a block and make it our best block if possible.
+	fn import_as_best(&mut self, origin: BlockOrigin, block: Block) -> Result<(), ConsensusError>;

-	fn import_as_best(&self, origin: BlockOrigin, block: Block)
-		-> Result<(), ConsensusError>
-	{
-		let (header, extrinsics) = block.deconstruct();
-		let import = BlockImportParams {
-			origin,
-			header,
-			justification: None,
-			post_digests: vec![],
-			body: Some(extrinsics),
-			finalized: false,
-			auxiliary: Vec::new(),
-			fork_choice: ForkChoiceStrategy::Custom(true),
-			allow_missing_state: false,
-			import_existing: false,
-		};
-
-		BlockImport::import_block(&mut (&*self), import, HashMap::new()).map(|_| ())
-	}
-
-	fn import_as_final(&self, origin: BlockOrigin, block: Block)
-		-> Result<(), ConsensusError>
-	{
-		let (header, extrinsics) = block.deconstruct();
-		let import = BlockImportParams {
-			origin,
-			header,
-			justification: None,
-			post_digests: vec![],
-			body: Some(extrinsics),
-			finalized: true,
-			auxiliary: Vec::new(),
-			fork_choice: ForkChoiceStrategy::Custom(true),
-			allow_missing_state: false,
-			import_existing: false,
-		};
-
-		BlockImport::import_block(&mut (&*self), import, HashMap::new()).map(|_| ())
-	}
+	/// Import a block and finalize it.
+	fn import_as_final(&mut self, origin: BlockOrigin, block: Block)
+		-> Result<(), ConsensusError>;

+	/// Import block with justification, finalizes block.
 	fn import_justified(
-		&self,
+		&mut self,
 		origin: BlockOrigin,
 		block: Block,
-		justification: Justification,
-	) -> Result<(), ConsensusError> {
-		let (header, extrinsics) = block.deconstruct();
-		let import = BlockImportParams {
-			origin,
-			header,
-			justification: Some(justification),
-			post_digests: vec![],
-			body: Some(extrinsics),
-			finalized: true,
-			auxiliary: Vec::new(),
-			fork_choice: ForkChoiceStrategy::LongestChain,
-			allow_missing_state: false,
-			import_existing: false,
-		};
-
-		BlockImport::import_block(&mut (&*self), import, HashMap::new()).map(|_| ())
-	}
+		justification: Justification
+	) -> Result<(), ConsensusError>;
+}

+impl<B, E, RA, Block> ClientExt<Block> for Client<B, E, Block, RA>
+	where
+		B: sc_client_api::backend::Backend<Block>,
+		E: sc_client::CallExecutor<Block>,
+		Self: BlockImport<Block, Error = ConsensusError>,
+		Block: BlockT,
+{
 	fn finalize_block(
 		&self,
 		id: BlockId<Block>,
@@ -164,3 +80,175 @@ impl<B, E, RA, Block> ClientExt<Block> for Client<B, E, Block, RA>
 		self.block_hash(0.into()).unwrap().unwrap()
 	}
 }
+
+/// This implementation is required, because of the weird api requirements around `BlockImport`.
+impl<Block: BlockT, T, Transaction> ClientBlockImportExt<Block> for std::sync::Arc<T>
+	where for<'r> &'r T: BlockImport<Block, Error = ConsensusError, Transaction = Transaction>
+{
+	fn import(&mut self, origin: BlockOrigin, block: Block) -> Result<(), ConsensusError> {
+		let (header, extrinsics) = block.deconstruct();
+		let import = BlockImportParams {
+			origin,
+			header,
+			justification: None,
+			post_digests: vec![],
+			body: Some(extrinsics),
+			storage_changes: None,
+			finalized: false,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::LongestChain,
+			allow_missing_state: false,
+			import_existing: false,
+		};
+
+		BlockImport::import_block(self, import, HashMap::new()).map(|_| ())
+	}
+
+	fn import_as_best(&mut self, origin: BlockOrigin, block: Block) -> Result<(), ConsensusError> {
+		let (header, extrinsics) = block.deconstruct();
+		let import = BlockImportParams {
+			origin,
+			header,
+			justification: None,
+			post_digests: vec![],
+			body: Some(extrinsics),
+			storage_changes: None,
+			finalized: false,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::Custom(true),
+			allow_missing_state: false,
+			import_existing: false,
+		};
+
+		BlockImport::import_block(self, import, HashMap::new()).map(|_| ())
+	}
+
+	fn import_as_final(&mut self, origin: BlockOrigin, block: Block) -> Result<(), ConsensusError> {
+		let (header, extrinsics) = block.deconstruct();
+		let import = BlockImportParams {
+			origin,
+			header,
+			justification: None,
+			post_digests: vec![],
+			body: Some(extrinsics),
+			storage_changes: None,
+			finalized: true,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::Custom(true),
+			allow_missing_state: false,
+			import_existing: false,
+		};
+
+		BlockImport::import_block(self, import, HashMap::new()).map(|_| ())
+	}
+
+	fn import_justified(
+		&mut self,
+		origin: BlockOrigin,
+		block: Block,
+		justification: Justification,
+	) -> Result<(), ConsensusError> {
+		let (header, extrinsics) = block.deconstruct();
+		let import = BlockImportParams {
+			origin,
+			header,
+			justification: Some(justification),
+			post_digests: vec![],
+			body: Some(extrinsics),
+			storage_changes: None,
+			finalized: true,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::LongestChain,
+			allow_missing_state: false,
+			import_existing: false,
+		};
+
+		BlockImport::import_block(self, import, HashMap::new()).map(|_| ())
+	}
+}
+
+impl<B, E, RA, Block: BlockT> ClientBlockImportExt<Block> for Client<B, E, Block, RA>
+	where
+		Self: BlockImport<Block, Error = ConsensusError>,
+{
+	fn import(&mut self, origin: BlockOrigin, block: Block) -> Result<(), ConsensusError> {
+		let (header, extrinsics) = block.deconstruct();
+		let import = BlockImportParams {
+			origin,
+			header,
+			justification: None,
+			post_digests: vec![],
+			body: Some(extrinsics),
+			storage_changes: None,
+			finalized: false,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::LongestChain,
+			allow_missing_state: false,
+			import_existing: false,
+		};
+
+		BlockImport::import_block(self, import, HashMap::new()).map(|_| ())
+	}
+
+	fn import_as_best(&mut self, origin: BlockOrigin, block: Block) -> Result<(), ConsensusError> {
+		let (header, extrinsics) = block.deconstruct();
+		let import = BlockImportParams {
+			origin,
+			header,
+			justification: None,
+			post_digests: vec![],
+			body: Some(extrinsics),
+			storage_changes: None,
+			finalized: false,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::Custom(true),
+			allow_missing_state: false,
+			import_existing: false,
+		};
+
+		BlockImport::import_block(self, import, HashMap::new()).map(|_| ())
+	}
+
+	fn import_as_final(&mut self, origin: BlockOrigin, block: Block) -> Result<(), ConsensusError> {
+		let (header, extrinsics) = block.deconstruct();
+		let import = BlockImportParams {
+			origin,
+			header,
+			justification: None,
+			post_digests: vec![],
+			body: Some(extrinsics),
+			storage_changes: None,
+			finalized: true,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::Custom(true),
+			allow_missing_state: false,
+			import_existing: false,
+		};
+
+		BlockImport::import_block(self, import, HashMap::new()).map(|_| ())
+	}
+
+	fn import_justified(
+		&mut self,
+		origin: BlockOrigin,
+		block: Block,
+		justification: Justification,
+	) -> Result<(), ConsensusError> {
+		let (header, extrinsics) = block.deconstruct();
+		let import = BlockImportParams {
+			origin,
+			header,
+			justification: Some(justification),
+			post_digests: vec![],
+			body: Some(extrinsics),
+			storage_changes: None,
+			finalized: true,
+			auxiliary: Vec::new(),
+			fork_choice: ForkChoiceStrategy::LongestChain,
+			allow_missing_state: false,
+			import_existing: false,
+		};
+
+		BlockImport::import_block(self, import, HashMap::new()).map(|_| ())
+	}
+}
@@ -34,11 +34,10 @@ pub use sp_core::{Blake2Hasher, traits::BareCryptoStorePtr};
 pub use sp_runtime::{Storage, StorageChild};
 pub use sp_state_machine::ExecutionStrategy;

-pub use self::client_ext::ClientExt;
+pub use self::client_ext::{ClientExt, ClientBlockImportExt};

 use std::sync::Arc;
 use std::collections::HashMap;
-use hash_db::Hasher;
 use sp_core::storage::{well_known_keys, ChildInfo};
 use sp_runtime::traits::Block as BlockT;
 use sc_client::LocalCallExecutor;
@@ -71,36 +70,20 @@ pub struct TestClientBuilder<Executor, Backend, G: GenesisInit> {
 	keystore: Option<BareCryptoStorePtr>,
 }

-impl<Block, Executor, G: GenesisInit> Default for TestClientBuilder<
-	Executor,
-	Backend<Block>,
-	G,
-> where
-	Block: BlockT<Hash=<Blake2Hasher as Hasher>::Out>,
-{
+impl<Block: BlockT, Executor, G: GenesisInit> Default
+	for TestClientBuilder<Executor, Backend<Block>, G> {
 	fn default() -> Self {
 		Self::with_default_backend()
 	}
 }

-impl<Block, Executor, G: GenesisInit> TestClientBuilder<
-	Executor,
-	Backend<Block>,
-	G,
-> where
-	Block: BlockT<Hash=<Blake2Hasher as Hasher>::Out>,
-{
+impl<Block: BlockT, Executor, G: GenesisInit> TestClientBuilder<Executor, Backend<Block>, G> {
 	/// Create new `TestClientBuilder` with default backend.
 	pub fn with_default_backend() -> Self {
 		let backend = Arc::new(Backend::new_test(std::u32::MAX, std::u64::MAX));
 		Self::with_backend(backend)
 	}

-	/// Give access to the underlying backend of these clients
-	pub fn backend(&self) -> Arc<Backend<Block>> {
-		self.backend.clone()
-	}
-
 	/// Create new `TestClientBuilder` with default backend and pruning window size
 	pub fn with_pruning_window(keep_blocks: u32) -> Self {
 		let backend = Arc::new(Backend::new_test(keep_blocks, 0));
@@ -132,6 +115,11 @@ impl<Executor, Backend, G: GenesisInit> TestClientBuilder<Executor, Backend, G>
 		&mut self.genesis_init
 	}

+	/// Give access to the underlying backend of these clients
+	pub fn backend(&self) -> Arc<Backend> {
+		self.backend.clone()
+	}
+
 	/// Extend child storage
 	pub fn add_child_storage(
 		mut self,
@@ -177,9 +165,9 @@ impl<Executor, Backend, G: GenesisInit> TestClientBuilder<Executor, Backend, G>
 		>,
 		sc_client::LongestChain<Backend, Block>,
 	) where
-		Executor: sc_client::CallExecutor<Block, Blake2Hasher>,
-		Backend: sc_client_api::backend::Backend<Block, Blake2Hasher>,
-		Block: BlockT<Hash=<Blake2Hasher as Hasher>::Out>,
+		Executor: sc_client::CallExecutor<Block>,
+		Backend: sc_client_api::backend::Backend<Block>,
+		Block: BlockT,
 	{

 		let storage = {
@@ -237,8 +225,8 @@ impl<E, Backend, G: GenesisInit> TestClientBuilder<
 	) where
 		I: Into<Option<NativeExecutor<E>>>,
 		E: sc_executor::NativeExecutionDispatch,
-		Backend: sc_client_api::backend::Backend<Block, Blake2Hasher>,
-		Block: BlockT<Hash=<Blake2Hasher as Hasher>::Out>,
+		Backend: sc_client_api::backend::Backend<Block>,
+		Block: BlockT,
 	{
 		let executor = executor.into().unwrap_or_else(||
 			NativeExecutor::new(WasmExecutionMethod::Interpreted, None)
@@ -10,6 +10,7 @@ substrate-test-client = { version = "2.0.0", path = "../../client" }
 sp-core = { version = "2.0.0", path = "../../../primitives/core" }
 substrate-test-runtime = { version = "2.0.0", path = "../../runtime" }
 sp-runtime = { version = "2.0.0", path = "../../../primitives/runtime" }
+sp-api = { version = "2.0.0", path = "../../../primitives/api" }
 sp-blockchain = { version = "2.0.0", path = "../../../primitives/blockchain" }
 codec = { package = "parity-scale-codec", version = "1.0.0" }
 sc-client-api = { version = "2.0.0", path = "../../../client/api" }
@@ -16,8 +16,9 @@

 //! Block Builder extensions for tests.

-use substrate_test_runtime;
-use sp_runtime::traits::ProvideRuntimeApi;
+use sp_api::{ApiExt, ProvideRuntimeApi};
+use sc_client_api::backend;
+use sp_runtime::traits::HasherFor;

 use sc_block_builder::BlockBuilderApi;

@@ -33,9 +34,17 @@ pub trait BlockBuilderExt {
 	) -> Result<(), sp_blockchain::Error>;
 }

-impl<'a, A> BlockBuilderExt for sc_block_builder::BlockBuilder<'a, substrate_test_runtime::Block, A> where
-	A: ProvideRuntimeApi + 'a,
-	A::Api: BlockBuilderApi<substrate_test_runtime::Block, Error = sp_blockchain::Error>,
+impl<'a, A, B> BlockBuilderExt for sc_block_builder::BlockBuilder<'a, substrate_test_runtime::Block, A, B> where
+	A: ProvideRuntimeApi<substrate_test_runtime::Block> + 'a,
+	A::Api: BlockBuilderApi<substrate_test_runtime::Block, Error = sp_blockchain::Error> +
+		ApiExt<
+			substrate_test_runtime::Block,
+			StateBackend = backend::StateBackendFor<B, substrate_test_runtime::Block>
+		>,
+	B: backend::Backend<substrate_test_runtime::Block>,
+	// Rust bug: https://github.com/rust-lang/rust/issues/24159
+	backend::StateBackendFor<B, substrate_test_runtime::Block>:
+		sp_api::StateBackend<HasherFor<substrate_test_runtime::Block>>,
 {
 	fn push_transfer(&mut self, transfer: substrate_test_runtime::Transfer) -> Result<(), sp_blockchain::Error> {
 		self.push(transfer.into_signed_tx())
@@ -32,7 +32,7 @@ pub use self::block_builder_ext::BlockBuilderExt;
 use sp_core::sr25519;
 use sp_core::storage::{ChildInfo, Storage, StorageChild};
 use substrate_test_runtime::genesismap::{GenesisConfig, additional_storage_with_genesis};
-use sp_runtime::traits::{Block as BlockT, Header as HeaderT, Hash as HashT, NumberFor};
+use sp_runtime::traits::{Block as BlockT, Header as HeaderT, Hash as HashT, NumberFor, HasherFor};
 use sc_client::{
 	light::fetcher::{
 		Fetcher,
@@ -45,7 +45,10 @@ use sc_client::{
 /// A prelude to import in tests.
 pub mod prelude {
 	// Trait extensions
-	pub use super::{BlockBuilderExt, DefaultTestClientBuilderExt, TestClientBuilderExt, ClientExt};
+	pub use super::{
+		BlockBuilderExt, DefaultTestClientBuilderExt, TestClientBuilderExt, ClientExt,
+		ClientBlockImportExt,
+	};
 	// Client structs
 	pub use super::{
 		TestClient, TestClientBuilder, Backend, LightBackend,
@@ -88,7 +91,7 @@ pub type LightExecutor = sc_client::light::call_executor::GenesisCallExecutor<
 	sc_client::LocalCallExecutor<
 		sc_client::light::backend::Backend<
 			sc_client_db::light::LightStorage<substrate_test_runtime::Block>,
-			Blake2Hasher,
+			HasherFor<substrate_test_runtime::Block>
 		>,
 		NativeExecutor<LocalExecutor>
 	>
@@ -165,10 +168,7 @@ pub trait DefaultTestClientBuilderExt: Sized {
 	fn new() -> Self;
 }

-impl DefaultTestClientBuilderExt for TestClientBuilder<
-	Executor,
-	Backend,
-> {
+impl DefaultTestClientBuilderExt for TestClientBuilder<Executor, Backend> {
 	fn new() -> Self {
 		Self::with_default_backend()
 	}
@@ -176,64 +176,26 @@ impl DefaultTestClientBuilderExt for TestClientBuilder<

 /// A `test-runtime` extensions to `TestClientBuilder`.
 pub trait TestClientBuilderExt<B>: Sized {
+	/// Returns a mutable reference to the genesis parameters.
+	fn genesis_init_mut(&mut self) -> &mut GenesisParameters;
+
 	/// Enable or disable support for changes trie in genesis.
-	fn set_support_changes_trie(self, support_changes_trie: bool) -> Self;
+	fn set_support_changes_trie(mut self, support_changes_trie: bool) -> Self {
+		self.genesis_init_mut().support_changes_trie = support_changes_trie;
+		self
+	}

 	/// Override the default value for Wasm heap pages.
-	fn set_heap_pages(self, heap_pages: u64) -> Self;
+	fn set_heap_pages(mut self, heap_pages: u64) -> Self {
+		self.genesis_init_mut().heap_pages_override = Some(heap_pages);
+		self
+	}

 	/// Add an extra value into the genesis storage.
 	///
 	/// # Panics
 	///
 	/// Panics if the key is empty.
-	fn add_extra_child_storage<SK: Into<Vec<u8>>, K: Into<Vec<u8>>, V: Into<Vec<u8>>>(
-		self,
-		storage_key: SK,
-		child_info: ChildInfo,
-		key: K,
-		value: V,
-	) -> Self;
-
-	/// Add an extra child value into the genesis storage.
-	///
-	/// # Panics
-	///
-	/// Panics if the key is empty.
-	fn add_extra_storage<K: Into<Vec<u8>>, V: Into<Vec<u8>>>(self, key: K, value: V) -> Self;
-
-	/// Build the test client.
-	fn build(self) -> Client<B> {
-		self.build_with_longest_chain().0
-	}
-
-	/// Build the test client and longest chain selector.
-	fn build_with_longest_chain(self) -> (Client<B>, sc_client::LongestChain<B, substrate_test_runtime::Block>);
-}
-
-impl<B> TestClientBuilderExt<B> for TestClientBuilder<
-	sc_client::LocalCallExecutor<B, sc_executor::NativeExecutor<LocalExecutor>>,
-	B
-> where
-	B: sc_client_api::backend::Backend<substrate_test_runtime::Block, Blake2Hasher>,
-{
-	fn set_heap_pages(mut self, heap_pages: u64) -> Self {
-		self.genesis_init_mut().heap_pages_override = Some(heap_pages);
-		self
-	}
-
-	fn set_support_changes_trie(mut self, support_changes_trie: bool) -> Self {
-		self.genesis_init_mut().support_changes_trie = support_changes_trie;
-		self
-	}
-
-	fn add_extra_storage<K: Into<Vec<u8>>, V: Into<Vec<u8>>>(mut self, key: K, value: V) -> Self {
-		let key = key.into();
-		assert!(!key.is_empty());
-		self.genesis_init_mut().extra_storage.top.insert(key, value.into());
-		self
-	}
-
 	fn add_extra_child_storage<SK: Into<Vec<u8>>, K: Into<Vec<u8>>, V: Into<Vec<u8>>>(
 		mut self,
 		storage_key: SK,
@@ -254,10 +216,51 @@ impl<B> TestClientBuilderExt<B> for TestClientBuilder<
 		self
 	}

+	/// Add an extra child value into the genesis storage.
+	///
+	/// # Panics
+	///
+	/// Panics if the key is empty.
+	fn add_extra_storage<K: Into<Vec<u8>>, V: Into<Vec<u8>>>(mut self, key: K, value: V) -> Self {
+		let key = key.into();
+		assert!(!key.is_empty());
+		self.genesis_init_mut().extra_storage.top.insert(key, value.into());
+		self
+	}
+
+	/// Build the test client.
+	fn build(self) -> Client<B> {
+		self.build_with_longest_chain().0
+	}
+
+	/// Build the test client and longest chain selector.
+	fn build_with_longest_chain(self) -> (Client<B>, sc_client::LongestChain<B, substrate_test_runtime::Block>);
+
+	/// Build the test client and the backend.
+	fn build_with_backend(self) -> (Client<B>, Arc<B>);
+}
+
+impl<B> TestClientBuilderExt<B> for TestClientBuilder<
+	sc_client::LocalCallExecutor<B, sc_executor::NativeExecutor<LocalExecutor>>,
+	B
+> where
+	B: sc_client_api::backend::Backend<substrate_test_runtime::Block>,
+	// Rust bug: https://github.com/rust-lang/rust/issues/24159
+	<B as sc_client_api::backend::Backend<substrate_test_runtime::Block>>::State:
+		sp_api::StateBackend<HasherFor<substrate_test_runtime::Block>>,
+{
+	fn genesis_init_mut(&mut self) -> &mut GenesisParameters {
+		Self::genesis_init_mut(self)
+	}

 	fn build_with_longest_chain(self) -> (Client<B>, sc_client::LongestChain<B, substrate_test_runtime::Block>) {
 		self.build_with_native_executor(None)
 	}
+
+	fn build_with_backend(self) -> (Client<B>, Arc<B>) {
+		let backend = self.backend();
+		(self.build_with_native_executor(None).0, backend)
+	}
 }

 /// Type of optional fetch callback.
@@ -21,19 +21,22 @@

 use std::sync::Arc;

-use sc_client_api::backend::LocalBackend;
-use crate::block_builder_ext::BlockBuilderExt;
+use crate::{
+	AccountKeyring, ClientBlockImportExt, BlockBuilderExt, TestClientBuilder, TestClientBuilderExt,
+};
+use sc_client_api::backend;
 use sc_client_api::blockchain::{Backend as BlockChainBackendT, HeaderBackend};
-use crate::{AccountKeyring, ClientExt, TestClientBuilder, TestClientBuilderExt};
 use substrate_test_client::sp_consensus::BlockOrigin;
-use sp_core::Blake2Hasher;
 use substrate_test_runtime::{self, Transfer};
 use sp_runtime::generic::BlockId;
-use sp_runtime::traits::Block as BlockT;
+use sp_runtime::traits::{Block as BlockT, HasherFor};

 /// helper to test the `leaves` implementation for various backends
 pub fn test_leaves_for_backend<B: 'static>(backend: Arc<B>) where
-	B: LocalBackend<substrate_test_runtime::Block, Blake2Hasher>,
+	B: backend::Backend<substrate_test_runtime::Block>,
+	// Rust bug: https://github.com/rust-lang/rust/issues/24159
+	backend::StateBackendFor<B, substrate_test_runtime::Block>:
+		sp_api::StateBackend<HasherFor<substrate_test_runtime::Block>>,
 {
 	// block tree:
 	// G -> A1 -> A2 -> A3 -> A4 -> A5
@@ -41,7 +44,7 @@ pub fn test_leaves_for_backend<B: 'static>(backend: Arc<B>) where
 	//			  B2 -> C3
 	//		A1 -> D2

-	let client = TestClientBuilder::with_backend(backend.clone()).build();
+	let mut client = TestClientBuilder::with_backend(backend.clone()).build();
 	let blockchain = backend.blockchain();

 	let genesis_hash = client.chain_info().genesis_hash;
@@ -51,44 +54,72 @@ pub fn test_leaves_for_backend<B: 'static>(backend: Arc<B>) where
 		vec![genesis_hash]);

 	// G -> A1
-	let a1 = client.new_block(Default::default()).unwrap().bake().unwrap();
+	let a1 = client.new_block(Default::default()).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a1.clone()).unwrap();
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a1.hash()]);
+		vec![a1.hash()],
+	);

 	// A1 -> A2
-	let a2 = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap().bake().unwrap();
+	let a2 = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a2.clone()).unwrap();

 	#[allow(deprecated)]
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a2.hash()]);
+		vec![a2.hash()],
+	);

 	// A2 -> A3
-	let a3 = client.new_block_at(&BlockId::Hash(a2.hash()), Default::default()).unwrap().bake().unwrap();
+	let a3 = client.new_block_at(
+		&BlockId::Hash(a2.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a3.clone()).unwrap();
+
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a3.hash()]);
+		vec![a3.hash()],
+	);

 	// A3 -> A4
-	let a4 = client.new_block_at(&BlockId::Hash(a3.hash()), Default::default()).unwrap().bake().unwrap();
+	let a4 = client.new_block_at(
+		&BlockId::Hash(a3.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a4.clone()).unwrap();
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a4.hash()]);
+		vec![a4.hash()],
+	);

 	// A4 -> A5
-	let a5 = client.new_block_at(&BlockId::Hash(a4.hash()), Default::default()).unwrap().bake().unwrap();
+	let a5 = client.new_block_at(
+		&BlockId::Hash(a4.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
+
 	client.import(BlockOrigin::Own, a5.clone()).unwrap();
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a5.hash()]);
+		vec![a5.hash()],
+	);

 	// A1 -> B2
-	let mut builder = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap();
+
 	// this push is required as otherwise B2 has the same hash as A2 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -96,28 +127,44 @@ pub fn test_leaves_for_backend<B: 'static>(backend: Arc<B>) where
 		amount: 41,
 		nonce: 0,
 	}).unwrap();
-	let b2 = builder.bake().unwrap();
+	let b2 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, b2.clone()).unwrap();
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a5.hash(), b2.hash()]);
+		vec![a5.hash(), b2.hash()],
+	);

 	// B2 -> B3
-	let b3 = client.new_block_at(&BlockId::Hash(b2.hash()), Default::default()).unwrap().bake().unwrap();
+	let b3 = client.new_block_at(
+		&BlockId::Hash(b2.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
+
 	client.import(BlockOrigin::Own, b3.clone()).unwrap();
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a5.hash(), b3.hash()]);
+		vec![a5.hash(), b3.hash()],
+	);

 	// B3 -> B4
-	let b4 = client.new_block_at(&BlockId::Hash(b3.hash()), Default::default()).unwrap().bake().unwrap();
+	let b4 = client.new_block_at(
+		&BlockId::Hash(b3.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, b4.clone()).unwrap();
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a5.hash(), b4.hash()]);
+		vec![a5.hash(), b4.hash()],
+	);

 	// // B2 -> C3
-	let mut builder = client.new_block_at(&BlockId::Hash(b2.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(b2.hash()),
+		Default::default(),
+		false,
+	).unwrap();
 	// this push is required as otherwise C3 has the same hash as B3 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -125,14 +172,19 @@ pub fn test_leaves_for_backend<B: 'static>(backend: Arc<B>) where
 		amount: 1,
 		nonce: 1,
 	}).unwrap();
-	let c3 = builder.bake().unwrap();
+	let c3 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, c3.clone()).unwrap();
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a5.hash(), b4.hash(), c3.hash()]);
+		vec![a5.hash(), b4.hash(), c3.hash()],
+	);

 	// A1 -> D2
-	let mut builder = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap();
 	// this push is required as otherwise D2 has the same hash as B2 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -140,16 +192,20 @@ pub fn test_leaves_for_backend<B: 'static>(backend: Arc<B>) where
 		amount: 1,
 		nonce: 0,
 	}).unwrap();
-	let d2 = builder.bake().unwrap();
+	let d2 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, d2.clone()).unwrap();
 	assert_eq!(
 		blockchain.leaves().unwrap(),
-		vec![a5.hash(), b4.hash(), c3.hash(), d2.hash()]);
+		vec![a5.hash(), b4.hash(), c3.hash(), d2.hash()],
+	);
 }

 /// helper to test the `children` implementation for various backends
 pub fn test_children_for_backend<B: 'static>(backend: Arc<B>) where
-	B: LocalBackend<substrate_test_runtime::Block, Blake2Hasher>,
+	B: backend::LocalBackend<substrate_test_runtime::Block>,
+	// Rust bug: https://github.com/rust-lang/rust/issues/24159
+	<B as backend::Backend<substrate_test_runtime::Block>>::State:
+		sp_api::StateBackend<HasherFor<substrate_test_runtime::Block>>,
 {
 	// block tree:
 	// G -> A1 -> A2 -> A3 -> A4 -> A5
@@ -157,31 +213,51 @@ pub fn test_children_for_backend<B: 'static>(backend: Arc<B>) where
 	//			  B2 -> C3
 	//		A1 -> D2

-	let client = TestClientBuilder::with_backend(backend.clone()).build();
+	let mut client = TestClientBuilder::with_backend(backend.clone()).build();
 	let blockchain = backend.blockchain();

 	// G -> A1
-	let a1 = client.new_block(Default::default()).unwrap().bake().unwrap();
+	let a1 = client.new_block(Default::default()).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a1.clone()).unwrap();

 	// A1 -> A2
-	let a2 = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap().bake().unwrap();
+	let a2 = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a2.clone()).unwrap();

 	// A2 -> A3
-	let a3 = client.new_block_at(&BlockId::Hash(a2.hash()), Default::default()).unwrap().bake().unwrap();
+	let a3 = client.new_block_at(
+		&BlockId::Hash(a2.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a3.clone()).unwrap();

 	// A3 -> A4
-	let a4 = client.new_block_at(&BlockId::Hash(a3.hash()), Default::default()).unwrap().bake().unwrap();
+	let a4 = client.new_block_at(
+		&BlockId::Hash(a3.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a4.clone()).unwrap();

 	// A4 -> A5
-	let a5 = client.new_block_at(&BlockId::Hash(a4.hash()), Default::default()).unwrap().bake().unwrap();
+	let a5 = client.new_block_at(
+		&BlockId::Hash(a4.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a5.clone()).unwrap();

 	// A1 -> B2
-	let mut builder = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap();
 	// this push is required as otherwise B2 has the same hash as A2 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -189,19 +265,31 @@ pub fn test_children_for_backend<B: 'static>(backend: Arc<B>) where
 		amount: 41,
 		nonce: 0,
 	}).unwrap();
-	let b2 = builder.bake().unwrap();
+	let b2 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, b2.clone()).unwrap();

 	// B2 -> B3
-	let b3 = client.new_block_at(&BlockId::Hash(b2.hash()), Default::default()).unwrap().bake().unwrap();
+	let b3 = client.new_block_at(
+		&BlockId::Hash(b2.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, b3.clone()).unwrap();

 	// B3 -> B4
-	let b4 = client.new_block_at(&BlockId::Hash(b3.hash()), Default::default()).unwrap().bake().unwrap();
+	let b4 = client.new_block_at(
+		&BlockId::Hash(b3.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, b4.clone()).unwrap();

 	// // B2 -> C3
-	let mut builder = client.new_block_at(&BlockId::Hash(b2.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(b2.hash()),
+		Default::default(),
+		false,
+	).unwrap();
 	// this push is required as otherwise C3 has the same hash as B3 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -209,11 +297,15 @@ pub fn test_children_for_backend<B: 'static>(backend: Arc<B>) where
 		amount: 1,
 		nonce: 1,
 	}).unwrap();
-	let c3 = builder.bake().unwrap();
+	let c3 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, c3.clone()).unwrap();

 	// A1 -> D2
-	let mut builder = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap();
 	// this push is required as otherwise D2 has the same hash as B2 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -221,7 +313,7 @@ pub fn test_children_for_backend<B: 'static>(backend: Arc<B>) where
 		amount: 1,
 		nonce: 0,
 	}).unwrap();
-	let d2 = builder.bake().unwrap();
+	let d2 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, d2.clone()).unwrap();

 	let genesis_hash = client.chain_info().genesis_hash;
@@ -240,38 +332,61 @@ pub fn test_children_for_backend<B: 'static>(backend: Arc<B>) where
 }

 pub fn test_blockchain_query_by_number_gets_canonical<B: 'static>(backend: Arc<B>) where
-	B: LocalBackend<substrate_test_runtime::Block, Blake2Hasher>,
+	B: backend::LocalBackend<substrate_test_runtime::Block>,
+	// Rust bug: https://github.com/rust-lang/rust/issues/24159
+	<B as backend::Backend<substrate_test_runtime::Block>>::State:
+		sp_api::StateBackend<HasherFor<substrate_test_runtime::Block>>,
 {
 	// block tree:
 	// G -> A1 -> A2 -> A3 -> A4 -> A5
 	//		A1 -> B2 -> B3 -> B4
 	//			  B2 -> C3
 	//		A1 -> D2
-	let client = TestClientBuilder::with_backend(backend.clone()).build();
+	let mut client = TestClientBuilder::with_backend(backend.clone()).build();
 	let blockchain = backend.blockchain();

 	// G -> A1
-	let a1 = client.new_block(Default::default()).unwrap().bake().unwrap();
+	let a1 = client.new_block(Default::default()).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a1.clone()).unwrap();

 	// A1 -> A2
-	let a2 = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap().bake().unwrap();
+	let a2 = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a2.clone()).unwrap();

 	// A2 -> A3
-	let a3 = client.new_block_at(&BlockId::Hash(a2.hash()), Default::default()).unwrap().bake().unwrap();
+	let a3 = client.new_block_at(
+		&BlockId::Hash(a2.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a3.clone()).unwrap();

 	// A3 -> A4
-	let a4 = client.new_block_at(&BlockId::Hash(a3.hash()), Default::default()).unwrap().bake().unwrap();
+	let a4 = client.new_block_at(
+		&BlockId::Hash(a3.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a4.clone()).unwrap();

 	// A4 -> A5
-	let a5 = client.new_block_at(&BlockId::Hash(a4.hash()), Default::default()).unwrap().bake().unwrap();
+	let a5 = client.new_block_at(
+		&BlockId::Hash(a4.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, a5.clone()).unwrap();

 	// A1 -> B2
-	let mut builder = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap();
 	// this push is required as otherwise B2 has the same hash as A2 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -279,19 +394,31 @@ pub fn test_blockchain_query_by_number_gets_canonical<B: 'static>(backend: Arc<B
 		amount: 41,
 		nonce: 0,
 	}).unwrap();
-	let b2 = builder.bake().unwrap();
+	let b2 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, b2.clone()).unwrap();

 	// B2 -> B3
-	let b3 = client.new_block_at(&BlockId::Hash(b2.hash()), Default::default()).unwrap().bake().unwrap();
+	let b3 = client.new_block_at(
+		&BlockId::Hash(b2.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, b3.clone()).unwrap();

 	// B3 -> B4
-	let b4 = client.new_block_at(&BlockId::Hash(b3.hash()), Default::default()).unwrap().bake().unwrap();
+	let b4 = client.new_block_at(
+		&BlockId::Hash(b3.hash()),
+		Default::default(),
+		false,
+	).unwrap().build().unwrap().block;
 	client.import(BlockOrigin::Own, b4.clone()).unwrap();

 	// // B2 -> C3
-	let mut builder = client.new_block_at(&BlockId::Hash(b2.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(b2.hash()),
+		Default::default(),
+		false,
+	).unwrap();
 	// this push is required as otherwise C3 has the same hash as B3 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -299,11 +426,15 @@ pub fn test_blockchain_query_by_number_gets_canonical<B: 'static>(backend: Arc<B
 		amount: 1,
 		nonce: 1,
 	}).unwrap();
-	let c3 = builder.bake().unwrap();
+	let c3 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, c3.clone()).unwrap();

 	// A1 -> D2
-	let mut builder = client.new_block_at(&BlockId::Hash(a1.hash()), Default::default()).unwrap();
+	let mut builder = client.new_block_at(
+		&BlockId::Hash(a1.hash()),
+		Default::default(),
+		false,
+	).unwrap();
 	// this push is required as otherwise D2 has the same hash as B2 and won't get imported
 	builder.push_transfer(Transfer {
 		from: AccountKeyring::Alice.into(),
@@ -311,7 +442,7 @@ pub fn test_blockchain_query_by_number_gets_canonical<B: 'static>(backend: Arc<B
 		amount: 1,
 		nonce: 0,
 	}).unwrap();
-	let d2 = builder.bake().unwrap();
+	let d2 = builder.build().unwrap().block;
 	client.import(BlockOrigin::Own, d2.clone()).unwrap();

 	let genesis_hash = client.chain_info().genesis_hash;
@@ -954,17 +954,17 @@ mod tests {
 		DefaultTestClientBuilderExt, TestClientBuilder,
 		runtime::TestAPI,
 	};
-	use sp_runtime::{
-		generic::BlockId,
-		traits::ProvideRuntimeApi,
-	};
+	use sp_api::ProvideRuntimeApi;
+	use sp_runtime::generic::BlockId;
 	use sp_core::storage::well_known_keys::HEAP_PAGES;
 	use sp_state_machine::ExecutionStrategy;
 	use codec::Encode;

 	#[test]
 	fn returns_mutable_static() {
-		let client = TestClientBuilder::new().set_execution_strategy(ExecutionStrategy::AlwaysWasm).build();
+		let client = TestClientBuilder::new()
+			.set_execution_strategy(ExecutionStrategy::AlwaysWasm)
+			.build();
 		let runtime_api = client.runtime_api();
 		let block_id = BlockId::Number(client.chain_info().best_number);

@@ -1013,31 +1013,31 @@ mod tests {
 		// This tests that the on-chain HEAP_PAGES parameter is respected.

 		// Create a client devoting only 8 pages of wasm memory. This gives us ~512k of heap memory.
-		let client = TestClientBuilder::new()
+		let mut client = TestClientBuilder::new()
 			.set_execution_strategy(ExecutionStrategy::AlwaysWasm)
 			.set_heap_pages(8)
 			.build();
-		let runtime_api = client.runtime_api();
 		let block_id = BlockId::Number(client.chain_info().best_number);

 		// Try to allocate 1024k of memory on heap. This is going to fail since it is twice larger
 		// than the heap.
-		let ret = runtime_api.vec_with_capacity(&block_id, 1048576);
+		let ret = client.runtime_api().vec_with_capacity(&block_id, 1048576);
 		assert!(ret.is_err());

 		// Create a block that sets the `:heap_pages` to 32 pages of memory which corresponds to
 		// ~2048k of heap memory.
-		let new_block_id = {
+		let (new_block_id, block) = {
 			let mut builder = client.new_block(Default::default()).unwrap();
 			builder.push_storage_change(HEAP_PAGES.to_vec(), Some(32u64.encode())).unwrap();
-			let block = builder.bake().unwrap();
+			let block = builder.build().unwrap().block;
 			let hash = block.header.hash();
-			client.import(BlockOrigin::Own, block).unwrap();
-			BlockId::Hash(hash)
+			(BlockId::Hash(hash), block)
 		};

+		client.import(BlockOrigin::Own, block).unwrap();
+
 		// Allocation of 1024k while having ~2048k should succeed.
-		let ret = runtime_api.vec_with_capacity(&new_block_id, 1048576);
+		let ret = client.runtime_api().vec_with_capacity(&new_block_id, 1048576);
 		assert!(ret.is_ok());
 	}

@@ -20,12 +20,12 @@
 use sp_std::prelude::*;
 use sp_io::{
 	storage::root as storage_root, storage::changes_root as storage_changes_root,
-	hashing::blake2_256,
+	hashing::blake2_256, trie,
 };
 use frame_support::storage;
 use frame_support::{decl_storage, decl_module};
 use sp_runtime::{
-	traits::{Hash as HashT, BlakeTwo256, Header as _}, generic, ApplyExtrinsicResult,
+	traits::Header as _, generic, ApplyExtrinsicResult,
 	transaction_validity::{
 		TransactionValidity, ValidTransaction, InvalidTransaction, TransactionValidityError,
 	},
@@ -109,60 +109,36 @@ pub fn execute_block(mut block: Block) {
 	execute_block_with_state_root_handler(&mut block, Mode::Verify);
 }

-fn execute_block_with_state_root_handler(
-	block: &mut Block,
-	mode: Mode,
-) {
+fn execute_block_with_state_root_handler(block: &mut Block, mode: Mode) {
 	let header = &mut block.header;

-	// check transaction trie root represents the transactions.
-	let txs = block.extrinsics.iter().map(Encode::encode).collect::<Vec<_>>();
-	let txs_root = BlakeTwo256::ordered_trie_root(txs);
-	info_expect_equal_hash(&txs_root, &header.extrinsics_root);
-	if let Mode::Overwrite = mode {
-		header.extrinsics_root = txs_root;
-	} else {
-		assert!(txs_root == header.extrinsics_root, "Transaction trie root must be valid.");
-	}
-
-	// try to read something that depends on current header digest
-	// so that it'll be included in execution proof
-	if let Some(generic::DigestItem::Other(v)) = header.digest().logs().iter().next() {
-		let _: Option<u32> = storage::unhashed::get(&v);
-	}
+	initialize_block(header);

 	// execute transactions
-	block.extrinsics.iter().enumerate().for_each(|(i, e)| {
-		storage::unhashed::put(well_known_keys::EXTRINSIC_INDEX, &(i as u32));
-		let _ = execute_transaction_backend(e).unwrap_or_else(|_| panic!("Invalid transaction"));
-		storage::unhashed::kill(well_known_keys::EXTRINSIC_INDEX);
+	block.extrinsics.iter().for_each(|e| {
+		let _ = execute_transaction(e.clone()).unwrap_or_else(|_| panic!("Invalid transaction"));
 	});

-	let o_new_authorities = <NewAuthorities>::take();
-	let storage_root = Hash::decode(&mut &storage_root()[..])
-		.expect("`storage_root` is a valid hash");
+	let new_header = finalize_block();

 	if let Mode::Overwrite = mode {
-		header.state_root = storage_root;
+		header.state_root = new_header.state_root;
 	} else {
-		// check storage root.
-		info_expect_equal_hash(&storage_root, &header.state_root);
-		assert!(storage_root == header.state_root, "Storage root must match that calculated.");
-	}
-
-	// check digest
-	let digest = &mut header.digest;
-	if let Some(storage_changes_root) = storage_changes_root(&header.parent_hash.encode()) {
-		digest.push(
-			generic::DigestItem::ChangesTrieRoot(
-				Hash::decode(&mut &storage_changes_root[..])
-					.expect("`storage_changes_root` is a valid hash")
-			)
+		info_expect_equal_hash(&new_header.state_root, &header.state_root);
+		assert!(
+			new_header.state_root == header.state_root,
+			"Storage root must match that calculated.",
 		);
 	}
-	if let Some(new_authorities) = o_new_authorities {
-		digest.push(generic::DigestItem::Consensus(*b"aura", new_authorities.encode()));
-		digest.push(generic::DigestItem::Consensus(*b"babe", new_authorities.encode()));
+
+	if let Mode::Overwrite = mode {
+		header.extrinsics_root = new_header.extrinsics_root;
+	} else {
+		info_expect_equal_hash(&new_header.extrinsics_root, &header.extrinsics_root);
+		assert!(
+			new_header.extrinsics_root == header.extrinsics_root,
+			"Transaction trie root must be valid.",
+		);
 	}
 }

@@ -224,7 +200,7 @@ pub fn execute_transaction(utx: Extrinsic) -> ApplyExtrinsicResult {
 pub fn finalize_block() -> Header {
 	let extrinsic_index: u32 = storage::unhashed::take(well_known_keys::EXTRINSIC_INDEX).unwrap();
 	let txs: Vec<_> = (0..extrinsic_index).map(ExtrinsicData::take).collect();
-	let extrinsics_root = BlakeTwo256::ordered_trie_root(txs).into();
+	let extrinsics_root = trie::blake2_256_ordered_root(txs).into();
 	let number = <Number>::take().expect("Number is set by `initialize_block`");
 	let parent_hash = <ParentHash>::take();
 	let mut digest = <StorageDigest>::take().expect("StorageDigest is set by `initialize_block`");