Fetching multiple remote values using single message (#3612)

* multiple remote values in single message

* keys_str: String ->keys_str: closure

substrate/core/client/src/cht.rs

@@ -23,8 +23,6 @@
 //! root has. A correct proof implies that the claimed block is identical to the one
 //! we discarded.
 
-use std::collections::HashSet;
-
 use hash_db;
 use codec::Encode;
 use trie;
@@ -105,15 +103,10 @@ pub fn build_proof<Header, Hasher, BlocksI, HashesI>(
 	let mut storage = InMemoryState::<Hasher>::default().update(transaction);
 	let trie_storage = storage.as_trie_backend()
 		.expect("InMemoryState::as_trie_backend always returns Some; qed");
-	let mut total_proof = HashSet::new();
-	for block in blocks.into_iter() {
-		debug_assert_eq!(block_to_cht_number(cht_size, block), Some(cht_num));
-
-		let (value, proof) = prove_read_on_trie_backend(trie_storage, &encode_cht_key(block))?;
-		assert!(value.is_some(), "we have just built trie that includes the value for block");
-		total_proof.extend(proof);
-	}
-	Ok(total_proof.into_iter().collect())
+	prove_read_on_trie_backend(
+		trie_storage,
+		blocks.into_iter().map(|number| encode_cht_key(number)),
+	).map_err(ClientError::Execution)
 }
 
 /// Check CHT-based header proof.
@@ -128,9 +121,21 @@ pub fn check_proof<Header, Hasher>(
 		Hasher: hash_db::Hasher,
 		Hasher::Out: Ord,
 {
-	do_check_proof::<Header, Hasher, _>(local_root, local_number, remote_hash, move |local_root, local_cht_key|
-		read_proof_check::<Hasher>(local_root, remote_proof,
-			local_cht_key).map_err(|e| ClientError::from(e)))
+	do_check_proof::<Header, Hasher, _>(
+		local_root,
+		local_number,
+		remote_hash,
+		move |local_root, local_cht_key|
+			read_proof_check::<Hasher, _>(
+				local_root,
+				remote_proof,
+				::std::iter::once(local_cht_key),
+			)
+			.map(|mut map| map
+				.remove(local_cht_key)
+				.expect("checked proof of local_cht_key; qed"))
+			.map_err(|e| ClientError::from(e)),
+	)
 }
 
 /// Check CHT-based header proof on pre-created proving backend.
@@ -145,9 +150,16 @@ pub fn check_proof_on_proving_backend<Header, Hasher>(
 		Hasher: hash_db::Hasher,
 		Hasher::Out: Ord,
 {
-	do_check_proof::<Header, Hasher, _>(local_root, local_number, remote_hash, |_, local_cht_key|
-		read_proof_check_on_proving_backend::<Hasher>(
-			proving_backend, local_cht_key).map_err(|e| ClientError::from(e)))
+	do_check_proof::<Header, Hasher, _>(
+		local_root,
+		local_number,
+		remote_hash,
+		|_, local_cht_key|
+			read_proof_check_on_proving_backend::<Hasher>(
+				proving_backend,
+				local_cht_key,
+			).map_err(|e| ClientError::from(e)),
+	)
 }
 
 /// Check CHT-based header proof using passed checker function.

substrate/core/client/src/client.rs

@@ -436,24 +436,28 @@ impl<B, E, Block, RA> Client<B, E, Block, RA> where
 	}
 
 	/// Reads storage value at a given block + key, returning read proof.
-	pub fn read_proof(&self, id: &BlockId<Block>, key: &[u8]) -> error::Result<Vec<Vec<u8>>> {
+	pub fn read_proof<I>(&self, id: &BlockId<Block>, keys: I) -> error::Result<Vec<Vec<u8>>> where
+		I: IntoIterator,
+		I::Item: AsRef<[u8]>,
+	{
 		self.state_at(id)
-			.and_then(|state| prove_read(state, key)
-				.map(|(_, proof)| proof)
+			.and_then(|state| prove_read(state, keys)
 				.map_err(Into::into))
 	}
 
 	/// Reads child storage value at a given block + storage_key + key, returning
 	/// read proof.
-	pub fn read_child_proof(
+	pub fn read_child_proof<I>(
 		&self,
 		id: &BlockId<Block>,
 		storage_key: &[u8],
-		key: &[u8]
-	) -> error::Result<Vec<Vec<u8>>> {
+		keys: I,
+	) -> error::Result<Vec<Vec<u8>>> where
+		I: IntoIterator,
+		I::Item: AsRef<[u8]>,
+	{
 		self.state_at(id)
-			.and_then(|state| prove_child_read(state, storage_key, key)
-				.map(|(_, proof)| proof)
+			.and_then(|state| prove_child_read(state, storage_key, keys)
 				.map_err(Into::into))
 	}
 

substrate/core/client/src/light/fetcher.rs

@@ -17,7 +17,7 @@
 //! Light client data fetcher. Fetches requested data from remote full nodes.
 
 use std::sync::Arc;
-use std::collections::BTreeMap;
+use std::collections::{BTreeMap, HashMap};
 use std::marker::PhantomData;
 use std::future::Future;
 
@@ -73,7 +73,7 @@ pub struct RemoteReadRequest<Header: HeaderT> {
 	/// Header of block at which read is performed.
 	pub header: Header,
 	/// Storage key to read.
-	pub key: Vec<u8>,
+	pub keys: Vec<Vec<u8>>,
 	/// Number of times to retry request. None means that default RETRY_COUNT is used.
 	pub retry_count: Option<usize>,
 }
@@ -88,7 +88,7 @@ pub struct RemoteReadChildRequest<Header: HeaderT> {
 	/// Storage key for child.
 	pub storage_key: Vec<u8>,
 	/// Child storage key to read.
-	pub key: Vec<u8>,
+	pub keys: Vec<Vec<u8>>,
 	/// Number of times to retry request. None means that default RETRY_COUNT is used.
 	pub retry_count: Option<usize>,
 }
@@ -147,7 +147,7 @@ pub trait Fetcher<Block: BlockT>: Send + Sync {
 	/// Remote header future.
 	type RemoteHeaderResult: Future<Output = Result<Block::Header, ClientError>> + Send + 'static;
 	/// Remote storage read future.
-	type RemoteReadResult: Future<Output = Result<Option<Vec<u8>>, ClientError>> + Send + 'static;
+	type RemoteReadResult: Future<Output = Result<HashMap<Vec<u8>, Option<Vec<u8>>>, ClientError>> + Send + 'static;
 	/// Remote call result future.
 	type RemoteCallResult: Future<Output = Result<Vec<u8>, ClientError>> + Send + 'static;
 	/// Remote changes result future.
@@ -193,13 +193,13 @@ pub trait FetchChecker<Block: BlockT>: Send + Sync {
 		&self,
 		request: &RemoteReadRequest<Block::Header>,
 		remote_proof: Vec<Vec<u8>>
-	) -> ClientResult<Option<Vec<u8>>>;
+	) -> ClientResult<HashMap<Vec<u8>, Option<Vec<u8>>>>;
 	/// Check remote storage read proof.
 	fn check_read_child_proof(
 		&self,
 		request: &RemoteReadChildRequest<Block::Header>,
 		remote_proof: Vec<Vec<u8>>
-	) -> ClientResult<Option<Vec<u8>>>;
+	) -> ClientResult<HashMap<Vec<u8>, Option<Vec<u8>>>>;
 	/// Check remote method execution proof.
 	fn check_execution_proof(
 		&self,
@@ -395,23 +395,26 @@ impl<E, Block, H, S> FetchChecker<Block> for LightDataChecker<E, H, Block, S>
 	fn check_read_proof(
 		&self,
 		request: &RemoteReadRequest<Block::Header>,
-		remote_proof: Vec<Vec<u8>>
-	) -> ClientResult<Option<Vec<u8>>> {
-		read_proof_check::<H>(convert_hash(request.header.state_root()), remote_proof, &request.key)
-			.map_err(Into::into)
+		remote_proof: Vec<Vec<u8>>,
+	) -> ClientResult<HashMap<Vec<u8>, Option<Vec<u8>>>> {
+		read_proof_check::<H, _>(
+			convert_hash(request.header.state_root()),
+			remote_proof,
+			request.keys.iter(),
+		).map_err(Into::into)
 	}
 
 	fn check_read_child_proof(
 		&self,
 		request: &RemoteReadChildRequest<Block::Header>,
 		remote_proof: Vec<Vec<u8>>
-	) -> ClientResult<Option<Vec<u8>>> {
-		read_child_proof_check::<H>(
+	) -> ClientResult<HashMap<Vec<u8>, Option<Vec<u8>>>> {
+		read_child_proof_check::<H, _>(
 			convert_hash(request.header.state_root()),
 			remote_proof,
 			&request.storage_key,
-			&request.key)
-			.map_err(Into::into)
+			request.keys.iter(),
+		).map_err(Into::into)
 	}
 
 	fn check_execution_proof(
@@ -529,7 +532,7 @@ pub mod tests {
 
 	impl Fetcher<Block> for OkCallFetcher {
 		type RemoteHeaderResult = Ready<Result<Header, ClientError>>;
-		type RemoteReadResult = Ready<Result<Option<Vec<u8>>, ClientError>>;
+		type RemoteReadResult = Ready<Result<HashMap<Vec<u8>, Option<Vec<u8>>>, ClientError>>;
 		type RemoteCallResult = Ready<Result<Vec<u8>, ClientError>>;
 		type RemoteChangesResult = Ready<Result<Vec<(NumberFor<Block>, u32)>, ClientError>>;
 		type RemoteBodyResult = Ready<Result<Vec<Extrinsic>, ClientError>>;
@@ -579,7 +582,10 @@ pub mod tests {
 		let heap_pages = remote_client.storage(&remote_block_id, &StorageKey(well_known_keys::HEAP_PAGES.to_vec()))
 			.unwrap()
 			.and_then(|v| Decode::decode(&mut &v.0[..]).ok()).unwrap();
-		let remote_read_proof = remote_client.read_proof(&remote_block_id, well_known_keys::HEAP_PAGES).unwrap();
+		let remote_read_proof = remote_client.read_proof(
+			&remote_block_id,
+			&[well_known_keys::HEAP_PAGES],
+		).unwrap();
 
 		// check remote read proof locally
 		let local_storage = InMemoryBlockchain::<Block>::new();
@@ -618,7 +624,7 @@ pub mod tests {
 		let remote_read_proof = remote_client.read_child_proof(
 			&remote_block_id,
 			b":child_storage:default:child1",
-			b"key1",
+			&[b"key1"],
 		).unwrap();
 
 		// check locally
@@ -676,9 +682,9 @@ pub mod tests {
 		assert_eq!((&local_checker as &dyn FetchChecker<Block>).check_read_proof(&RemoteReadRequest::<Header> {
 			block: remote_block_header.hash(),
 			header: remote_block_header,
-			key: well_known_keys::HEAP_PAGES.to_vec(),
+			keys: vec![well_known_keys::HEAP_PAGES.to_vec()],
 			retry_count: None,
-		}, remote_read_proof).unwrap().unwrap()[0], heap_pages as u8);
+		}, remote_read_proof).unwrap().remove(well_known_keys::HEAP_PAGES).unwrap().unwrap()[0], heap_pages as u8);
 	}
 
 	#[test]
@@ -694,11 +700,11 @@ pub mod tests {
 				block: remote_block_header.hash(),
 				header: remote_block_header,
 				storage_key: b":child_storage:default:child1".to_vec(),
-				key: b"key1".to_vec(),
+				keys: vec![b"key1".to_vec()],
 				retry_count: None,
 			},
 			remote_read_proof
-		).unwrap().unwrap(), result);
+		).unwrap().remove(b"key1".as_ref()).unwrap().unwrap(), result);
 	}
 
 	#[test]