use crate::{
fatdb::FatDBDoubleEndedIterator,
iterator::{TrieDBNodeDoubleEndedIterator, TrieDBRawIterator},
lookup::Lookup,
nibble::NibbleSlice,
node::{decode_hash, NodeHandle, OwnedNode},
rstd::boxed::Box,
CError, DBValue, MerkleValue, Query, Result, Trie, TrieAccess, TrieCache,
TrieDoubleEndedIterator, TrieError, TrieHash, TrieItem, TrieIterator, TrieKeyItem, TrieLayout,
TrieRecorder,
};
#[cfg(feature = "std")]
use crate::{
nibble::NibbleVec,
node::Node,
rstd::{fmt, vec::Vec},
};
use hash_db::{HashDBRef, Prefix, EMPTY_PREFIX};
pub struct TrieDBBuilder<'db, 'cache, L: TrieLayout> {
db: &'db dyn HashDBRef<L::Hash, DBValue>,
root: &'db TrieHash<L>,
cache: Option<&'cache mut dyn TrieCache<L::Codec>>,
recorder: Option<&'cache mut dyn TrieRecorder<TrieHash<L>>>,
}
impl<'db, 'cache, L: TrieLayout> TrieDBBuilder<'db, 'cache, L> {
#[inline]
pub fn new(db: &'db dyn HashDBRef<L::Hash, DBValue>, root: &'db TrieHash<L>) -> Self {
Self { db, root, cache: None, recorder: None }
}
#[inline]
pub fn with_cache(mut self, cache: &'cache mut dyn TrieCache<L::Codec>) -> Self {
self.cache = Some(cache);
self
}
#[inline]
pub fn with_optional_cache<'ocache: 'cache>(
mut self,
cache: Option<&'ocache mut dyn TrieCache<L::Codec>>,
) -> Self {
self.cache = cache.map(|c| c as _);
self
}
#[inline]
pub fn with_recorder(mut self, recorder: &'cache mut dyn TrieRecorder<TrieHash<L>>) -> Self {
self.recorder = Some(recorder);
self
}
#[inline]
pub fn with_optional_recorder<'recorder: 'cache>(
mut self,
recorder: Option<&'recorder mut dyn TrieRecorder<TrieHash<L>>>,
) -> Self {
self.recorder = recorder.map(|r| r as _);
self
}
#[inline]
pub fn build(self) -> TrieDB<'db, 'cache, L> {
TrieDB {
db: self.db,
root: self.root,
cache: self.cache.map(core::cell::RefCell::new),
recorder: self.recorder.map(core::cell::RefCell::new),
}
}
}
pub struct TrieDB<'db, 'cache, L>
where
L: TrieLayout,
{
db: &'db dyn HashDBRef<L::Hash, DBValue>,
root: &'db TrieHash<L>,
cache: Option<core::cell::RefCell<&'cache mut dyn TrieCache<L::Codec>>>,
recorder: Option<core::cell::RefCell<&'cache mut dyn TrieRecorder<TrieHash<L>>>>,
}
impl<'db, 'cache, L> TrieDB<'db, 'cache, L>
where
L: TrieLayout,
{
pub fn db(&'db self) -> &'db dyn HashDBRef<L::Hash, DBValue> {
self.db
}
pub fn into_double_ended_iter(
&'db self,
) -> Result<TrieDBDoubleEndedIterator<'db, 'cache, L>, TrieHash<L>, CError<L>> {
TrieDBDoubleEndedIterator::new(&self)
}
pub fn into_node_double_ended_iter(
&'db self,
) -> Result<TrieDBNodeDoubleEndedIterator<'db, 'cache, L>, TrieHash<L>, CError<L>> {
TrieDBNodeDoubleEndedIterator::new(&self)
}
pub fn into_key_double_ended_iter(
&'db self,
) -> Result<TrieDBKeyDoubleEndedIterator<'db, 'cache, L>, TrieHash<L>, CError<L>> {
TrieDBKeyDoubleEndedIterator::new(&self)
}
pub fn into_fat_double_ended_iter(
&'db self,
) -> Result<FatDBDoubleEndedIterator<'db, 'cache, L>, TrieHash<L>, CError<L>> {
FatDBDoubleEndedIterator::new(&self)
}
pub(crate) fn get_raw_or_lookup(
&self,
parent_hash: TrieHash<L>,
node_handle: NodeHandle,
partial_key: Prefix,
record_access: bool,
) -> Result<(OwnedNode<DBValue>, Option<TrieHash<L>>), TrieHash<L>, CError<L>> {
let (node_hash, node_data) = match node_handle {
NodeHandle::Hash(data) => {
let node_hash = decode_hash::<L::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, data.to_vec()),
};
let owned_node = OwnedNode::new::<L::Codec>(node_data)
.map_err(|e| Box::new(TrieError::DecoderError(node_hash.unwrap_or(parent_hash), e)))?;
if record_access {
if let Some((hash, recorder)) =
node_hash.as_ref().and_then(|h| self.recorder.as_ref().map(|r| (h, r)))
{
recorder.borrow_mut().record(TrieAccess::EncodedNode {
hash: *hash,
encoded_node: owned_node.data().into(),
});
}
}
Ok((owned_node, node_hash))
}
pub(crate) fn fetch_value(
&self,
hash: TrieHash<L>,
prefix: Prefix,
) -> Result<DBValue, TrieHash<L>, CError<L>> {
let value = self
.db
.get(&hash, prefix)
.ok_or_else(|| Box::new(TrieError::IncompleteDatabase(hash)))?;
if let Some(recorder) = self.recorder.as_ref() {
debug_assert!(prefix.1.is_none(), "A value has never a partial key; qed");
recorder.borrow_mut().record(TrieAccess::Value {
hash,
value: value.as_slice().into(),
full_key: prefix.0,
});
}
Ok(value)
}
}
impl<'db, 'cache, L> Trie<L> for TrieDB<'db, 'cache, L>
where
L: TrieLayout,
{
fn root(&self) -> &TrieHash<L> {
self.root
}
fn get_hash(&self, key: &[u8]) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
let mut cache = self.cache.as_ref().map(|c| c.borrow_mut());
let mut recorder = self.recorder.as_ref().map(|r| r.borrow_mut());
Lookup::<L, _> {
db: self.db,
query: |_: &[u8]| (),
hash: *self.root,
cache: cache.as_mut().map(|c| &mut ***c as &mut dyn TrieCache<L::Codec>),
recorder: recorder.as_mut().map(|r| &mut ***r as &mut dyn TrieRecorder<TrieHash<L>>),
}
.look_up_hash(key, NibbleSlice::new(key))
}
fn get_with<Q: Query<L::Hash>>(
&self,
key: &[u8],
query: Q,
) -> Result<Option<Q::Item>, TrieHash<L>, CError<L>> {
let mut cache = self.cache.as_ref().map(|c| c.borrow_mut());
let mut recorder = self.recorder.as_ref().map(|r| r.borrow_mut());
Lookup::<L, Q> {
db: self.db,
query,
hash: *self.root,
cache: cache.as_mut().map(|c| &mut ***c as &mut dyn TrieCache<L::Codec>),
recorder: recorder.as_mut().map(|r| &mut ***r as &mut dyn TrieRecorder<TrieHash<L>>),
}
.look_up(key, NibbleSlice::new(key))
}
fn lookup_first_descendant(
&self,
key: &[u8],
) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>> {
let mut cache = self.cache.as_ref().map(|c| c.borrow_mut());
let mut recorder = self.recorder.as_ref().map(|r| r.borrow_mut());
Lookup::<L, _> {
db: self.db,
query: |_: &[u8]| (),
hash: *self.root,
cache: cache.as_mut().map(|c| &mut ***c as &mut dyn TrieCache<L::Codec>),
recorder: recorder.as_mut().map(|r| &mut ***r as &mut dyn TrieRecorder<TrieHash<L>>),
}
.lookup_first_descendant(key, NibbleSlice::new(key))
}
fn iter<'a>(
&'a self,
) -> Result<
Box<dyn TrieIterator<L, Item = TrieItem<TrieHash<L>, CError<L>>> + 'a>,
TrieHash<L>,
CError<L>,
> {
TrieDBIterator::new(self).map(|iter| Box::new(iter) as Box<_>)
}
fn key_iter<'a>(
&'a self,
) -> Result<
Box<dyn TrieIterator<L, Item = TrieKeyItem<TrieHash<L>, CError<L>>> + 'a>,
TrieHash<L>,
CError<L>,
> {
TrieDBKeyIterator::new(self).map(|iter| Box::new(iter) as Box<_>)
}
}
#[cfg(feature = "std")]
struct TrieAwareDebugNode<'db, 'cache, 'a, L>
where
L: TrieLayout,
{
trie: &'db TrieDB<'db, 'cache, L>,
node_key: NodeHandle<'a>,
partial_key: NibbleVec,
index: Option<u8>,
}
#[cfg(feature = "std")]
impl<'db, 'cache, 'a, L> fmt::Debug for TrieAwareDebugNode<'db, 'cache, 'a, L>
where
L: TrieLayout,
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
match self.trie.get_raw_or_lookup(
<TrieHash<L>>::default(),
self.node_key,
self.partial_key.as_prefix(),
false,
) {
Ok((owned_node, _node_hash)) => match owned_node.node() {
Node::Leaf(slice, value) => {
let mut disp = f.debug_struct("Node::Leaf");
if let Some(i) = self.index {
disp.field("index", &i);
}
disp.field("slice", &slice).field("value", &value);
disp.finish()
},
Node::Extension(slice, item) => {
let mut disp = f.debug_struct("Node::Extension");
if let Some(i) = self.index {
disp.field("index", &i);
}
disp.field("slice", &slice).field(
"item",
&TrieAwareDebugNode {
trie: self.trie,
node_key: item,
partial_key: self
.partial_key
.clone_append_optional_slice_and_nibble(Some(&slice), None),
index: None,
},
);
disp.finish()
},
Node::Branch(ref nodes, ref value) => {
let nodes: Vec<TrieAwareDebugNode<L>> = nodes
.into_iter()
.enumerate()
.filter_map(|(i, n)| n.map(|n| (i, n)))
.map(|(i, n)| TrieAwareDebugNode {
trie: self.trie,
index: Some(i as u8),
node_key: n,
partial_key: self
.partial_key
.clone_append_optional_slice_and_nibble(None, Some(i as u8)),
})
.collect();
let mut disp = f.debug_struct("Node::Branch");
if let Some(i) = self.index {
disp.field("index", &i);
}
disp.field("nodes", &nodes).field("value", &value);
disp.finish()
},
Node::NibbledBranch(slice, nodes, value) => {
let nodes: Vec<TrieAwareDebugNode<L>> = nodes
.iter()
.enumerate()
.filter_map(|(i, n)| n.map(|n| (i, n)))
.map(|(i, n)| TrieAwareDebugNode {
trie: self.trie,
index: Some(i as u8),
node_key: n,
partial_key: self.partial_key.clone_append_optional_slice_and_nibble(
Some(&slice),
Some(i as u8),
),
})
.collect();
let mut disp = f.debug_struct("Node::NibbledBranch");
if let Some(i) = self.index {
disp.field("index", &i);
}
disp.field("slice", &slice).field("nodes", &nodes).field("value", &value);
disp.finish()
},
Node::Empty => {
let mut disp = f.debug_struct("Node::Empty");
disp.finish()
},
},
Err(e) => f
.debug_struct("BROKEN_NODE")
.field("index", &self.index)
.field("key", &self.node_key)
.field("error", &format!("ERROR fetching node: {}", e))
.finish(),
}
}
}
#[cfg(feature = "std")]
impl<'db, 'cache, L> fmt::Debug for TrieDB<'db, 'cache, L>
where
L: TrieLayout,
{
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("TrieDB")
.field(
"root",
&TrieAwareDebugNode {
trie: self,
node_key: NodeHandle::Hash(self.root().as_ref()),
partial_key: NibbleVec::new(),
index: None,
},
)
.finish()
}
}
pub struct TrieDBIterator<'a, 'cache, L: TrieLayout> {
db: &'a TrieDB<'a, 'cache, L>,
raw_iter: TrieDBRawIterator<L>,
}
pub struct TrieDBKeyIterator<'a, 'cache, L: TrieLayout> {
db: &'a TrieDB<'a, 'cache, L>,
raw_iter: TrieDBRawIterator<L>,
}
pub struct TrieDBKeyDoubleEndedIterator<'a, 'cache, L: TrieLayout> {
db: &'a TrieDB<'a, 'cache, L>,
raw_iter: TrieDBRawIterator<L>,
back_raw_iter: TrieDBRawIterator<L>,
}
impl<'a, 'cache, L: TrieLayout> TrieDBKeyDoubleEndedIterator<'a, 'cache, L> {
pub fn new(db: &'a TrieDB<'a, 'cache, L>) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self {
db,
raw_iter: TrieDBRawIterator::new(db)?,
back_raw_iter: TrieDBRawIterator::new(db)?,
})
}
}
pub struct TrieDBDoubleEndedIterator<'a, 'cache, L: TrieLayout> {
db: &'a TrieDB<'a, 'cache, L>,
raw_iter: TrieDBRawIterator<L>,
back_raw_iter: TrieDBRawIterator<L>,
}
impl<'a, 'cache, L: TrieLayout> TrieDBDoubleEndedIterator<'a, 'cache, L> {
pub fn new(db: &'a TrieDB<'a, 'cache, L>) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self {
db,
raw_iter: TrieDBRawIterator::new(db)?,
back_raw_iter: TrieDBRawIterator::new(db)?,
})
}
pub fn new_prefixed(
db: &'a TrieDB<'a, 'cache, L>,
prefix: &[u8],
) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self {
db,
raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)?,
back_raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)?,
})
}
}
impl<L: TrieLayout> TrieDoubleEndedIterator<L> for TrieDBDoubleEndedIterator<'_, '_, L> {}
impl<'a, 'cache, L: TrieLayout> TrieDBIterator<'a, 'cache, L> {
pub fn new(db: &'a TrieDB<'a, 'cache, L>) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self { db, raw_iter: TrieDBRawIterator::new(db)? })
}
pub fn new_prefixed(
db: &'a TrieDB<'a, 'cache, L>,
prefix: &[u8],
) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self { db, raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)? })
}
pub fn new_prefixed_then_seek(
db: &'a TrieDB<'a, 'cache, L>,
prefix: &[u8],
start_at: &[u8],
) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self { db, raw_iter: TrieDBRawIterator::new_prefixed_then_seek(db, prefix, start_at)? })
}
pub fn from_raw(db: &'a TrieDB<'a, 'cache, L>, raw_iter: TrieDBRawIterator<L>) -> Self {
Self { db, raw_iter }
}
pub fn into_raw(self) -> TrieDBRawIterator<L> {
self.raw_iter
}
}
impl<'a, 'cache, L: TrieLayout> TrieIterator<L> for TrieDBIterator<'a, 'cache, L> {
fn seek(&mut self, key: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
self.raw_iter.seek(self.db, key, true).map(|_| ())
}
}
impl<'a, 'cache, L: TrieLayout> Iterator for TrieDBIterator<'a, 'cache, L> {
type Item = TrieItem<TrieHash<L>, CError<L>>;
fn next(&mut self) -> Option<Self::Item> {
self.raw_iter.next_item(self.db)
}
}
impl<'a, 'cache, L: TrieLayout> TrieIterator<L> for TrieDBDoubleEndedIterator<'a, 'cache, L> {
fn seek(&mut self, key: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
self.raw_iter.seek(self.db, key, true).map(|_| ())?;
self.back_raw_iter.seek(self.db, key, false).map(|_| ())
}
}
impl<'a, 'cache, L: TrieLayout> Iterator for TrieDBDoubleEndedIterator<'a, 'cache, L> {
type Item = TrieItem<TrieHash<L>, CError<L>>;
fn next(&mut self) -> Option<Self::Item> {
self.raw_iter.next_item(self.db)
}
}
impl<'a, 'cache, L: TrieLayout> DoubleEndedIterator for TrieDBDoubleEndedIterator<'a, 'cache, L> {
fn next_back(&mut self) -> Option<Self::Item> {
self.back_raw_iter.prev_item(self.db)
}
}
impl<'a, 'cache, L: TrieLayout> TrieDBKeyIterator<'a, 'cache, L> {
pub fn new(db: &'a TrieDB<'a, 'cache, L>) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self { db, raw_iter: TrieDBRawIterator::new(db)? })
}
pub fn new_prefixed(
db: &'a TrieDB<'a, 'cache, L>,
prefix: &[u8],
) -> Result<Self, TrieHash<L>, CError<L>> {
Ok(Self { db, raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)? })
}
pub fn new_prefixed_then_seek(
db: &'a TrieDB<'a, 'cache, L>,
prefix: &[u8],
start_at: &[u8],
) -> Result<TrieDBKeyIterator<'a, 'cache, L>, TrieHash<L>, CError<L>> {
Ok(Self { db, raw_iter: TrieDBRawIterator::new_prefixed_then_seek(db, prefix, start_at)? })
}
pub fn from_raw(db: &'a TrieDB<'a, 'cache, L>, raw_iter: TrieDBRawIterator<L>) -> Self {
Self { db, raw_iter }
}
pub fn into_raw(self) -> TrieDBRawIterator<L> {
self.raw_iter
}
}
impl<'a, 'cache, L: TrieLayout> TrieIterator<L> for TrieDBKeyIterator<'a, 'cache, L> {
fn seek(&mut self, key: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
self.raw_iter.seek(self.db, key, true).map(|_| ())
}
}
impl<'a, 'cache, L: TrieLayout> Iterator for TrieDBKeyIterator<'a, 'cache, L> {
type Item = TrieKeyItem<TrieHash<L>, CError<L>>;
fn next(&mut self) -> Option<Self::Item> {
self.raw_iter.next_key(self.db)
}
}
impl<'a, 'cache, L: TrieLayout> TrieIterator<L> for TrieDBKeyDoubleEndedIterator<'a, 'cache, L> {
fn seek(&mut self, key: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
self.raw_iter.seek(self.db, key, true).map(|_| ())?;
self.back_raw_iter.seek(self.db, key, false).map(|_| ())
}
}
impl<'a, 'cache, L: TrieLayout> Iterator for TrieDBKeyDoubleEndedIterator<'a, 'cache, L> {
type Item = TrieKeyItem<TrieHash<L>, CError<L>>;
fn next(&mut self) -> Option<Self::Item> {
self.raw_iter.next_key(self.db)
}
}
impl<'a, 'cache, L: TrieLayout> DoubleEndedIterator
for TrieDBKeyDoubleEndedIterator<'a, 'cache, L>
{
fn next_back(&mut self) -> Option<Self::Item> {
self.back_raw_iter.prev_key(self.db)
}
}