Trie simplification. (#2815)

* switch to simple codec, trie broken for now * Actualy use trie_root_noext * align some hash, failing test on EMCH comment * Fix trie code over layout instead of hash, revert legacy code for legacy mainnet ?? * stub behind LayOut * fix no_std * temp solution for legacy trie behind feature legacy-key in various crate * use remote project * rc client db need prefix * update trie deps * bum spec runtime version * Removing legacy as default. * Switch mode to non legacy. * bump runtime version * Remove legacy trie compatibility features. * fix warning * bump version * change hash on new test. * Move dependency (#11 trie PR) patched to a parity repo. Bench reverted to correct hasher. Some renaming and doc improvments. * ChildBitmap renaming to BitMap. * Renaming of LayOut to Layout. * formatting. * Removing abreviation such as _ix nb_ or bm. * Update deps and apply renaming 'Buff' -> 'Buffer'. * Align to latest trie crates naming changes. * Update trie dependency. * Update trie dependency. * change block_import test hash * update trie deps (trie use new scale codec but it does not seems to be an issue). * update to use latest trie version (no mgmt of multiple radix). * tabify * Restoring test to 10 000. * Use published crate, trie bench is currently down until publishing (require another pr to update version). * Update trie-bench.
2026-06-12 10:01:17 +00:00 · 2019-08-02 19:51:59 +02:00
parent 7927e80bc6
commit da8b91ae7b
43 changed files with 892 additions and 590 deletions
@@ -33,57 +33,104 @@ pub use trie_stream::TrieStream;
 /// The Substrate format implementation of `NodeCodec`.
 pub use node_codec::NodeCodec;
 /// Various re-exports from the `trie-db` crate.
-pub use trie_db::{Trie, TrieMut, DBValue, Recorder, Query};
+pub use trie_db::{Trie, TrieMut, DBValue, Recorder, CError,
+	Query, TrieLayout, TrieConfiguration, nibble_ops};
 /// Various re-exports from the `memory-db` crate.
-pub use memory_db::{KeyFunction, prefixed_key};
+pub use memory_db::KeyFunction;
+pub use memory_db::prefixed_key;
 /// Various re-exports from the `hash-db` crate.
-pub use hash_db::HashDB as HashDBT;
+pub use hash_db::{HashDB as HashDBT, EMPTY_PREFIX};

-/// As in `trie_db`, but less generic, error type for the crate.
-pub type TrieError<H> = trie_db::TrieError<H, Error>;
-/// As in `hash_db`, but less generic, trait exposed.
+
+#[derive(Default)]
+/// substrate trie layout
+pub struct Layout<H>(rstd::marker::PhantomData<H>);
+
+impl<H: Hasher> TrieLayout for Layout<H> {
+	const USE_EXTENSION: bool = false;
+	type Hash = H;
+	type Codec = NodeCodec<Self::Hash>;
+}
+
+impl<H: Hasher> TrieConfiguration for Layout<H> {
+	fn trie_root<I, A, B>(input: I) -> <Self::Hash as Hasher>::Out where
+		I: IntoIterator<Item = (A, B)>,
+		A: AsRef<[u8]> + Ord,
+		B: AsRef<[u8]>,
+	{
+		trie_root::trie_root_no_extension::<H, TrieStream, _, _, _>(input)
+	}
+
+	fn trie_root_unhashed<I, A, B>(input: I) -> Vec<u8> where
+		I: IntoIterator<Item = (A, B)>,
+		A: AsRef<[u8]> + Ord,
+		B: AsRef<[u8]>,
+	{
+		trie_root::unhashed_trie_no_extension::<H, TrieStream, _, _, _>(input)
+	}
+
+	fn encode_index(input: u32) -> Vec<u8> {
+		codec::Encode::encode(&codec::Compact(input))
+	}
+}
+
+/// TrieDB error over `TrieConfiguration` trait.
+pub type TrieError<L> = trie_db::TrieError<TrieHash<L>, CError<L>>;
+/// Reexport from `hash_db`, with genericity set for `Hasher` trait.
 pub trait AsHashDB<H: Hasher>: hash_db::AsHashDB<H, trie_db::DBValue> {}
 impl<H: Hasher, T: hash_db::AsHashDB<H, trie_db::DBValue>> AsHashDB<H> for T {}
-/// As in `hash_db`, but less generic, trait exposed.
+/// Reexport from `hash_db`, with genericity set for `Hasher` trait.
 pub type HashDB<'a, H> = dyn hash_db::HashDB<H, trie_db::DBValue> + 'a;
-/// As in `hash_db`, but less generic, trait exposed.
+/// Reexport from `hash_db`, with genericity set for key only.
 pub type PlainDB<'a, K> = dyn hash_db::PlainDB<K, trie_db::DBValue> + 'a;
-/// As in `memory_db::MemoryDB` that uses prefixed storage key scheme.
+/// Reexport from `hash_db`, with genericity set for `Hasher` trait.
+/// This uses a `KeyFunction` for prefixing keys internally (avoiding
+/// key conflict for non random keys).
 pub type PrefixedMemoryDB<H> = memory_db::MemoryDB<H, memory_db::PrefixedKey<H>, trie_db::DBValue>;
-/// As in `memory_db::MemoryDB` that uses prefixed storage key scheme.
+/// Reexport from `hash_db`, with genericity set for `Hasher` trait.
+/// This uses the `KeyFunction` for prefixing keys internally (avoiding
+/// This uses a noops `KeyFunction` (key addressing must be hashed or using
+/// an encoding scheme that avoid key conflict).
 pub type MemoryDB<H> = memory_db::MemoryDB<H, memory_db::HashKey<H>, trie_db::DBValue>;
-/// As in `memory_db`, but less generic, trait exposed.
+/// Reexport from `hash_db`, with genericity set for `Hasher` trait.
 pub type GenericMemoryDB<H, KF> = memory_db::MemoryDB<H, KF, trie_db::DBValue>;

 /// Persistent trie database read-access interface for the a given hasher.
-pub type TrieDB<'a, H> = trie_db::TrieDB<'a, H, NodeCodec<H>>;
+pub type TrieDB<'a, L> = trie_db::TrieDB<'a, L>;
 /// Persistent trie database write-access interface for the a given hasher.
-pub type TrieDBMut<'a, H> = trie_db::TrieDBMut<'a, H, NodeCodec<H>>;
+pub type TrieDBMut<'a, L> = trie_db::TrieDBMut<'a, L>;
 /// Querying interface, as in `trie_db` but less generic.
-pub type Lookup<'a, H, Q> = trie_db::Lookup<'a, H, NodeCodec<H>, Q>;
+pub type Lookup<'a, L, Q> = trie_db::Lookup<'a, L, Q>;
+/// Hash type for a trie layout.
+pub type TrieHash<L> = <<L as TrieLayout>::Hash as Hasher>::Out;

-/// Determine a trie root given its ordered contents, closed form.
-pub fn trie_root<H: Hasher, I, A, B>(input: I) -> H::Out where
-	I: IntoIterator<Item = (A, B)>,
-	A: AsRef<[u8]> + Ord,
-	B: AsRef<[u8]>,
-{
-	trie_root::trie_root::<H, TrieStream, _, _, _>(input)
+/// This module is for non generic definition of trie type.
+/// Only the `Hasher` trait is generic in this case.
+pub mod trie_types {
+	pub type Layout<H> = super::Layout<H>;
+	/// Persistent trie database read-access interface for the a given hasher.
+	pub type TrieDB<'a, H> = super::TrieDB<'a, Layout<H>>;
+	/// Persistent trie database write-access interface for the a given hasher.
+	pub type TrieDBMut<'a, H> = super::TrieDBMut<'a, Layout<H>>;
+	/// Querying interface, as in `trie_db` but less generic.
+	pub type Lookup<'a, H, Q> = trie_db::Lookup<'a, Layout<H>, Q>;
+	/// As in `trie_db`, but less generic, error type for the crate.
+	pub type TrieError<H> = trie_db::TrieError<H, super::Error>;
 }

 /// Determine a trie root given a hash DB and delta values.
-pub fn delta_trie_root<H: Hasher, I, A, B, DB>(
+pub fn delta_trie_root<L: TrieConfiguration, I, A, B, DB>(
 	db: &mut DB,
-	mut root: H::Out,
+	mut root: TrieHash<L>,
 	delta: I
-) -> Result<H::Out, Box<TrieError<H::Out>>> where
+) -> Result<TrieHash<L>, Box<TrieError<L>>> where
 	I: IntoIterator<Item = (A, Option<B>)>,
 	A: AsRef<[u8]> + Ord,
 	B: AsRef<[u8]>,
-	DB: hash_db::HashDB<H, trie_db::DBValue>,
+	DB: hash_db::HashDB<L::Hash, trie_db::DBValue>,
 {
 	{
-		let mut trie = TrieDBMut::<H>::from_existing(&mut *db, &mut root)?;
+		let mut trie = TrieDBMut::<L>::from_existing(&mut *db, &mut root)?;

 		for (key, change) in delta {
 			match change {
@@ -97,45 +144,26 @@ pub fn delta_trie_root<H: Hasher, I, A, B, DB>(
 }

 /// Read a value from the trie.
-pub fn read_trie_value<H: Hasher, DB: hash_db::HashDBRef<H, trie_db::DBValue>>(
+pub fn read_trie_value<L: TrieConfiguration, DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>>(
 	db: &DB,
-	root: &H::Out,
+	root: &TrieHash<L>,
 	key: &[u8]
-) -> Result<Option<Vec<u8>>, Box<TrieError<H::Out>>> {
-	Ok(TrieDB::<H>::new(&*db, root)?.get(key).map(|x| x.map(|val| val.to_vec()))?)
+) -> Result<Option<Vec<u8>>, Box<TrieError<L>>> {
+	Ok(TrieDB::<L>::new(&*db, root)?.get(key).map(|x| x.map(|val| val.to_vec()))?)
 }

 /// Read a value from the trie with given Query.
-pub fn read_trie_value_with<H: Hasher, Q: Query<H, Item=DBValue>, DB: hash_db::HashDBRef<H, trie_db::DBValue>>(
+pub fn read_trie_value_with<
+	L: TrieConfiguration,
+	Q: Query<L::Hash, Item=DBValue>,
+	DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>
+>(
 	db: &DB,
-	root: &H::Out,
+	root: &TrieHash<L>,
 	key: &[u8],
 	query: Q
-) -> Result<Option<Vec<u8>>, Box<TrieError<H::Out>>> {
-	Ok(TrieDB::<H>::new(&*db, root)?.get_with(key, query).map(|x| x.map(|val| val.to_vec()))?)
-}
-
-/// Determine a trie root node's data given its ordered contents, closed form.
-pub fn unhashed_trie<H: Hasher, I, A, B>(input: I) -> Vec<u8> where
-	I: IntoIterator<Item = (A, B)>,
-	A: AsRef<[u8]> + Ord,
-	B: AsRef<[u8]>,
-{
-	trie_root::unhashed_trie::<H, TrieStream, _, _, _>(input)
-}
-
-/// A trie root formed from the items, with keys attached according to their
-/// compact-encoded index (using `parity-codec` crate).
-pub fn ordered_trie_root<H: Hasher, I, A>(input: I) -> H::Out
-where
-	I: IntoIterator<Item = A>,
-	A: AsRef<[u8]>,
-{
-	trie_root::<H, _, _, _>(input
-		.into_iter()
-		.enumerate()
-		.map(|(i, v)| (codec::Encode::encode(&codec::Compact(i as u32)), v))
-	)
+) -> Result<Option<Vec<u8>>, Box<TrieError<L>>> {
+	Ok(TrieDB::<L>::new(&*db, root)?.get_with(key, query).map(|x| x.map(|val| val.to_vec()))?)
 }

 /// Determine whether a child trie key is valid.
@@ -143,7 +171,7 @@ where
 /// For now, the only valid child trie key is `:child_storage:default:`.
 ///
 /// `child_trie_root` and `child_delta_trie_root` can panic if invalid value is provided to them.
-pub fn is_child_trie_key_valid<H: Hasher>(storage_key: &[u8]) -> bool {
+pub fn is_child_trie_key_valid<L: TrieConfiguration>(storage_key: &[u8]) -> bool {
 	use primitives::storage::well_known_keys;
 	let has_right_prefix = storage_key.starts_with(b":child_storage:default:");
 	if has_right_prefix {
@@ -158,37 +186,42 @@ pub fn is_child_trie_key_valid<H: Hasher>(storage_key: &[u8]) -> bool {
 }

 /// Determine the default child trie root.
-pub fn default_child_trie_root<H: Hasher>(_storage_key: &[u8]) -> Vec<u8> {
-	trie_root::<H, _, Vec<u8>, Vec<u8>>(core::iter::empty()).as_ref().iter().cloned().collect()
+pub fn default_child_trie_root<L: TrieConfiguration>(_storage_key: &[u8]) -> Vec<u8> {
+	L::trie_root::<_, Vec<u8>, Vec<u8>>(core::iter::empty()).as_ref().iter().cloned().collect()
 }

-/// Determine a child trie root given its ordered contents, closed form. H is the default hasher, but a generic
-/// implementation may ignore this type parameter and use other hashers.
-pub fn child_trie_root<H: Hasher, I, A, B>(_storage_key: &[u8], input: I) -> Vec<u8> where
-	I: IntoIterator<Item = (A, B)>,
-	A: AsRef<[u8]> + Ord,
-	B: AsRef<[u8]>,
+/// Determine a child trie root given its ordered contents, closed form. H is the default hasher,
+/// but a generic implementation may ignore this type parameter and use other hashers.
+pub fn child_trie_root<L: TrieConfiguration, I, A, B>(_storage_key: &[u8], input: I) -> Vec<u8>
+	where
+		I: IntoIterator<Item = (A, B)>,
+		A: AsRef<[u8]> + Ord,
+		B: AsRef<[u8]>,
 {
-	trie_root::<H, _, _, _>(input).as_ref().iter().cloned().collect()
+	L::trie_root(input).as_ref().iter().cloned().collect()
 }

-/// Determine a child trie root given a hash DB and delta values. H is the default hasher, but a generic implementation may ignore this type parameter and use other hashers.
-pub fn child_delta_trie_root<H: Hasher, I, A, B, DB>(
+/// Determine a child trie root given a hash DB and delta values. H is the default hasher,
+/// but a generic implementation may ignore this type parameter and use other hashers.
+pub fn child_delta_trie_root<L: TrieConfiguration, I, A, B, DB>(
 	_storage_key: &[u8],
 	db: &mut DB,
 	root_vec: Vec<u8>,
 	delta: I
-) -> Result<Vec<u8>, Box<TrieError<H::Out>>> where
-	I: IntoIterator<Item = (A, Option<B>)>,
-	A: AsRef<[u8]> + Ord,
-	B: AsRef<[u8]>,
-	DB: hash_db::HashDB<H, trie_db::DBValue> + hash_db::PlainDB<H::Out, trie_db::DBValue>,
+) -> Result<Vec<u8>, Box<TrieError<L>>>
+	where
+		I: IntoIterator<Item = (A, Option<B>)>,
+		A: AsRef<[u8]> + Ord,
+		B: AsRef<[u8]>,
+		DB: hash_db::HashDB<L::Hash, trie_db::DBValue>
+			+ hash_db::PlainDB<TrieHash<L>, trie_db::DBValue>,
 {
-	let mut root = H::Out::default();
-	root.as_mut().copy_from_slice(&root_vec); // root is fetched from DB, not writable by runtime, so it's always valid.
+	let mut root = TrieHash::<L>::default();
+	// root is fetched from DB, not writable by runtime, so it's always valid.
+	root.as_mut().copy_from_slice(&root_vec);

 	{
-		let mut trie = TrieDBMut::<H>::from_existing(&mut *db, &mut root)?;
+		let mut trie = TrieDBMut::<L>::from_existing(&mut *db, &mut root)?;

 		for (key, change) in delta {
 			match change {
@@ -202,18 +235,21 @@ pub fn child_delta_trie_root<H: Hasher, I, A, B, DB>(
 }

 /// Call `f` for all keys in a child trie.
-pub fn for_keys_in_child_trie<H: Hasher, F: FnMut(&[u8]), DB>(
+pub fn for_keys_in_child_trie<L: TrieConfiguration, F: FnMut(&[u8]), DB>(
 	_storage_key: &[u8],
 	db: &DB,
 	root_slice: &[u8],
 	mut f: F
-) -> Result<(), Box<TrieError<H::Out>>> where
-	DB: hash_db::HashDBRef<H, trie_db::DBValue> + hash_db::PlainDBRef<H::Out, trie_db::DBValue>,
+) -> Result<(), Box<TrieError<L>>>
+	where
+		DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>
+			+ hash_db::PlainDBRef<TrieHash<L>, trie_db::DBValue>,
 {
-	let mut root = H::Out::default();
-	root.as_mut().copy_from_slice(root_slice); // root is fetched from DB, not writable by runtime, so it's always valid.
+	let mut root = TrieHash::<L>::default();
+	// root is fetched from DB, not writable by runtime, so it's always valid.
+	root.as_mut().copy_from_slice(root_slice);

-	let trie = TrieDB::<H>::new(&*db, &root)?;
+	let trie = TrieDB::<L>::new(&*db, &root)?;
 	let iter = trie.iter()?;

 	for x in iter {
@@ -225,14 +261,14 @@ pub fn for_keys_in_child_trie<H: Hasher, F: FnMut(&[u8]), DB>(
 }

 /// Record all keys for a given root.
-pub fn record_all_keys<H: Hasher, DB>(
+pub fn record_all_keys<L: TrieConfiguration, DB>(
 	db: &DB,
-	root: &H::Out,
-	recorder: &mut Recorder<H::Out>
-) -> Result<(), Box<TrieError<H::Out>>> where
-	DB: hash_db::HashDBRef<H, trie_db::DBValue>
+	root: &TrieHash<L>,
+	recorder: &mut Recorder<TrieHash<L>>
+) -> Result<(), Box<TrieError<L>>> where
+	DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>
 {
-	let trie = TrieDB::<H>::new(&*db, root)?;
+	let trie = TrieDB::<L>::new(&*db, root)?;
 	let iter = trie.iter()?;

 	for x in iter {
@@ -248,84 +284,49 @@ pub fn record_all_keys<H: Hasher, DB>(
 }

 /// Read a value from the child trie.
-pub fn read_child_trie_value<H: Hasher, DB>(
+pub fn read_child_trie_value<L: TrieConfiguration, DB>(
 	_storage_key: &[u8],
 	db: &DB,
 	root_slice: &[u8],
 	key: &[u8]
-) -> Result<Option<Vec<u8>>, Box<TrieError<H::Out>>> where
-	DB: hash_db::HashDBRef<H, trie_db::DBValue> + hash_db::PlainDBRef<H::Out, trie_db::DBValue>,
+) -> Result<Option<Vec<u8>>, Box<TrieError<L>>>
+	where
+		DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>
+			+ hash_db::PlainDBRef<TrieHash<L>, trie_db::DBValue>,
 {
-	let mut root = H::Out::default();
-	root.as_mut().copy_from_slice(root_slice); // root is fetched from DB, not writable by runtime, so it's always valid.
+	let mut root = TrieHash::<L>::default();
+	// root is fetched from DB, not writable by runtime, so it's always valid.
+	root.as_mut().copy_from_slice(root_slice);

-	Ok(TrieDB::<H>::new(&*db, &root)?.get(key).map(|x| x.map(|val| val.to_vec()))?)
+	Ok(TrieDB::<L>::new(&*db, &root)?.get(key).map(|x| x.map(|val| val.to_vec()))?)
 }

 /// Read a value from the child trie with given query.
-pub fn read_child_trie_value_with<H: Hasher, Q: Query<H, Item=DBValue>, DB>(
+pub fn read_child_trie_value_with<L: TrieConfiguration, Q: Query<L::Hash, Item=DBValue>, DB>(
 	_storage_key: &[u8],
 	db: &DB,
 	root_slice: &[u8],
 	key: &[u8],
 	query: Q
-) -> Result<Option<Vec<u8>>, Box<TrieError<H::Out>>> where
-	DB: hash_db::HashDBRef<H, trie_db::DBValue> + hash_db::PlainDBRef<H::Out, trie_db::DBValue>,
+) -> Result<Option<Vec<u8>>, Box<TrieError<L>>>
+	where
+		DB: hash_db::HashDBRef<L::Hash, trie_db::DBValue>
+			+ hash_db::PlainDBRef<TrieHash<L>, trie_db::DBValue>,
 {
-	let mut root = H::Out::default();
-	root.as_mut().copy_from_slice(root_slice); // root is fetched from DB, not writable by runtime, so it's always valid.
+	let mut root = TrieHash::<L>::default();
+	// root is fetched from DB, not writable by runtime, so it's always valid.
+	root.as_mut().copy_from_slice(root_slice);

-	Ok(TrieDB::<H>::new(&*db, &root)?.get_with(key, query).map(|x| x.map(|val| val.to_vec()))?)
+	Ok(TrieDB::<L>::new(&*db, &root)?.get_with(key, query).map(|x| x.map(|val| val.to_vec()))?)
 }

-// Utilities (not exported):
-
-const EMPTY_TRIE: u8 = 0;
-const LEAF_NODE_OFFSET: u8 = 1;
-const LEAF_NODE_BIG: u8 = 127;
-const EXTENSION_NODE_OFFSET: u8 = 128;
-const EXTENSION_NODE_BIG: u8 = 253;
-const BRANCH_NODE_NO_VALUE: u8 = 254;
-const BRANCH_NODE_WITH_VALUE: u8 = 255;
-const LEAF_NODE_THRESHOLD: u8 = LEAF_NODE_BIG - LEAF_NODE_OFFSET;
-const EXTENSION_NODE_THRESHOLD: u8 = EXTENSION_NODE_BIG - EXTENSION_NODE_OFFSET;	//125
-const LEAF_NODE_SMALL_MAX: u8 = LEAF_NODE_BIG - 1;
-const EXTENSION_NODE_SMALL_MAX: u8 = EXTENSION_NODE_BIG - 1;
-
-fn take<'a>(input: &mut &'a[u8], count: usize) -> Option<&'a[u8]> {
-	if input.len() < count {
-		return None
-	}
-	let r = &(*input)[..count];
-	*input = &(*input)[count..];
-	Some(r)
-}
-
-fn partial_to_key(partial: &[u8], offset: u8, big: u8) -> Vec<u8> {
-	let nibble_count = (partial.len() - 1) * 2 + if partial[0] & 16 == 16 { 1 } else { 0 };
-	let (first_byte_small, big_threshold) = (offset, (big - offset) as usize);
-	let mut output = [first_byte_small + nibble_count.min(big_threshold) as u8].to_vec();
-	if nibble_count >= big_threshold { output.push((nibble_count - big_threshold) as u8) }
-	if nibble_count % 2 == 1 {
-		output.push(partial[0] & 0x0f);
-	}
-	output.extend_from_slice(&partial[1..]);
-	output
-}
-
-fn branch_node(has_value: bool, has_children: impl Iterator<Item = bool>) -> [u8; 3] {
-	let first = if has_value {
-		BRANCH_NODE_WITH_VALUE
-	} else {
-		BRANCH_NODE_NO_VALUE
-	};
-	let mut bitmap: u16 = 0;
-	let mut cursor: u16 = 1;
-	for v in has_children {
-		if v { bitmap |= cursor }
-		cursor <<= 1;
-	}
-	[first, (bitmap % 256 ) as u8, (bitmap / 256 ) as u8]
+/// Constants used into trie simplification codec.
+mod trie_constants {
+	pub const EMPTY_TRIE: u8 = 0;
+	pub const NIBBLE_SIZE_BOUND: usize = u16::max_value() as usize;
+	pub const LEAF_PREFIX_MASK: u8 = 0b_01 << 6;
+	pub const BRANCH_WITHOUT_MASK: u8 = 0b_10 << 6;
+	pub const BRANCH_WITH_MASK: u8 = 0b_11 << 6;
 }

 #[cfg(test)]
@@ -334,19 +335,25 @@ mod tests {
 	use codec::{Encode, Compact};
 	use primitives::Blake2Hasher;
 	use hash_db::{HashDB, Hasher};
-	use trie_db::{DBValue, TrieMut, Trie};
+	use trie_db::{DBValue, TrieMut, Trie, NodeCodec as NodeCodecT};
 	use trie_standardmap::{Alphabet, ValueMode, StandardMap};
 	use hex_literal::hex;

-	fn check_equivalent(input: &Vec<(&[u8], &[u8])>) {
+	type Layout = super::Layout<Blake2Hasher>;
+
+	fn hashed_null_node<T: TrieConfiguration>() -> TrieHash<T> {
+		<T::Codec as NodeCodecT<_>>::hashed_null_node()
+	}
+
+	fn check_equivalent<T: TrieConfiguration>(input: &Vec<(&[u8], &[u8])>) {
 		{
-			let closed_form = trie_root::<Blake2Hasher, _, _, _>(input.clone());
-			let d = unhashed_trie::<Blake2Hasher, _, _, _>(input.clone());
+			let closed_form = T::trie_root(input.clone());
+			let d = T::trie_root_unhashed(input.clone());
 			println!("Data: {:#x?}, {:#x?}", d, Blake2Hasher::hash(&d[..]));
 			let persistent = {
 				let mut memdb = MemoryDB::default();
 				let mut root = Default::default();
-				let mut t = TrieDBMut::<Blake2Hasher>::new(&mut memdb, &mut root);
+				let mut t = TrieDBMut::<T>::new(&mut memdb, &mut root);
 				for (x, y) in input.iter().rev() {
 					t.insert(x, y).unwrap();
 				}
@@ -356,20 +363,22 @@ mod tests {
 		}
 	}

-	fn check_iteration(input: &Vec<(&[u8], &[u8])>) {
+	fn check_iteration<T: TrieConfiguration>(input: &Vec<(&[u8], &[u8])>) {
 		let mut memdb = MemoryDB::default();
 		let mut root = Default::default();
 		{
-			let mut t = TrieDBMut::<Blake2Hasher>::new(&mut memdb, &mut root);
+			let mut t = TrieDBMut::<T>::new(&mut memdb, &mut root);
 			for (x, y) in input.clone() {
 				t.insert(x, y).unwrap();
 			}
 		}
 		{
-			let t = TrieDB::<Blake2Hasher>::new(&mut memdb, &root).unwrap();
+			let t = TrieDB::<T>::new(&mut memdb, &root).unwrap();
 			assert_eq!(
 				input.iter().map(|(i, j)| (i.to_vec(), j.to_vec())).collect::<Vec<_>>(),
-				t.iter().unwrap().map(|x| x.map(|y| (y.0, y.1.to_vec())).unwrap()).collect::<Vec<_>>()
+				t.iter().unwrap()
+					.map(|x| x.map(|y| (y.0, y.1.to_vec())).unwrap())
+					.collect::<Vec<_>>()
 			);
 		}
 	}
@@ -377,11 +386,11 @@ mod tests {
 	#[test]
 	fn default_trie_root() {
 		let mut db = MemoryDB::default();
-		let mut root = <Blake2Hasher as Hasher>::Out::default();
-		let mut empty = TrieDBMut::<Blake2Hasher>::new(&mut db, &mut root);
+		let mut root = TrieHash::<Layout>::default();
+		let mut empty = TrieDBMut::<Layout>::new(&mut db, &mut root);
 		empty.commit();
 		let root1 = empty.root().as_ref().to_vec();
-		let root2: Vec<u8> = trie_root::<Blake2Hasher, _, Vec<u8>, Vec<u8>>(
+		let root2: Vec<u8> = Layout::trie_root::<_, Vec<u8>, Vec<u8>>(
 			std::iter::empty(),
 		).as_ref().iter().cloned().collect();

@@ -391,29 +400,35 @@ mod tests {
 	#[test]
 	fn empty_is_equivalent() {
 		let input: Vec<(&[u8], &[u8])> = vec![];
-		check_equivalent(&input);
-		check_iteration(&input);
+		check_equivalent::<Layout>(&input);
+		check_iteration::<Layout>(&input);
 	}

 	#[test]
 	fn leaf_is_equivalent() {
 		let input: Vec<(&[u8], &[u8])> = vec![(&[0xaa][..], &[0xbb][..])];
-		check_equivalent(&input);
-		check_iteration(&input);
+		check_equivalent::<Layout>(&input);
+		check_iteration::<Layout>(&input);
 	}

 	#[test]
 	fn branch_is_equivalent() {
-		let input: Vec<(&[u8], &[u8])> = vec![(&[0xaa][..], &[0x10][..]), (&[0xba][..], &[0x11][..])];
-		check_equivalent(&input);
-		check_iteration(&input);
+		let input: Vec<(&[u8], &[u8])> = vec![
+			(&[0xaa][..], &[0x10][..]),
+			(&[0xba][..], &[0x11][..]),
+		];
+		check_equivalent::<Layout>(&input);
+		check_iteration::<Layout>(&input);
 	}

 	#[test]
 	fn extension_and_branch_is_equivalent() {
-		let input: Vec<(&[u8], &[u8])> = vec![(&[0xaa][..], &[0x10][..]), (&[0xab][..], &[0x11][..])];
-		check_equivalent(&input);
-		check_iteration(&input);
+		let input: Vec<(&[u8], &[u8])> = vec![
+			(&[0xaa][..], &[0x10][..]),
+			(&[0xab][..], &[0x11][..]),
+		];
+		check_equivalent::<Layout>(&input);
+		check_iteration::<Layout>(&input);
 	}

 	#[test]
@@ -428,8 +443,8 @@ mod tests {
 		let mut d = st.make();
 		d.sort_unstable_by(|&(ref a, _), &(ref b, _)| a.cmp(b));
 		let dr = d.iter().map(|v| (&v.0[..], &v.1[..])).collect();
-		check_equivalent(&dr);
-		check_iteration(&dr);
+		check_equivalent::<Layout>(&dr);
+		check_iteration::<Layout>(&dr);
 	}

 	#[test]
@@ -439,8 +454,8 @@ mod tests {
 			(&[0xaa, 0xaa][..], &[0xaa][..]),
 			(&[0xaa, 0xbb][..], &[0xab][..])
 		];
-		check_equivalent(&input);
-		check_iteration(&input);
+		check_equivalent::<Layout>(&input);
+		check_iteration::<Layout>(&input);
 	}

 	#[test]
@@ -453,8 +468,8 @@ mod tests {
 			(&[0xbb, 0xbb][..], &[0xbb][..]),
 			(&[0xbb, 0xcc][..], &[0xbc][..]),
 		];
-		check_equivalent(&input);
-		check_iteration(&input);
+		check_equivalent::<Layout>(&input);
+		check_iteration::<Layout>(&input);
 	}

 	#[test]
@@ -463,8 +478,8 @@ mod tests {
 			(&[0xaa][..], &b"ABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABC"[..]),
 			(&[0xba][..], &[0x11][..]),
 		];
-		check_equivalent(&input);
-		check_iteration(&input);
+		check_equivalent::<Layout>(&input);
+		check_iteration::<Layout>(&input);
 	}

 	#[test]
@@ -473,16 +488,16 @@ mod tests {
 			(&[0xaa][..], &b"ABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABC"[..]),
 			(&[0xba][..], &b"ABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABCABC"[..])
 		];
-		check_equivalent(&input);
-		check_iteration(&input);
+		check_equivalent::<Layout>(&input);
+		check_iteration::<Layout>(&input);
 	}

-	fn populate_trie<'db>(
-		db: &'db mut dyn HashDB<Blake2Hasher, DBValue>,
-		root: &'db mut <Blake2Hasher as Hasher>::Out,
+	fn populate_trie<'db, T: TrieConfiguration>(
+		db: &'db mut dyn HashDB<T::Hash, DBValue>,
+		root: &'db mut TrieHash<T>,
 		v: &[(Vec<u8>, Vec<u8>)]
-	) -> TrieDBMut<'db, Blake2Hasher> {
-		let mut t = TrieDBMut::<Blake2Hasher>::new(db, root);
+	) -> TrieDBMut<'db, T> {
+		let mut t = TrieDBMut::<T>::new(db, root);
 		for i in 0..v.len() {
 			let key: &[u8]= &v[i].0;
 			let val: &[u8] = &v[i].1;
@@ -491,7 +506,10 @@ mod tests {
 		t
 	}

-	fn unpopulate_trie<'db>(t: &mut TrieDBMut<'db, Blake2Hasher>, v: &[(Vec<u8>, Vec<u8>)]) {
+	fn unpopulate_trie<'db, T: TrieConfiguration>(
+		t: &mut TrieDBMut<'db, T>,
+		v: &[(Vec<u8>, Vec<u8>)],
+	) {
 		for i in v {
 			let key: &[u8]= &i.0;
 			t.remove(key).unwrap();
@@ -513,10 +531,10 @@ mod tests {
 				count: 100,
 			}.make_with(seed.as_fixed_bytes_mut());

-			let real = trie_root::<Blake2Hasher,_, _, _>(x.clone());
+			let real = Layout::trie_root(x.clone());
 			let mut memdb = MemoryDB::default();
 			let mut root = Default::default();
-			let mut memtrie = populate_trie(&mut memdb, &mut root, &x);
+			let mut memtrie = populate_trie::<Layout>(&mut memdb, &mut root, &x);

 			memtrie.commit();
 			if *memtrie.root() != real {
@@ -528,17 +546,18 @@ mod tests {
 				}
 			}
 			assert_eq!(*memtrie.root(), real);
-			unpopulate_trie(&mut memtrie, &x);
+			unpopulate_trie::<Layout>(&mut memtrie, &x);
 			memtrie.commit();
-			if *memtrie.root() != <NodeCodec<Blake2Hasher> as trie_db::NodeCodec<Blake2Hasher>>::hashed_null_node() {
+			let hashed_null_node = hashed_null_node::<Layout>();
+			if *memtrie.root() != hashed_null_node {
 				println!("- TRIE MISMATCH");
 				println!("");
-				println!("{:?} vs {:?}", memtrie.root(), <NodeCodec<Blake2Hasher> as trie_db::NodeCodec<Blake2Hasher>>::hashed_null_node());
+				println!("{:?} vs {:?}", memtrie.root(), hashed_null_node);
 				for i in &x {
 					println!("{:#x?} -> {:#x?}", i.0, i.1);
 				}
 			}
-			assert_eq!(*memtrie.root(), <NodeCodec<Blake2Hasher> as trie_db::NodeCodec<Blake2Hasher>>::hashed_null_node());
+			assert_eq!(*memtrie.root(), hashed_null_node);
 		}
 	}

@@ -549,7 +568,7 @@ mod tests {
 	#[test]
 	fn codec_trie_empty() {
 		let input: Vec<(&[u8], &[u8])> = vec![];
-		let trie = unhashed_trie::<Blake2Hasher, _, _, _>(input);
+		let trie = Layout::trie_root_unhashed::<_, _, _>(input);
 		println!("trie: {:#x?}", trie);
 		assert_eq!(trie, vec![0x0]);
 	}
@@ -559,11 +578,10 @@ mod tests {
 		let input = vec![
 			(vec![0xaa], vec![0xbb])
 		];
-		let trie = unhashed_trie::<Blake2Hasher, _, _, _>(input);
+		let trie = Layout::trie_root_unhashed::<_, _, _>(input);
 		println!("trie: {:#x?}", trie);
-
 		assert_eq!(trie, vec![
-			0x03,					// leaf (0x01) with (+) key of 2 nibbles (0x02)
+			0x42,					// leaf 0x40 (2^6) with (+) key of 2 nibbles (0x02)
 			0xaa,					// key data
 			to_compact(1),			// length of value in bytes as Compact
 			0xbb					// value data
@@ -573,21 +591,20 @@ mod tests {
 	#[test]
 	fn codec_trie_two_tuples_disjoint_keys() {
 		let input = vec![(&[0x48, 0x19], &[0xfe]), (&[0x13, 0x14], &[0xff])];
-		let trie = unhashed_trie::<Blake2Hasher, _, _, _>(input);
+		let trie = Layout::trie_root_unhashed::<_, _, _>(input);
 		println!("trie: {:#x?}", trie);
-
 		let mut ex = Vec::<u8>::new();
-		ex.push(0xfe);									// branch, no value
+		ex.push(0x80);									// branch, no value (0b_10..) no nibble
 		ex.push(0x12);									// slots 1 & 4 are taken from 0-7
 		ex.push(0x00);									// no slots from 8-15
 		ex.push(to_compact(0x05));						// first slot: LEAF, 5 bytes long.
-		ex.push(0x04);									// leaf with 3 nibbles
+		ex.push(0x43);									// leaf 0x40 with 3 nibbles
 		ex.push(0x03);									// first nibble
 		ex.push(0x14);									// second & third nibble
 		ex.push(to_compact(0x01));						// 1 byte data
 		ex.push(0xff);									// value data
 		ex.push(to_compact(0x05));						// second slot: LEAF, 5 bytes long.
-		ex.push(0x04);									// leaf with 3 nibbles
+		ex.push(0x43);									// leaf with 3 nibbles
 		ex.push(0x08);									// first nibble
 		ex.push(0x19);									// second & third nibble
 		ex.push(to_compact(0x01));						// 1 byte data
@@ -605,9 +622,9 @@ mod tests {

 		let mut mdb = MemoryDB::default();
 		let mut root = Default::default();
-		let _ = populate_trie(&mut mdb, &mut root, &pairs);
+		let _ = populate_trie::<Layout>(&mut mdb, &mut root, &pairs);

-		let trie = TrieDB::<Blake2Hasher>::new(&mdb, &root).unwrap();
+		let trie = TrieDB::<Layout>::new(&mdb, &root).unwrap();

 		let iter = trie.iter().unwrap();
 		let mut iter_pairs = Vec::new();
@@ -18,72 +18,95 @@

 use rstd::marker::PhantomData;
 use rstd::vec::Vec;
+use rstd::borrow::Borrow;
 use codec::{Encode, Decode, Compact};
 use hash_db::Hasher;
-use trie_db::{self, DBValue, NibbleSlice, node::Node, ChildReference};
+use trie_db::{self, NibbleSlice, node::Node, ChildReference,
+	nibble_ops, Partial, NodeCodec as NodeCodecT};
 use crate::error::Error;
-use super::{EMPTY_TRIE, LEAF_NODE_OFFSET, LEAF_NODE_BIG, EXTENSION_NODE_OFFSET,
-	EXTENSION_NODE_BIG, take, partial_to_key, node_header::NodeHeader, branch_node};
+use crate::trie_constants;
+use super::{node_header::{NodeHeader, NodeKind}};
+
+fn take<'a>(input: &mut &'a[u8], count: usize) -> Option<&'a[u8]> {
+	if input.len() < count {
+		return None
+	}
+	let r = &(*input)[..count];
+	*input = &(*input)[count..];
+	Some(r)
+}

 /// Concrete implementation of a `NodeCodec` with Parity Codec encoding, generic over the `Hasher`
 #[derive(Default, Clone)]
-pub struct NodeCodec<H: Hasher>(PhantomData<H>);
+pub struct NodeCodec<H>(PhantomData<H>);

-impl<H: Hasher> trie_db::NodeCodec<H> for NodeCodec<H> {
+impl<H: Hasher> NodeCodecT<H> for NodeCodec<H> {
 	type Error = Error;

-	fn hashed_null_node() -> H::Out {
-		H::hash(&[0u8][..])
+	fn hashed_null_node() -> <H as Hasher>::Out {
+		H::hash(<Self as NodeCodecT<_>>::empty_node())
 	}

-	fn decode(data: &[u8]) -> ::rstd::result::Result<Node, Self::Error> {
-		use Error::BadFormat;
+	fn decode(data: &[u8]) -> rstd::result::Result<Node, Self::Error> {
 		let input = &mut &*data;
-		match NodeHeader::decode(input).ok_or(BadFormat)? {
+		let head = NodeHeader::decode(input).ok_or(Error::BadFormat)?;
+		match head {
 			NodeHeader::Null => Ok(Node::Empty),
-			NodeHeader::Branch(has_value) => {
-				let bitmap = u16::decode(input).ok_or(BadFormat)?;
+			NodeHeader::Branch(has_value, nibble_count) => {
+				let padding = nibble_count % nibble_ops::NIBBLE_PER_BYTE != 0;
+				// check that the padding is valid (if any)
+				if padding && nibble_ops::pad_left(input[0]) != 0 {
+					return Err(Error::BadFormat);
+				}
+				let nibble_data = take(
+					input,
+					(nibble_count + (nibble_ops::NIBBLE_PER_BYTE - 1)) / nibble_ops::NIBBLE_PER_BYTE,
+				).ok_or(Error::BadFormat)?;
+				let nibble_slice = NibbleSlice::new_offset(
+					nibble_data,
+					nibble_ops::number_padding(nibble_count),
+				);
+				let bitmap_slice = take(input, BITMAP_LENGTH).ok_or(Error::BadFormat)?;
+				let bitmap = Bitmap::decode(&bitmap_slice[..])?;
 				let value = if has_value {
-					let count = <Compact<u32>>::decode(input).ok_or(BadFormat)?.0 as usize;
-					Some(take(input, count).ok_or(BadFormat)?)
+					let count = <Compact<u32>>::decode(input).ok_or(Error::BadFormat)?.0 as usize;
+					Some(take(input, count).ok_or(Error::BadFormat)?)
 				} else {
 					None
 				};
 				let mut children = [None; 16];
-				let mut pot_cursor = 1;
-				for i in 0..16 {
-					if bitmap & pot_cursor != 0 {
-						let count = <Compact<u32>>::decode(input).ok_or(BadFormat)?.0 as usize;
-						children[i] = Some(take(input, count).ok_or(BadFormat)?);
+
+				for i in 0..nibble_ops::NIBBLE_LENGTH {
+					if bitmap.value_at(i) {
+						let count = <Compact<u32>>::decode(input).ok_or(Error::BadFormat)?.0 as usize;
+						children[i] = Some(take(input, count).ok_or(Error::BadFormat)?);
 					}
-					pot_cursor <<= 1;
 				}
-				Ok(Node::Branch(children, value))
-			}
-			NodeHeader::Extension(nibble_count) => {
-				if nibble_count % 2 == 1 && input[0] & 0xf0 != 0x00 {
-					return Err(BadFormat);
-				}
-				let nibble_data = take(input, (nibble_count + 1) / 2).ok_or(BadFormat)?;
-				let nibble_slice = NibbleSlice::new_offset(nibble_data, nibble_count % 2);
-				let count = <Compact<u32>>::decode(input).ok_or(BadFormat)?.0 as usize;
-				Ok(Node::Extension(nibble_slice, take(input, count).ok_or(BadFormat)?))
+				Ok(Node::NibbledBranch(nibble_slice, children, value))
 			}
 			NodeHeader::Leaf(nibble_count) => {
-				if nibble_count % 2 == 1 && input[0] & 0xf0 != 0x00 {
-					return Err(BadFormat);
+				let padding = nibble_count % nibble_ops::NIBBLE_PER_BYTE != 0;
+				// check that the padding is valid (if any)
+				if padding && nibble_ops::pad_left(input[0]) != 0 {
+					return Err(Error::BadFormat);
 				}
-				let nibble_data = take(input, (nibble_count + 1) / 2).ok_or(BadFormat)?;
-				let nibble_slice = NibbleSlice::new_offset(nibble_data, nibble_count % 2);
-				let count = <Compact<u32>>::decode(input).ok_or(BadFormat)?.0 as usize;
-				Ok(Node::Leaf(nibble_slice, take(input, count).ok_or(BadFormat)?))
+				let nibble_data = take(
+					input,
+					(nibble_count + (nibble_ops::NIBBLE_PER_BYTE - 1)) / nibble_ops::NIBBLE_PER_BYTE,
+				).ok_or(Error::BadFormat)?;
+				let nibble_slice = NibbleSlice::new_offset(
+					nibble_data,
+					nibble_ops::number_padding(nibble_count),
+				);
+				let count = <Compact<u32>>::decode(input).ok_or(Error::BadFormat)?.0 as usize;
+				Ok(Node::Leaf(nibble_slice, take(input, count).ok_or(Error::BadFormat)?))
 			}
 		}
 	}

-	fn try_decode_hash(data: &[u8]) -> Option<H::Out> {
+	fn try_decode_hash(data: &[u8]) -> Option<<H as Hasher>::Out> {
 		if data.len() == H::LENGTH {
-			let mut r = H::Out::default();
+			let mut r = <H as Hasher>::Out::default();
 			r.as_mut().copy_from_slice(data);
 			Some(r)
 		} else {
@@ -92,53 +115,140 @@ impl<H: Hasher> trie_db::NodeCodec<H> for NodeCodec<H> {
 	}

 	fn is_empty_node(data: &[u8]) -> bool {
-		data == &[EMPTY_TRIE][..]
-	}
-	fn empty_node() -> Vec<u8> {
-		[EMPTY_TRIE].to_vec()
+		data == <Self as NodeCodecT<_>>::empty_node()
 	}

-	// FIXME: refactor this so that `partial` isn't already encoded with HPE. Should just be an `impl Iterator<Item=u8>`.
-	fn leaf_node(partial: &[u8], value: &[u8]) -> Vec<u8> {
-		let mut output = partial_to_key(partial, LEAF_NODE_OFFSET, LEAF_NODE_BIG);
+	fn empty_node() -> &'static [u8] {
+		&[trie_constants::EMPTY_TRIE]
+	}
+
+	fn leaf_node(partial: Partial, value: &[u8]) -> Vec<u8> {
+		let mut output = partial_encode(partial, NodeKind::Leaf);
 		value.encode_to(&mut output);
 		output
 	}

-	// FIXME: refactor this so that `partial` isn't already encoded with HPE. Should just be an `impl Iterator<Item=u8>`.
-	fn ext_node(partial: &[u8], child: ChildReference<H::Out>) -> Vec<u8> {
-		let mut output = partial_to_key(partial, EXTENSION_NODE_OFFSET, EXTENSION_NODE_BIG);
-		match child {
-			ChildReference::Hash(h) => 
-				h.as_ref().encode_to(&mut output),
-			ChildReference::Inline(inline_data, len) =>
-				(&AsRef::<[u8]>::as_ref(&inline_data)[..len]).encode_to(&mut output),
-		};
-		output
+	fn extension_node(
+		_partial: impl Iterator<Item = u8>,
+		_nbnibble: usize,
+		_child: ChildReference<<H as Hasher>::Out>,
+	) -> Vec<u8> {
+		unreachable!()
 	}

-	fn branch_node<I>(children: I, maybe_value: Option<DBValue>) -> Vec<u8>
-		where I: IntoIterator<Item=Option<ChildReference<H::Out>>> + Iterator<Item=Option<ChildReference<H::Out>>>
-	{
-		let mut output = [0, 0, 0].to_vec();
-		let have_value = if let Some(value) = maybe_value {
-			(&*value).encode_to(&mut output);
-			true
+	fn branch_node(
+		_children: impl Iterator<Item = impl Borrow<Option<ChildReference<<H as Hasher>::Out>>>>,
+		_maybe_value: Option<&[u8]>,
+	) -> Vec<u8> {
+		unreachable!()
+	}
+
+	fn branch_node_nibbled(
+		partial: impl Iterator<Item = u8>,
+		number_nibble: usize,
+		children: impl Iterator<Item = impl Borrow<Option<ChildReference<<H as Hasher>::Out>>>>,
+		maybe_value: Option<&[u8]>,
+	) -> Vec<u8> {
+		let mut output = if maybe_value.is_some() {
+			partial_from_iterator_encode(partial, number_nibble, NodeKind::BranchWithValue)
 		} else {
-			false
+			partial_from_iterator_encode(partial, number_nibble, NodeKind::BranchNoValue)
 		};
-		let prefix = branch_node(have_value, children.map(|maybe_child| match maybe_child {
+		let bitmap_index = output.len();
+		let mut bitmap: [u8; BITMAP_LENGTH] = [0; BITMAP_LENGTH];
+		(0..BITMAP_LENGTH).for_each(|_|output.push(0));
+		if let Some(value) = maybe_value {
+			value.encode_to(&mut output);
+		};
+		Bitmap::encode(children.map(|maybe_child| match maybe_child.borrow() {
 			Some(ChildReference::Hash(h)) => {
 				h.as_ref().encode_to(&mut output);
 				true
 			}
-			Some(ChildReference::Inline(inline_data, len)) => {
-				(&AsRef::<[u8]>::as_ref(&inline_data)[..len]).encode_to(&mut output);
+			&Some(ChildReference::Inline(inline_data, len)) => {
+				inline_data.as_ref()[..len].encode_to(&mut output);
 				true
 			}
 			None => false,
-		}));
-		output[0..3].copy_from_slice(&prefix[..]);
+		}), bitmap.as_mut());
+		output[bitmap_index..bitmap_index + BITMAP_LENGTH]
+			.copy_from_slice(&bitmap.as_ref()[..BITMAP_LENGTH]);
 		output
 	}
+
 }
+
+// utils
+
+/// Encode and allocate node type header (type and size), and partial value.
+/// It uses an iterator over encoded partial bytes as input.
+fn partial_from_iterator_encode<I: Iterator<Item = u8>>(
+	partial: I,
+	nibble_count: usize,
+	node_kind: NodeKind,
+) -> Vec<u8> {
+	let nibble_count = rstd::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, nibble_count);
+
+	let mut output = Vec::with_capacity(3 + (nibble_count / nibble_ops::NIBBLE_PER_BYTE));
+	match node_kind {
+		NodeKind::Leaf => NodeHeader::Leaf(nibble_count).encode_to(&mut output),
+		NodeKind::BranchWithValue => NodeHeader::Branch(true, nibble_count).encode_to(&mut output),
+		NodeKind::BranchNoValue => NodeHeader::Branch(false, nibble_count).encode_to(&mut output),
+	};
+	output.extend(partial);
+	output
+}
+
+/// Encode and allocate node type header (type and size), and partial value.
+/// Same as `partial_from_iterator_encode` but uses non encoded `Partial` as input.
+fn partial_encode(partial: Partial, node_kind: NodeKind) -> Vec<u8> {
+	let number_nibble_encoded = (partial.0).0 as usize;
+	let nibble_count = partial.1.len() * nibble_ops::NIBBLE_PER_BYTE + number_nibble_encoded;
+
+	let nibble_count = rstd::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, nibble_count);
+
+	let mut output = Vec::with_capacity(3 + partial.1.len());
+	match node_kind {
+		NodeKind::Leaf => NodeHeader::Leaf(nibble_count).encode_to(&mut output),
+		NodeKind::BranchWithValue => NodeHeader::Branch(true, nibble_count).encode_to(&mut output),
+		NodeKind::BranchNoValue => NodeHeader::Branch(false, nibble_count).encode_to(&mut output),
+	};
+	if number_nibble_encoded > 0 {
+		output.push(nibble_ops::pad_right((partial.0).1));
+	}
+	output.extend_from_slice(&partial.1[..]);
+	output
+}
+
+const BITMAP_LENGTH: usize = 2;
+
+/// Radix 16 trie, bitmap encoding implementation,
+/// it contains children mapping information for a branch
+/// (children presence only), it encodes into
+/// a compact bitmap encoding representation.
+pub(crate) struct Bitmap(u16);
+
+impl Bitmap {
+
+	pub fn decode(data: &[u8]) -> Result<Self, Error> {
+		u16::decode(&mut &data[..])
+			.ok_or(Error::BadFormat)
+			.map(|v|Bitmap(v))
+	}
+
+	pub fn value_at(&self, i: usize) -> bool {
+		self.0 & (1u16 << i) != 0
+	}
+
+	pub fn encode<I: Iterator<Item = bool>>(has_children: I , dest: &mut [u8]) {
+		let mut bitmap: u16 = 0;
+		let mut cursor: u16 = 1;
+		for v in has_children {
+			if v { bitmap |= cursor }
+			cursor <<= 1;
+		}
+		dest[0] = (bitmap % 256) as u8;
+		dest[1] = (bitmap / 256) as u8;
+	}
+}
+
@@ -16,62 +16,105 @@

 //! The node header.

+use crate::trie_constants;
 use codec::{Encode, Decode, Input, Output};
-use super::{EMPTY_TRIE, LEAF_NODE_OFFSET, LEAF_NODE_BIG, EXTENSION_NODE_OFFSET,
-	EXTENSION_NODE_BIG, BRANCH_NODE_NO_VALUE, BRANCH_NODE_WITH_VALUE, LEAF_NODE_THRESHOLD,
-	EXTENSION_NODE_THRESHOLD, LEAF_NODE_SMALL_MAX, EXTENSION_NODE_SMALL_MAX};
+use rstd::iter::once;

-/// A node header.
-#[derive(Copy, Clone, PartialEq, Eq, Debug)]
-pub enum NodeHeader {
+/// A node header
+#[derive(Copy, Clone, PartialEq, Eq)]
+#[cfg_attr(feature = "std", derive(Debug))]
+pub(crate) enum NodeHeader {
 	Null,
-	Branch(bool),
-	Extension(usize),
+	Branch(bool, usize),
 	Leaf(usize),
 }

+/// NodeHeader without content
+pub(crate) enum NodeKind {
+	Leaf,
+	BranchNoValue,
+	BranchWithValue,
+}
+
 impl Encode for NodeHeader {
 	fn encode_to<T: Output>(&self, output: &mut T) {
 		match self {
-			NodeHeader::Null => output.push_byte(EMPTY_TRIE),
-
-			NodeHeader::Branch(true) => output.push_byte(BRANCH_NODE_WITH_VALUE),
-			NodeHeader::Branch(false) => output.push_byte(BRANCH_NODE_NO_VALUE),
-
-			NodeHeader::Leaf(nibble_count) if *nibble_count < LEAF_NODE_THRESHOLD as usize =>
-				output.push_byte(LEAF_NODE_OFFSET + *nibble_count as u8),
-			NodeHeader::Leaf(nibble_count) => {
-				output.push_byte(LEAF_NODE_BIG);
-				output.push_byte((*nibble_count - LEAF_NODE_THRESHOLD as usize) as u8);
-			}
-
-			NodeHeader::Extension(nibble_count) if *nibble_count < EXTENSION_NODE_THRESHOLD as usize =>
-				output.push_byte(EXTENSION_NODE_OFFSET + *nibble_count as u8),
-			NodeHeader::Extension(nibble_count) => {
-				output.push_byte(EXTENSION_NODE_BIG);
-				output.push_byte((*nibble_count - EXTENSION_NODE_THRESHOLD as usize) as u8);
-			}
+			NodeHeader::Null => output.push_byte(trie_constants::EMPTY_TRIE),
+			NodeHeader::Branch(true, nibble_count)	=>
+				encode_size_and_prefix(*nibble_count, trie_constants::BRANCH_WITH_MASK, output),
+			NodeHeader::Branch(false, nibble_count) =>
+				encode_size_and_prefix(*nibble_count, trie_constants::BRANCH_WITHOUT_MASK, output),
+			NodeHeader::Leaf(nibble_count) =>
+				encode_size_and_prefix(*nibble_count, trie_constants::LEAF_PREFIX_MASK, output),
 		}
 	}
 }

 impl Decode for NodeHeader {
 	fn decode<I: Input>(input: &mut I) -> Option<Self> {
-		Some(match input.read_byte()? {
-			EMPTY_TRIE => NodeHeader::Null,							// 0
-
-			i @ LEAF_NODE_OFFSET ..= LEAF_NODE_SMALL_MAX =>			// 1 ... (127 - 1)
-				NodeHeader::Leaf((i - LEAF_NODE_OFFSET) as usize),
-			LEAF_NODE_BIG =>										// 127
-				NodeHeader::Leaf(input.read_byte()? as usize + LEAF_NODE_THRESHOLD as usize),
-
-			i @ EXTENSION_NODE_OFFSET ..= EXTENSION_NODE_SMALL_MAX =>// 128 ... (253 - 1)
-				NodeHeader::Extension((i - EXTENSION_NODE_OFFSET) as usize),
-			EXTENSION_NODE_BIG =>									// 253
-				NodeHeader::Extension(input.read_byte()? as usize + EXTENSION_NODE_THRESHOLD as usize),
-
-			BRANCH_NODE_NO_VALUE => NodeHeader::Branch(false),		// 254
-			BRANCH_NODE_WITH_VALUE => NodeHeader::Branch(true),		// 255
-		})
+		let i = input.read_byte()?;
+		if i == trie_constants::EMPTY_TRIE {
+			return Some(NodeHeader::Null);
+		}
+		match i & (0b11 << 6) {
+			trie_constants::LEAF_PREFIX_MASK => Some(NodeHeader::Leaf(decode_size(i, input)?)),
+			trie_constants::BRANCH_WITHOUT_MASK => Some(NodeHeader::Branch(false, decode_size(i, input)?)),
+			trie_constants::BRANCH_WITH_MASK => Some(NodeHeader::Branch(true, decode_size(i, input)?)),
+			// do not allow any special encoding
+			_ => None,
+		}
 	}
 }
+
+/// Returns an iterator over encoded bytes for node header and size.
+/// Size encoding allows unlimited, length unefficient, representation, but
+/// is bounded to 16 bit maximum value to avoid possible DOS.
+pub(crate) fn size_and_prefix_iterator(size: usize, prefix: u8) -> impl Iterator<Item = u8> {
+	let size = rstd::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, size);
+
+	let l1 = rstd::cmp::min(62, size);
+	let (first_byte, mut rem) = if size == l1 {
+		(once(prefix + l1 as u8), 0)
+	} else {
+		(once(prefix + 63), size - l1)
+	};
+	let next_bytes = move || {
+		if rem > 0 {
+			if rem < 256 {
+				let result = rem - 1;
+				rem = 0;
+				Some(result as u8)
+			} else {
+				rem = rem.saturating_sub(255);
+				Some(255)
+			}
+		} else {
+			None
+		}
+	};
+	first_byte.chain(rstd::iter::from_fn(next_bytes))
+}
+
+/// Encodes size and prefix to a stream output.
+fn encode_size_and_prefix(size: usize, prefix: u8, out: &mut impl Output) {
+	for b in size_and_prefix_iterator(size, prefix) {
+		out.push_byte(b)
+	}
+}
+
+/// Decode size only from stream input and header byte.
+fn decode_size(first: u8, input: &mut impl Input) -> Option<usize> {
+	let mut result = (first & 255u8 >> 2) as usize;
+	if result < 63 {
+		return Some(result);
+	}
+	result -= 1;
+	while result <= trie_constants::NIBBLE_SIZE_BOUND {
+		let n = input.read_byte()? as usize;
+		if n < 255 {
+			return Some(result + n + 1);
+		}
+		result += 255;
+	}
+	Some(trie_constants::NIBBLE_SIZE_BOUND)
+}
@@ -16,16 +16,19 @@

 //! `TrieStream` implementation for Substrate's trie format.

-use rstd::iter::once;
 use hash_db::Hasher;
 use trie_root;
 use codec::Encode;
 use rstd::vec::Vec;
+use crate::trie_constants;
+use crate::node_header::{NodeKind, size_and_prefix_iterator};
+use crate::node_codec::Bitmap;

-use super::{EMPTY_TRIE, LEAF_NODE_OFFSET, LEAF_NODE_BIG, EXTENSION_NODE_OFFSET,
-	EXTENSION_NODE_BIG, branch_node};
+const BRANCH_NODE_NO_VALUE: u8 = 254;
+const BRANCH_NODE_WITH_VALUE: u8 = 255;

-/// Codec-flavored TrieStream
+#[derive(Default, Clone)]
+/// Codec-flavored TrieStream.
 pub struct TrieStream {
 	buffer: Vec<u8>,
 }
@@ -35,62 +38,102 @@ impl TrieStream {
 	pub fn as_raw(&self) -> &[u8] { &self.buffer }
 }

-/// Create a leaf/extension node, encoding a number of nibbles. Note that this
-/// cannot handle a number of nibbles that is zero or greater than 127 and if
-/// you attempt to do so *IT WILL PANIC*.
-fn fuse_nibbles_node<'a>(nibbles: &'a [u8], leaf: bool) -> impl Iterator<Item = u8> + 'a {
-	debug_assert!(nibbles.len() < 255 + 126, "nibbles length too long. what kind of size of key are you trying to include in the trie!?!");
-	// We use two ranges of possible values; one for leafs and the other for extensions.
-	// Each range encodes zero following nibbles up to some maximum. If the maximum is
-	// reached, then it is considered "big" and a second byte follows it in order to
-	// encode a further offset to the number of nibbles of up to 255. Beyond that, we
-	// cannot encode. This shouldn't be a problem though since that allows for keys of
-	// up to 380 nibbles (190 bytes) and we expect key sizes to be generally 128-bit (16
-	// bytes) or, at a push, 384-bit (48 bytes).
+fn branch_node_bit_mask(has_children: impl Iterator<Item = bool>) -> (u8, u8) {
+	let mut bitmap: u16 = 0;
+	let mut cursor: u16 = 1;
+	for v in has_children {
+		if v { bitmap |= cursor }
+		cursor <<= 1;
+	}
+	((bitmap % 256 ) as u8, (bitmap / 256 ) as u8)
+}

-	let (first_byte_small, big_threshold) = if leaf {
-		(LEAF_NODE_OFFSET, (LEAF_NODE_BIG - LEAF_NODE_OFFSET) as usize)
-	} else {
-		(EXTENSION_NODE_OFFSET, (EXTENSION_NODE_BIG - EXTENSION_NODE_OFFSET) as usize)
+
+/// Create a leaf/branch node, encoding a number of nibbles.
+fn fuse_nibbles_node<'a>(nibbles: &'a [u8], kind: NodeKind) -> impl Iterator<Item = u8> + 'a {
+	let size = rstd::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, nibbles.len());
+
+	let iter_start = match kind {
+		NodeKind::Leaf => size_and_prefix_iterator(size, trie_constants::LEAF_PREFIX_MASK),
+		NodeKind::BranchNoValue => size_and_prefix_iterator(size, trie_constants::BRANCH_WITHOUT_MASK),
+		NodeKind::BranchWithValue => size_and_prefix_iterator(size, trie_constants::BRANCH_WITH_MASK),
 	};
-	let first_byte = first_byte_small + nibbles.len().min(big_threshold) as u8;
-	once(first_byte)
-		.chain(if nibbles.len() >= big_threshold { Some((nibbles.len() - big_threshold) as u8) } else { None })
+	iter_start
 		.chain(if nibbles.len() % 2 == 1 { Some(nibbles[0]) } else { None })
 		.chain(nibbles[nibbles.len() % 2..].chunks(2).map(|ch| ch[0] << 4 | ch[1]))
 }

+
 impl trie_root::TrieStream for TrieStream {
-	fn new() -> Self { Self {buffer: Vec::new() } }
+
+	fn new() -> Self {
+		TrieStream {
+			buffer: Vec::new()
+		}
+	}
+
 	fn append_empty_data(&mut self) {
-		self.buffer.push(EMPTY_TRIE);
+		self.buffer.push(trie_constants::EMPTY_TRIE);
 	}

 	fn append_leaf(&mut self, key: &[u8], value: &[u8]) {
-		self.buffer.extend(fuse_nibbles_node(key, true));
+		self.buffer.extend(fuse_nibbles_node(key, NodeKind::Leaf));
 		value.encode_to(&mut self.buffer);
 	}
-	fn begin_branch(&mut self, maybe_value: Option<&[u8]>, has_children: impl Iterator<Item = bool>) {
-		self.buffer.extend(&branch_node(maybe_value.is_some(), has_children));
-		// Push the value if one exists.
+
+	fn begin_branch(
+		&mut self,
+		maybe_partial: Option<&[u8]>,
+		maybe_value: Option<&[u8]>,
+		has_children: impl Iterator<Item = bool>,
+	) {
+		if let Some(partial) = maybe_partial {
+			if maybe_value.is_some() {
+				self.buffer.extend(fuse_nibbles_node(partial, NodeKind::BranchWithValue));
+			} else {
+				self.buffer.extend(fuse_nibbles_node(partial, NodeKind::BranchNoValue));
+			}
+			let bm = branch_node_bit_mask(has_children);
+			self.buffer.extend([bm.0,bm.1].iter());
+		} else {
+			debug_assert!(false, "trie stream codec only for no extension trie");
+			self.buffer.extend(&branch_node(maybe_value.is_some(), has_children));
+		}
 		if let Some(value) = maybe_value {
 			value.encode_to(&mut self.buffer);
 		}
 	}
-	fn append_extension(&mut self, key: &[u8]) {
-		self.buffer.extend(fuse_nibbles_node(key, false));
+
+	fn append_extension(&mut self, _key: &[u8]) {
+		debug_assert!(false, "trie stream codec only for no extension trie");
 	}
+
 	fn append_substream<H: Hasher>(&mut self, other: Self) {
 		let data = other.out();
 		match data.len() {
-			0..=31 => {
-				data.encode_to(&mut self.buffer)
-			},
-			_ => {
-				H::hash(&data).as_ref().encode_to(&mut self.buffer)
-			}
+			0..=31 => data.encode_to(&mut self.buffer),
+			_ => H::hash(&data).as_ref().encode_to(&mut self.buffer),
 		}
 	}

 	fn out(self) -> Vec<u8> { self.buffer }
 }
+
+fn branch_node(has_value: bool, has_children: impl Iterator<Item = bool>) -> [u8; 3] {
+	let mut result = [0, 0, 0];
+	branch_node_buffered(has_value, has_children, &mut result[..]);
+	result
+}
+
+fn branch_node_buffered<I>(has_value: bool, has_children: I, output: &mut[u8]) 
+	where
+		I: Iterator<Item = bool>,
+{
+	let first = if has_value {
+		BRANCH_NODE_WITH_VALUE
+	} else {
+		BRANCH_NODE_NO_VALUE
+	};
+	output[0] = first;
+	Bitmap::encode(has_children, &mut output[1..]);
+}