diff --git a/test-support/reference-trie/src/lib.rs b/test-support/reference-trie/src/lib.rs index 3c546e47..43c94e64 100644 --- a/test-support/reference-trie/src/lib.rs +++ b/test-support/reference-trie/src/lib.rs @@ -21,7 +21,7 @@ use std::ops::Range; use parity_scale_codec::{Decode, Input, Output, Encode, Compact, Error as CodecError}; use trie_root::Hasher; use trie_db::{ - node::{NibbleSlicePlan, NodePlan}, + node::{NibbleSlicePlan, NodePlan, NodeHandlePlan}, triedbmut::ChildReference, DBValue, trie_visit, @@ -47,8 +47,8 @@ pub struct ExtensionLayout; impl TrieLayout for ExtensionLayout { const USE_EXTENSION: bool = true; - type Hash = keccak_hasher::KeccakHasher; - type Codec = ReferenceNodeCodec; + type Hash = KeccakHasher; + type Codec = ReferenceNodeCodec; } impl TrieConfiguration for ExtensionLayout { } @@ -60,7 +60,7 @@ pub struct GenericNoExtensionLayout(PhantomData); impl TrieLayout for GenericNoExtensionLayout { const USE_EXTENSION: bool = false; type Hash = H; - type Codec = ReferenceNodeCodecNoExt; + type Codec = ReferenceNodeCodecNoExt; } impl TrieConfiguration for GenericNoExtensionLayout { } @@ -481,14 +481,14 @@ impl Decode for NodeHeaderNoExt { /// Simple reference implementation of a `NodeCodec`. #[derive(Default, Clone)] -pub struct ReferenceNodeCodec; +pub struct ReferenceNodeCodec(PhantomData); /// Simple reference implementation of a `NodeCodec`. /// Even if implementation follows initial specification of /// https://github.com/w3f/polkadot-re-spec/issues/8, this may /// not follow it in the future, it is mainly the testing codec without extension node. #[derive(Default, Clone)] -pub struct ReferenceNodeCodecNoExt; +pub struct ReferenceNodeCodecNoExt(PhantomData); fn partial_to_key(partial: Partial, offset: u8, over: u8) -> Vec { let number_nibble_encoded = (partial.0).0 as usize; @@ -613,11 +613,12 @@ impl<'a> Input for ByteSliceInput<'a> { // but due to the current limitations of Rust const evaluation we can't do // `const HASHED_NULL_NODE: ::Out = ::Out( … … )`. // Perhaps one day soon? -impl NodeCodec for ReferenceNodeCodec { +impl NodeCodec for ReferenceNodeCodec { type Error = CodecError; + type HashOut = H::Out; fn hashed_null_node() -> ::Out { - H::hash(>::empty_node()) + H::hash(::empty_node()) } fn decode_plan(data: &[u8]) -> ::std::result::Result { @@ -641,7 +642,12 @@ impl NodeCodec for ReferenceNodeCodec { for i in 0..nibble_ops::NIBBLE_LENGTH { if bitmap.value_at(i) { let count = >::decode(&mut input)?.0 as usize; - children[i] = Some(input.take(count)?); + let range = input.take(count)?; + children[i] = Some(if count == H::LENGTH { + NodeHandlePlan::Hash(range) + } else { + NodeHandlePlan::Inline(range) + }); } } Ok(NodePlan::Branch { value, children }) @@ -652,7 +658,12 @@ impl NodeCodec for ReferenceNodeCodec { )?; let partial_padding = nibble_ops::number_padding(nibble_count); let count = >::decode(&mut input)?.0 as usize; - let child = input.take(count)?; + let range = input.take(count)?; + let child = if count == H::LENGTH { + NodeHandlePlan::Hash(range) + } else { + NodeHandlePlan::Inline(range) + }; Ok(NodePlan::Extension { partial: NibbleSlicePlan::new(partial, partial_padding), child @@ -673,18 +684,8 @@ impl NodeCodec for ReferenceNodeCodec { } } - fn try_decode_hash(data: &[u8]) -> Option<::Out> { - if data.len() == H::LENGTH { - let mut r = ::Out::default(); - r.as_mut().copy_from_slice(data); - Some(r) - } else { - None - } - } - fn is_empty_node(data: &[u8]) -> bool { - data == >::empty_node() + data == ::empty_node() } fn empty_node() -> &'static[u8] { @@ -700,7 +701,7 @@ impl NodeCodec for ReferenceNodeCodec { fn extension_node( partial: impl Iterator, number_nibble: usize, - child: ChildReference<::Out>, + child: ChildReference, ) -> Vec { let mut output = partial_from_iterator_to_key( partial, @@ -717,7 +718,7 @@ impl NodeCodec for ReferenceNodeCodec { } fn branch_node( - children: impl Iterator::Out>>>>, + children: impl Iterator>>>, maybe_value: Option<&[u8]>, ) -> Vec { let mut output = vec![0; BITMAP_LENGTH + 1]; @@ -747,18 +748,19 @@ impl NodeCodec for ReferenceNodeCodec { fn branch_node_nibbled( _partial: impl Iterator, _number_nibble: usize, - _children: impl Iterator::Out>>>>, + _children: impl Iterator>>>, _maybe_value: Option<&[u8]>) -> Vec { unreachable!() } } -impl NodeCodec for ReferenceNodeCodecNoExt { +impl NodeCodec for ReferenceNodeCodecNoExt { type Error = CodecError; + type HashOut = ::Out; fn hashed_null_node() -> ::Out { - H::hash(>::empty_node()) + H::hash(::empty_node()) } fn decode_plan(data: &[u8]) -> ::std::result::Result { @@ -790,7 +792,12 @@ impl NodeCodec for ReferenceNodeCodecNoExt { for i in 0..nibble_ops::NIBBLE_LENGTH { if bitmap.value_at(i) { let count = >::decode(&mut input)?.0 as usize; - children[i] = Some(input.take(count)?); + let range = input.take(count)?; + children[i] = Some(if count == H::LENGTH { + NodeHandlePlan::Hash(range) + } else { + NodeHandlePlan::Inline(range) + }); } } Ok(NodePlan::NibbledBranch { @@ -819,12 +826,8 @@ impl NodeCodec for ReferenceNodeCodecNoExt { } } - fn try_decode_hash(data: &[u8]) -> Option<::Out> { - >::try_decode_hash(data) - } - fn is_empty_node(data: &[u8]) -> bool { - data == >::empty_node() + data == ::empty_node() } fn empty_node() -> &'static [u8] { @@ -855,7 +858,7 @@ impl NodeCodec for ReferenceNodeCodecNoExt { fn branch_node_nibbled( partial: impl Iterator, number_nibble: usize, - children: impl Iterator::Out>>>>, + children: impl Iterator>>>, maybe_value: Option<&[u8]>, ) -> Vec { let mut output = if maybe_value.is_some() { @@ -1211,9 +1214,9 @@ mod tests { fn too_big_nibble_length() { // + 1 for 0 added byte of nibble encode let input = vec![0u8; (NIBBLE_SIZE_BOUND_NO_EXT as usize + 1) / 2 + 1]; - let enc = > + let enc = as NodeCodec> ::leaf_node(((0, 0), &input), &[1]); - let dec = > + let dec = as NodeCodec> ::decode(&enc).unwrap(); let o_sl = if let Node::Leaf(sl, _) = dec { Some(sl) diff --git a/trie-db/src/iterator.rs b/trie-db/src/iterator.rs index d9af230f..0d8232a4 100644 --- a/trie-db/src/iterator.rs +++ b/trie-db/src/iterator.rs @@ -15,7 +15,7 @@ use super::{CError, DBValue, Result, Trie, TrieHash, TrieIterator, TrieLayout}; use hash_db::{Hasher, EMPTY_PREFIX}; use triedb::TrieDB; -use node::{NodePlan, OwnedNode}; +use node::{NodePlan, NodeHandle, OwnedNode}; use nibble::{NibbleSlice, NibbleVec, nibble_ops}; #[cfg(feature = "std")] @@ -76,7 +76,11 @@ impl<'a, L: TrieLayout> TrieDBNodeIterator<'a, L> { trail: Vec::with_capacity(8), key_nibbles: NibbleVec::new(), }; - let (root_node, root_hash) = db.get_raw_or_lookup(db.root().as_ref(), EMPTY_PREFIX)?; + let (root_node, root_hash) = db.get_raw_or_lookup( + *db.root(), + NodeHandle::Hash(db.root().as_ref()), + EMPTY_PREFIX + )?; r.descend(root_node, root_hash); Ok(r) } @@ -97,8 +101,11 @@ impl<'a, L: TrieLayout> TrieIterator for TrieDBNodeIterator<'a, L> { self.key_nibbles.clear(); let key = NibbleSlice::new(key); - let (mut node, mut node_hash) = - self.db.get_raw_or_lookup(self.db.root().as_ref(), EMPTY_PREFIX)?; + let (mut node, mut node_hash) = self.db.get_raw_or_lookup( + >::default(), + NodeHandle::Hash(self.db.root().as_ref()), + EMPTY_PREFIX + )?; let mut partial = key; let mut full_key_nibbles = 0; loop { @@ -135,7 +142,11 @@ impl<'a, L: TrieLayout> TrieIterator for TrieDBNodeIterator<'a, L> { self.key_nibbles.append_partial(slice.right()); let prefix = key.back(full_key_nibbles); - self.db.get_raw_or_lookup(&node_data[child.clone()], prefix.left())? + self.db.get_raw_or_lookup( + node_hash.unwrap_or_default(), + child.build(node_data), + prefix.left() + )? }, NodePlan::Branch { value: _, children } => { if partial.is_empty() { @@ -151,7 +162,11 @@ impl<'a, L: TrieLayout> TrieIterator for TrieDBNodeIterator<'a, L> { partial = partial.mid(1); let prefix = key.back(full_key_nibbles); - self.db.get_raw_or_lookup(&node_data[child.clone()], prefix.left())? + self.db.get_raw_or_lookup( + node_hash.unwrap_or_default(), + child.build(node_data), + prefix.left() + )? } else { return Ok(()) } @@ -184,7 +199,11 @@ impl<'a, L: TrieLayout> TrieIterator for TrieDBNodeIterator<'a, L> { partial = partial.mid(1); let prefix = key.back(full_key_nibbles); - self.db.get_raw_or_lookup(&node_data[child.clone()], prefix.left())? + self.db.get_raw_or_lookup( + node_hash.unwrap_or_default(), + child.build(node_data), + prefix.left() + )? } else { return Ok(()) } @@ -239,7 +258,8 @@ impl<'a, L: TrieLayout> Iterator for TrieDBNodeIterator<'a, L> { self.key_nibbles.append_partial(partial.right()); IterStep::Descend::, CError>( self.db.get_raw_or_lookup( - &node_data[child.clone()], + b.hash.unwrap_or_default(), + child.build(node_data), self.key_nibbles.as_prefix() ) ) @@ -261,7 +281,8 @@ impl<'a, L: TrieLayout> Iterator for TrieDBNodeIterator<'a, L> { self.key_nibbles.push(i as u8); IterStep::Descend::, CError>( self.db.get_raw_or_lookup( - &node_data[child.clone()], + b.hash.unwrap_or_default(), + child.build(node_data), self.key_nibbles.as_prefix() ) ) diff --git a/trie-db/src/lib.rs b/trie-db/src/lib.rs index 2b8333bd..622fcedc 100644 --- a/trie-db/src/lib.rs +++ b/trie-db/src/lib.rs @@ -123,6 +123,7 @@ pub enum TrieError { ValueAtIncompleteKey(Vec, u8), /// Corrupt Trie item DecoderError(T, E), + InvalidHash(T, Vec), } #[cfg(feature = "std")] @@ -138,6 +139,12 @@ impl fmt::Display for TrieError where T: MaybeDebug, E: MaybeDebug { TrieError::DecoderError(ref hash, ref decoder_err) => { write!(f, "Decoding failed for hash {:?}; err: {:?}", hash, decoder_err) } + TrieError::InvalidHash(ref hash, ref data) => + write!( + f, + "Encoded node {:?} contains invalid hash reference with length: {}", + hash, data.len() + ), } } } @@ -150,6 +157,7 @@ impl Error for TrieError where T: fmt::Debug, E: Error { TrieError::IncompleteDatabase(_) => "Incomplete database", TrieError::ValueAtIncompleteKey(_, _) => "Value at incomplete key", TrieError::DecoderError(_, ref err) => err.description(), + TrieError::InvalidHash(_, _) => "Encoded node contains invalid hash reference", } } } @@ -415,7 +423,7 @@ pub trait TrieLayout { /// Hasher to use for this trie. type Hash: Hasher; /// Codec to use (needs to match hasher and nibble ops). - type Codec: NodeCodec; + type Codec: NodeCodec::Out>; } /// This traits associates a trie definition with prefered methods. @@ -477,6 +485,4 @@ pub trait TrieConfiguration: Sized + TrieLayout { /// Alias accessor to hasher hash output type from a `TrieLayout`. pub type TrieHash = <::Hash as Hasher>::Out; /// Alias accessor to `NodeCodec` associated `Error` type from a `TrieLayout`. -pub type CError = < - ::Codec as NodeCodec<::Hash> ->::Error; +pub type CError = <::Codec as NodeCodec>::Error; diff --git a/trie-db/src/lookup.rs b/trie-db/src/lookup.rs index 5eb0ec06..62fd0bde 100644 --- a/trie-db/src/lookup.rs +++ b/trie-db/src/lookup.rs @@ -16,7 +16,7 @@ use hash_db::HashDBRef; use nibble::NibbleSlice; -use node::Node; +use node::{Node, NodeHandle, decode_hash}; use node_codec::NodeCodec; use super::{DBValue, Result, TrieError, Query, TrieLayout, CError, TrieHash}; @@ -70,7 +70,7 @@ where return Err(Box::new(TrieError::DecoderError(hash, e))) } }; - match decoded { + let next_node = match decoded { Node::Leaf(slice, value) => { return Ok(match slice == partial { true => Some(self.query.decode(value)), @@ -79,9 +79,9 @@ where } Node::Extension(slice, item) => { if partial.starts_with(&slice) { - node_data = item; partial = partial.mid(slice.len()); key_nibbles += slice.len(); + item } else { return Ok(None) } @@ -90,9 +90,9 @@ where true => return Ok(value.map(move |val| self.query.decode(val))), false => match children[partial.at(0) as usize] { Some(x) => { - node_data = x; partial = partial.mid(1); key_nibbles += 1; + x } None => return Ok(None) } @@ -106,21 +106,27 @@ where true => return Ok(value.map(move |val| self.query.decode(val))), false => match children[partial.at(slice.len()) as usize] { Some(x) => { - node_data = x; partial = partial.mid(slice.len() + 1); key_nibbles += slice.len() + 1; + x } None => return Ok(None) } } }, Node::Empty => return Ok(None), - } + }; // check if new node data is inline or hash. - if let Some(h) = L::Codec::try_decode_hash(&node_data) { - hash = h; - break + match next_node { + NodeHandle::Hash(data) => { + hash = decode_hash::(data) + .ok_or_else(|| Box::new(TrieError::InvalidHash(hash, data.to_vec())))?; + break; + }, + NodeHandle::Inline(data) => { + node_data = data; + }, } } } diff --git a/trie-db/src/node.rs b/trie-db/src/node.rs index f126b4cc..0a318e84 100644 --- a/trie-db/src/node.rs +++ b/trie-db/src/node.rs @@ -27,6 +27,23 @@ use alloc::vec::Vec; /// Offset is applied on first byte of array (bytes are right aligned). pub type NodeKey = (usize, ElasticArray36); +/// A reference to a trie node which may be stored within another trie node. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum NodeHandle<'a> { + Hash(&'a [u8]), + Inline(&'a [u8]), +} + +/// Read a hash from a slice into a Hasher output. Returns None if the slice is the wrong length. +pub fn decode_hash(data: &[u8]) -> Option { + if data.len() != H::LENGTH { + return None; + } + let mut hash = H::Out::default(); + hash.as_mut().copy_from_slice(data); + Some(hash) +} + /// Type of node in the trie and essential information thereof. #[derive(Eq, PartialEq, Clone)] #[cfg_attr(feature = "std", derive(Debug))] @@ -36,12 +53,32 @@ pub enum Node<'a> { /// Leaf node; has key slice and value. Value may not be empty. Leaf(NibbleSlice<'a>, &'a [u8]), /// Extension node; has key slice and node data. Data may not be null. - Extension(NibbleSlice<'a>, &'a [u8]), + Extension(NibbleSlice<'a>, NodeHandle<'a>), /// Branch node; has slice of child nodes (each possibly null) /// and an optional immediate node data. - Branch([Option<&'a [u8]>; nibble_ops::NIBBLE_LENGTH], Option<&'a [u8]>), + Branch([Option>; nibble_ops::NIBBLE_LENGTH], Option<&'a [u8]>), /// Branch node with support for a nibble (when extension nodes are not used). - NibbledBranch(NibbleSlice<'a>, [Option<&'a [u8]>; nibble_ops::NIBBLE_LENGTH], Option<&'a [u8]>), + NibbledBranch(NibbleSlice<'a>, [Option>; nibble_ops::NIBBLE_LENGTH], Option<&'a [u8]>), +} + +/// A `NodeHandlePlan` is a decoding plan for constructing a `NodeHandle` from an encoded trie +/// node. This is used as a substructure of `NodePlan`. See `NodePlan` for details. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum NodeHandlePlan { + Hash(Range), + Inline(Range), +} + +impl NodeHandlePlan { + /// Build a node handle by decoding a byte slice according to the node handle plan. It is the + /// responsibility of the caller to ensure that the node plan was created for the argument + /// data, otherwise the call may decode incorrectly or panic. + pub fn build<'a, 'b>(&'a self, data: &'b [u8]) -> NodeHandle<'b> { + match self { + NodeHandlePlan::Hash(range) => NodeHandle::Hash(&data[range.clone()]), + NodeHandlePlan::Inline(range) => NodeHandle::Inline(&data[range.clone()]), + } + } } /// A `NibbleSlicePlan` is a blueprint for decoding a nibble slice from a byte slice. The @@ -69,7 +106,7 @@ impl NibbleSlicePlan { /// Build a nibble slice by decoding a byte slice according to the plan. It is the /// responsibility of the caller to ensure that the node plan was created for the argument - /// data, otherwise the call decode incorrectly or panic. + /// data, otherwise the call may decode incorrectly or panic. pub fn build<'a, 'b>(&'a self, data: &'b [u8]) -> NibbleSlice<'b> { NibbleSlice::new_offset(&data[self.bytes.clone()], self.offset) } @@ -94,37 +131,37 @@ pub enum NodePlan { /// Extension node; has a partial key plan and child data. Extension { partial: NibbleSlicePlan, - child: Range, + child: NodeHandlePlan, }, /// Branch node; has slice of child nodes (each possibly null) /// and an optional immediate node data. Branch { value: Option>, - children: [Option>; nibble_ops::NIBBLE_LENGTH], + children: [Option; nibble_ops::NIBBLE_LENGTH], }, /// Branch node with support for a nibble (when extension nodes are not used). NibbledBranch { partial: NibbleSlicePlan, value: Option>, - children: [Option>; nibble_ops::NIBBLE_LENGTH], + children: [Option; nibble_ops::NIBBLE_LENGTH], }, } impl NodePlan { /// Build a node by decoding a byte slice according to the node plan. It is the responsibility /// of the caller to ensure that the node plan was created for the argument data, otherwise the - /// call decode incorrectly or panic. + /// call may decode incorrectly or panic. pub fn build<'a, 'b>(&'a self, data: &'b [u8]) -> Node<'b> { match self { NodePlan::Empty => Node::Empty, NodePlan::Leaf { partial, value } => Node::Leaf(partial.build(data), &data[value.clone()]), NodePlan::Extension { partial, child } => - Node::Extension(partial.build(data), &data[child.clone()]), + Node::Extension(partial.build(data), child.build(data)), NodePlan::Branch { value, children } => { let mut child_slices = [None; nibble_ops::NIBBLE_LENGTH]; for i in 0..nibble_ops::NIBBLE_LENGTH { - child_slices[i] = children[i].clone().map(|child| &data[child]); + child_slices[i] = children[i].as_ref().map(|child| child.build(data)); } let value_slice = value.clone().map(|value| &data[value]); Node::Branch(child_slices, value_slice) @@ -132,7 +169,7 @@ impl NodePlan { NodePlan::NibbledBranch { partial, value, children } => { let mut child_slices = [None; nibble_ops::NIBBLE_LENGTH]; for i in 0..nibble_ops::NIBBLE_LENGTH { - child_slices[i] = children[i].clone().map(|child| &data[child]); + child_slices[i] = children[i].as_ref().map(|child| child.build(data)); } let value_slice = value.clone().map(|value| &data[value]); Node::NibbledBranch(partial.build(data), child_slices, value_slice) @@ -152,7 +189,7 @@ pub struct OwnedNode> { impl> OwnedNode { /// Construct an `OwnedNode` by decoding an owned data source according to some codec. - pub fn new>(data: D) -> Result { + pub fn new(data: D) -> Result { let plan = C::decode_plan(data.borrow())?; Ok(OwnedNode { data, plan }) } diff --git a/trie-db/src/node_codec.rs b/trie-db/src/node_codec.rs index 1bab5f55..4d837b37 100644 --- a/trie-db/src/node_codec.rs +++ b/trie-db/src/node_codec.rs @@ -15,14 +15,17 @@ //! Generic trait for trie node encoding/decoding. Takes a `hash_db::Hasher` //! to parametrize the hashes used in the codec. -use hash_db::Hasher; +use hash_db::MaybeDebug; use node::{Node, NodePlan}; use ChildReference; #[cfg(feature = "std")] use std::borrow::Borrow; - #[cfg(not(feature = "std"))] use core::borrow::Borrow; +#[cfg(feature = "std")] +use std::hash; +#[cfg(not(feature = "std"))] +use core::hash; #[cfg(feature = "std")] use std::error::Error; @@ -44,12 +47,16 @@ impl Error for T {} pub type Partial<'a> = ((u8, u8), &'a[u8]); /// Trait for trie node encoding/decoding. -pub trait NodeCodec: Sized { +pub trait NodeCodec: Sized { /// Codec error type. type Error: Error; + /// Output type of encoded node hasher. + type HashOut: AsRef<[u8]> + AsMut<[u8]> + Default + MaybeDebug + PartialEq + Eq + + hash::Hash + Send + Sync + Clone + Copy; + /// Get the hashed null node. - fn hashed_null_node() -> H::Out; + fn hashed_null_node() -> Self::HashOut; /// Decode bytes to a `NodePlan`. Returns `Self::E` on failure. fn decode_plan(data: &[u8]) -> Result; @@ -59,9 +66,6 @@ pub trait NodeCodec: Sized { Ok(Self::decode_plan(data)?.build(data)) } - /// Decode bytes to the `Hasher`s output type. Returns `None` on failure. - fn try_decode_hash(data: &[u8]) -> Option; - /// Check if the provided bytes correspond to the codecs "empty" node. fn is_empty_node(data: &[u8]) -> bool; @@ -78,13 +82,13 @@ pub trait NodeCodec: Sized { fn extension_node( partial: impl Iterator, number_nibble: usize, - child_ref: ChildReference, + child_ref: ChildReference, ) -> Vec; /// Returns an encoded branch node. - /// Takes an iterator yielding `ChildReference` and an optional value. + /// Takes an iterator yielding `ChildReference` and an optional value. fn branch_node( - children: impl Iterator>>>, + children: impl Iterator>>>, value: Option<&[u8]>, ) -> Vec; @@ -93,8 +97,7 @@ pub trait NodeCodec: Sized { fn branch_node_nibbled( partial: impl Iterator, number_nibble: usize, - children: impl Iterator>>>, + children: impl Iterator>>>, value: Option<&[u8]> ) -> Vec; } - diff --git a/trie-db/src/triedb.rs b/trie-db/src/triedb.rs index e17bb470..679aad72 100644 --- a/trie-db/src/triedb.rs +++ b/trie-db/src/triedb.rs @@ -15,8 +15,7 @@ use hash_db::{HashDBRef, Prefix, EMPTY_PREFIX}; use nibble::NibbleSlice; use iterator::TrieDBNodeIterator; -use super::node::{Node, OwnedNode}; -use node_codec::NodeCodec; +use super::node::{NodeHandle, Node, OwnedNode, decode_hash}; use super::lookup::Lookup; use super::{Result, DBValue, Trie, TrieItem, TrieError, TrieIterator, Query, TrieLayout, CError, TrieHash}; @@ -99,25 +98,31 @@ where /// /// `partial_key` is encoded nibble slice that addresses the node. pub(crate) fn get_raw_or_lookup( - &self, node: &[u8], + &self, + parent_hash: TrieHash, + node_handle: NodeHandle, partial_key: Prefix, ) -> Result<(OwnedNode, Option>), TrieHash, CError> { - let node_hash = L::Codec::try_decode_hash(node); - let node_data = if let Some(key) = node_hash { - self.db - .get(&key, partial_key) - .ok_or_else(|| { - if partial_key == EMPTY_PREFIX { - Box::new(TrieError::InvalidStateRoot(key)) - } else { - Box::new(TrieError::IncompleteDatabase(key)) - } - })? - } else { - DBValue::from_slice(node) + let (node_hash, node_data) = match node_handle { + NodeHandle::Hash(data) => { + let node_hash = decode_hash::(data) + .ok_or_else(|| Box::new(TrieError::InvalidHash(parent_hash, data.to_vec())))?; + let node_data = self.db + .get(&node_hash, partial_key) + .ok_or_else(|| { + if partial_key == EMPTY_PREFIX { + Box::new(TrieError::InvalidStateRoot(node_hash)) + } else { + Box::new(TrieError::IncompleteDatabase(node_hash)) + } + })?; + + (Some(node_hash), node_data) + } + NodeHandle::Inline(data) => (None, DBValue::from_slice(data)), }; - let owned_node = OwnedNode::new::(node_data) - .map_err(|e| Box::new(TrieError::DecoderError(node_hash.unwrap_or_default(), e)))?; + let owned_node = OwnedNode::new::(node_data) + .map_err(|e| Box::new(TrieError::DecoderError(node_hash.unwrap_or(parent_hash), e)))?; Ok((owned_node, node_hash)) } } @@ -159,7 +164,7 @@ where L: TrieLayout, { trie: &'db TrieDB<'db, L>, - node_key: &'a[u8], + node_key: NodeHandle<'a>, partial_key: NibbleVec, index: Option, } @@ -170,7 +175,11 @@ where L: TrieLayout, { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { - match self.trie.get_raw_or_lookup(self.node_key, self.partial_key.as_prefix()) { + match self.trie.get_raw_or_lookup( + >::default(), + self.node_key, + self.partial_key.as_prefix() + ) { Ok((owned_node, _node_hash)) => match owned_node.node() { Node::Leaf(slice, value) => match (f.debug_struct("Node::Leaf"), self.index) { @@ -180,7 +189,7 @@ where .field("slice", &slice) .field("value", &value) .finish(), - Node::Extension(slice, item) => + Node::Extension(slice, item) => { match (f.debug_struct("Node::Extension"), self.index) { (ref mut d, Some(i)) => d.field("index", &i), (ref mut d, _) => d, @@ -193,7 +202,8 @@ where .clone_append_optional_slice_and_nibble(Some(&slice), None), index: None, }) - .finish(), + .finish() + }, Node::Branch(ref nodes, ref value) => { let nodes: Vec> = nodes.into_iter() .enumerate() @@ -255,7 +265,7 @@ where .field("hash_count", &self.hash_count) .field("root", &TrieAwareDebugNode { trie: self, - node_key: self.root().as_ref(), + node_key: NodeHandle::Hash(self.root().as_ref()), partial_key: NibbleVec::new(), index: None, }) diff --git a/trie-db/src/triedbmut.rs b/trie-db/src/triedbmut.rs index 592c99cc..fc8f47ee 100644 --- a/trie-db/src/triedbmut.rs +++ b/trie-db/src/triedbmut.rs @@ -16,7 +16,7 @@ use super::{Result, TrieError, TrieMut, TrieLayout, TrieHash, CError}; use super::lookup::Lookup; -use super::node::Node as EncodedNode; +use super::node::{NodeHandle as EncodedNodeHandle, Node as EncodedNode, decode_hash}; use node_codec::NodeCodec; use super::{DBValue, node::NodeKey}; @@ -100,78 +100,90 @@ where { // load an inline node into memory or get the hash to do the lookup later. fn inline_or_hash( - node: &[u8], + parent_hash: H::Out, + child: EncodedNodeHandle, db: &dyn HashDB, storage: &mut NodeStorage - ) -> NodeHandle + ) -> Result, H::Out, C::Error> where - C: NodeCodec, - H: Hasher, + C: NodeCodec, + H: Hasher, { - C::try_decode_hash(&node) - .map(NodeHandle::Hash) - .unwrap_or_else(|| { - let child = Node::from_encoded::(node, db, storage); + let handle = match child { + EncodedNodeHandle::Hash(data) => { + let hash = decode_hash::(data) + .ok_or_else(|| Box::new(TrieError::InvalidHash(parent_hash, data.to_vec())))?; + NodeHandle::Hash(hash) + }, + EncodedNodeHandle::Inline(data) => { + let child = Node::from_encoded::(parent_hash, data, db, storage)?; NodeHandle::InMemory(storage.alloc(Stored::New(child))) - }) + }, + }; + Ok(handle) } // Decode a node from encoded bytes. fn from_encoded<'a, 'b, C, H>( + node_hash: H::Out, data: &'a[u8], db: &dyn HashDB, storage: &'b mut NodeStorage, - ) -> Self + ) -> Result where - C: NodeCodec, H: Hasher, + C: NodeCodec, H: Hasher, { - match C::decode(data).unwrap_or(EncodedNode::Empty) { + let encoded_node = C::decode(data) + .map_err(|e| Box::new(TrieError::DecoderError(node_hash, e)))?; + let node = match encoded_node { EncodedNode::Empty => Node::Empty, EncodedNode::Leaf(k, v) => Node::Leaf(k.into(), DBValue::from_slice(&v)), EncodedNode::Extension(key, cb) => { Node::Extension( key.into(), - Self::inline_or_hash::(cb, db, storage)) - }, + Self::inline_or_hash::(node_hash, cb, db, storage)? + ) + }, EncodedNode::Branch(encoded_children, val) => { - let mut child = |i:usize| { - encoded_children[i].map(|data| - Self::inline_or_hash::(data, db, storage) - ) + let mut child = |i:usize| match encoded_children[i] { + Some(child) => Self::inline_or_hash::(node_hash, child, db, storage) + .map(Some), + None => Ok(None), }; let children = Box::new([ - child(0), child(1), child(2), child(3), - child(4), child(5), child(6), child(7), - child(8), child(9), child(10), child(11), - child(12), child(13), child(14), child(15), + child(0)?, child(1)?, child(2)?, child(3)?, + child(4)?, child(5)?, child(6)?, child(7)?, + child(8)?, child(9)?, child(10)?, child(11)?, + child(12)?, child(13)?, child(14)?, child(15)?, ]); Node::Branch(children, val.map(DBValue::from_slice)) }, EncodedNode::NibbledBranch(k, encoded_children, val) => { - let mut child = |i:usize| { - encoded_children[i].map(|data| - Self::inline_or_hash::(data, db, storage) - ) + let mut child = |i:usize| match encoded_children[i] { + Some(child) => Self::inline_or_hash::(node_hash, child, db, storage) + .map(Some), + None => Ok(None), }; let children = Box::new([ - child(0), child(1), child(2), child(3), - child(4), child(5), child(6), child(7), - child(8), child(9), child(10), child(11), - child(12), child(13), child(14), child(15), + child(0)?, child(1)?, child(2)?, child(3)?, + child(4)?, child(5)?, child(6)?, child(7)?, + child(8)?, child(9)?, child(10)?, child(11)?, + child(12)?, child(13)?, child(14)?, child(15)?, ]); Node::NibbledBranch(k.into(), children, val.map(DBValue::from_slice)) }, - } + }; + Ok(node) } // TODO: parallelize fn into_encoded(self, mut child_cb: F) -> Vec where - C: NodeCodec, + C: NodeCodec, F: FnMut(NodeHandle, Option<&NibbleSlice>, Option) -> ChildReference, H: Hasher, { @@ -428,10 +440,11 @@ where let node_encoded = self.db.get(&hash, key) .ok_or_else(|| Box::new(TrieError::IncompleteDatabase(hash)))?; let node = Node::from_encoded::( + hash, &node_encoded, &*self.db, &mut self.storage - ); + )?; Ok(self.storage.alloc(Stored::Cached(node, hash))) } @@ -1605,7 +1618,7 @@ mod tests { } fn reference_hashed_null_node() -> ::Out { - >::hashed_null_node() + as NodeCodec>::hashed_null_node() } #[test]