merkleize chunks and yield branches for each

erasure-coding/Cargo.toml

@@ -8,3 +8,5 @@ edition = "2018"
 polkadot-primitives = { path = "../primitives" }
 reed-solomon-erasure = "3.1.1"
 parity-codec = "2.1"
+substrate-primitives = { git = "https://github.com/paritytech/substrate" }
+substrate-trie = { git = "https://github.com/paritytech/substrate" }

erasure-coding/src/lib.rs

@@ -27,10 +27,15 @@
 extern crate polkadot_primitives as primitives;
 extern crate reed_solomon_erasure as reed_solomon;
 extern crate parity_codec as codec;
+extern crate substrate_primitives;
+extern crate substrate_trie as trie;
 
 use codec::{Encode, Decode};
 use reed_solomon::ReedSolomon;
+use primitives::{Hash as H256, BlakeTwo256, HashT};
 use primitives::parachain::{BlockData, Extrinsic};
+use substrate_primitives::Blake2Hasher;
+use trie::{MemoryDB, Trie, TrieMut, TrieDB, TrieDBMut};
 
 // unfortunate requirement due to use of GF(256) in the reed-solomon
 // implementation.
@@ -171,6 +176,74 @@ pub fn reconstruct<'a, I: 'a>(n_validators: usize, chunks: I)
 	}).ok_or_else(|| Error::BadPayload)
 }
 
+/// An iterator that yields merkle branches and chunk data for all chunks to
+/// be sent to other validators.
+pub struct Branches<'a> {
+	trie_storage: MemoryDB<Blake2Hasher>,
+	root: H256,
+	chunks: Vec<&'a [u8]>,
+	current_pos: usize,
+}
+
+impl<'a> Branches<'a> {
+	/// Get the trie root.
+	pub fn root(&self) -> H256 { self.root.clone() }
+}
+
+impl<'a> Iterator for Branches<'a> {
+	type Item = (Vec<Vec<u8>>, &'a [u8]);
+
+	fn next(&mut self) -> Option<Self::Item> {
+		use trie::Recorder;
+
+		let trie = TrieDB::new(&self.trie_storage, &self.root)
+			.expect("`Branches` is only created with a valid memorydb that contains all nodes for the trie with given root; qed");
+
+		let mut recorder = Recorder::new();
+		let res = (self.current_pos as u32).using_encoded(|s|
+			trie.get_with(s, &mut recorder)
+		);
+
+		match res.expect("all nodes in trie present; qed") {
+			Some(_) => {
+				let nodes = recorder.drain().into_iter().map(|r| r.data).collect();
+				let chunk = &self.chunks.get(self.current_pos)
+					.expect("there is a one-to-one mapping of chunks to valid merkle branches; qed");
+
+				self.current_pos += 1;
+				Some((nodes, chunk))
+			}
+			None => None,
+		}
+	}
+}
+
+/// Construct a trie from chunks of an erasure-coded value. This returns the root hash and an
+/// iterator of merkle proofs, one for each validator.
+pub fn branches<'a>(chunks: Vec<&'a [u8]>) -> Branches<'a> {
+	let mut trie_storage: MemoryDB<Blake2Hasher> = MemoryDB::default();
+	let mut root = H256::default();
+
+	// construct trie mapping each chunk's index to its hash.
+	{
+		let mut trie = TrieDBMut::new(&mut trie_storage, &mut root);
+		for (i, &chunk) in chunks.iter().enumerate() {
+			(i as u32).using_encoded(|encoded_index| {
+				let chunk_hash = BlakeTwo256::hash(chunk);
+				trie.insert(encoded_index, chunk_hash.as_ref())
+					.expect("a fresh trie stored in memory cannot have errors loading nodes; qed");
+			})
+		}
+	}
+
+	Branches {
+		trie_storage,
+		root,
+		chunks,
+		current_pos: 0,
+	}
+}
+
 // input for `parity_codec` which draws data from the data shards
 struct ShardInput {
 	shards: Vec<Option<Box<[u8]>>>,

primitives/src/lib.rs

@@ -44,7 +44,9 @@ extern crate serde;
 extern crate substrate_client;
 
 use rstd::prelude::*;
-use runtime_primitives::{generic, traits::{Extrinsic, BlakeTwo256}};
+use runtime_primitives::{generic, traits::Extrinsic};
+
+pub use runtime_primitives::traits::{BlakeTwo256, Hash as HashT};
 
 pub mod parachain;