use hash_db::HashDB;
use std::{convert::TryInto, marker::PhantomData};
use trie_db::{
nibble_ops::NIBBLE_LENGTH,
node::{Node, NodeHandle, Value},
CError, ChildReference, DBValue, NibbleVec, NodeCodec, TrieError, TrieHash, TrieLayout,
};
struct DecoderStackEntry<'a, C: NodeCodec> {
node: Node<'a>,
child_index: usize,
children: Vec<Option<ChildReference<C::HashOut>>>,
attached_value: Option<&'a [u8]>,
_marker: PhantomData<C>,
}
impl<'a, C: NodeCodec> DecoderStackEntry<'a, C> {
fn advance_child_index(&mut self) -> trie_db::Result<bool, C::HashOut, C::Error> {
match self.node {
Node::Extension(_, child) if self.child_index == 0 => {
match child {
NodeHandle::Inline(data) if data.is_empty() => return Ok(false),
_ => {
let child_ref = child.try_into().map_err(|hash| {
Box::new(TrieError::InvalidHash(C::HashOut::default(), hash))
})?;
self.children[self.child_index] = Some(child_ref);
},
}
self.child_index += 1;
},
Node::Branch(children, _) | Node::NibbledBranch(_, children, _) => {
while self.child_index < NIBBLE_LENGTH {
match children[self.child_index] {
Some(NodeHandle::Inline(data)) if data.is_empty() => return Ok(false),
Some(child) => {
let child_ref = child.try_into().map_err(|hash| {
Box::new(TrieError::InvalidHash(C::HashOut::default(), hash))
})?;
self.children[self.child_index] = Some(child_ref);
},
None => {},
}
self.child_index += 1;
}
},
_ => {},
}
Ok(true)
}
fn push_to_prefix(&self, prefix: &mut NibbleVec) {
match self.node {
Node::Empty => {},
Node::Leaf(partial, _) | Node::Extension(partial, _) => {
prefix.append_partial(partial.right());
},
Node::Branch(_, _) => {
prefix.push(self.child_index as u8);
},
Node::NibbledBranch(partial, _, _) => {
prefix.append_partial(partial.right());
prefix.push(self.child_index as u8);
},
}
}
fn pop_from_prefix(&self, prefix: &mut NibbleVec) {
match self.node {
Node::Empty => {},
Node::Leaf(partial, _) | Node::Extension(partial, _) => {
prefix.drop_lasts(partial.len());
},
Node::Branch(_, _) => {
prefix.pop();
},
Node::NibbledBranch(partial, _, _) => {
prefix.pop();
prefix.drop_lasts(partial.len());
},
}
}
fn encode_node(self, attached_hash: Option<&[u8]>) -> Vec<u8> {
let attached_hash = attached_hash.map(|h| Value::Node(h));
match self.node {
Node::Empty => C::empty_node().to_vec(),
Node::Leaf(partial, value) =>
C::leaf_node(partial.right_iter(), partial.len(), attached_hash.unwrap_or(value)),
Node::Extension(partial, _) => C::extension_node(
partial.right_iter(),
partial.len(),
self.children[0].expect("required by method precondition; qed"),
),
Node::Branch(_, value) => C::branch_node(
self.children.into_iter(),
if attached_hash.is_some() { attached_hash } else { value },
),
Node::NibbledBranch(partial, _, value) => C::branch_node_nibbled(
partial.right_iter(),
partial.len(),
self.children.iter(),
if attached_hash.is_some() { attached_hash } else { value },
),
}
}
}
pub fn decode_compact_from_iter<'a, L, DB, I>(
db: &mut DB,
encoded: I,
) -> trie_db::Result<(TrieHash<L>, usize), TrieHash<L>, CError<L>>
where
L: TrieLayout,
DB: HashDB<L::Hash, DBValue>,
I: IntoIterator<Item = &'a [u8]>,
{
let mut stack: Vec<DecoderStackEntry<L::Codec>> = Vec::new();
let mut prefix = NibbleVec::new();
let mut iter = encoded.into_iter().enumerate();
while let Some((i, encoded_node)) = iter.next() {
let mut attached_node = 0;
if let Some(header) = L::Codec::ESCAPE_HEADER {
if encoded_node.starts_with(&[header]) {
attached_node = 1;
}
}
let node = L::Codec::decode(&encoded_node[attached_node..])
.map_err(|err| Box::new(TrieError::DecoderError(<TrieHash<L>>::default(), err)))?;
let children_len = match node {
Node::Empty | Node::Leaf(..) => 0,
Node::Extension(..) => 1,
Node::Branch(..) | Node::NibbledBranch(..) => NIBBLE_LENGTH,
};
let mut last_entry = DecoderStackEntry {
node,
child_index: 0,
children: vec![None; children_len],
attached_value: None,
_marker: PhantomData::default(),
};
if attached_node > 0 {
if let Some((_, fetched_value)) = iter.next() {
last_entry.attached_value = Some(fetched_value);
} else {
return Err(Box::new(TrieError::IncompleteDatabase(<TrieHash<L>>::default())))
}
}
loop {
if !last_entry.advance_child_index()? {
last_entry.push_to_prefix(&mut prefix);
stack.push(last_entry);
break
}
let hash = last_entry
.attached_value
.as_ref()
.map(|value| db.insert(prefix.as_prefix(), value));
let node_data = last_entry.encode_node(hash.as_ref().map(|h| h.as_ref()));
let node_hash = db.insert(prefix.as_prefix(), node_data.as_ref());
if let Some(entry) = stack.pop() {
last_entry = entry;
last_entry.pop_from_prefix(&mut prefix);
last_entry.children[last_entry.child_index] = Some(ChildReference::Hash(node_hash));
last_entry.child_index += 1;
} else {
return Ok((node_hash, i + 1))
}
}
}
Err(Box::new(TrieError::IncompleteDatabase(<TrieHash<L>>::default())))
}