use super::{DBValue, node::NodeKey};
use super::{Result, TrieError, TrieMut, TrieLayout, TrieHash, CError};
use super::lookup::Lookup;
use super::node::{NodeHandle as EncodedNodeHandle, Node as EncodedNode, decode_hash};
use hash_db::{HashDB, Hasher, Prefix, EMPTY_PREFIX};
use hashbrown::HashSet;
use crate::node_codec::NodeCodec;
use crate::nibble::{NibbleVec, NibbleSlice, nibble_ops, BackingByteVec};
use crate::rstd::{
boxed::Box, convert::TryFrom, hash::Hash, mem, ops::Index, result, vec::Vec, VecDeque,
};
#[cfg(feature = "std")]
use log::trace;
#[cfg(feature = "std")]
use crate::rstd::fmt::{self, Debug};
#[cfg_attr(feature = "std", derive(Debug))]
struct StorageHandle(usize);
#[cfg_attr(feature = "std", derive(Debug))]
enum NodeHandle<H> {
InMemory(StorageHandle),
Hash(H),
}
impl<H> From<StorageHandle> for NodeHandle<H> {
fn from(handle: StorageHandle) -> Self {
NodeHandle::InMemory(handle)
}
}
fn empty_children<H>() -> Box<[Option<NodeHandle<H>>; 16]> {
Box::new([
None, None, None, None, None, None, None, None,
None, None, None, None, None, None, None, None,
])
}
type NibbleFullKey<'key> = NibbleSlice<'key>;
enum Node<H> {
Empty,
Leaf(NodeKey, DBValue),
Extension(NodeKey, NodeHandle<H>),
Branch(Box<[Option<NodeHandle<H>>; 16]>, Option<DBValue>),
NibbledBranch(NodeKey, Box<[Option<NodeHandle<H>>; 16]>, Option<DBValue>),
}
#[cfg(feature = "std")]
struct ToHex<'a>(&'a [u8]);
#[cfg(feature = "std")]
impl<'a> Debug for ToHex<'a> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
let hex = rustc_hex::ToHexIter::new(self.0.iter());
for b in hex {
write!(fmt, "{}", b)?;
}
Ok(())
}
}
#[cfg(feature = "std")]
impl<H: Debug> Debug for Node<H> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
match *self {
Self::Empty => write!(fmt, "Empty"),
Self::Leaf((ref a, ref b), ref c) =>
write!(fmt, "Leaf({:?}, {:?})", (a, ToHex(&*b)), ToHex(&*c)),
Self::Extension((ref a, ref b), ref c) =>
write!(fmt, "Extension({:?}, {:?})", (a, ToHex(&*b)), c),
Self::Branch(ref a, ref b) =>
write!(fmt, "Branch({:?}, {:?}", a, b.as_ref().map(Vec::as_slice).map(ToHex)),
Self::NibbledBranch((ref a, ref b), ref c, ref d) =>
write!(fmt, "NibbledBranch({:?}, {:?}, {:?})", (a, ToHex(&*b)), c, d.as_ref().map(Vec::as_slice).map(ToHex)),
}
}
}
impl<O> Node<O>
where
O: AsRef<[u8]> + AsMut<[u8]> + Default + crate::MaybeDebug
+ PartialEq + Eq + Hash + Send + Sync + Clone + Copy
{
fn inline_or_hash<C, H>(
parent_hash: H::Out,
child: EncodedNodeHandle,
db: &dyn HashDB<H, DBValue>,
storage: &mut NodeStorage<H::Out>
) -> Result<NodeHandle<H::Out>, H::Out, C::Error>
where
C: NodeCodec<HashOut=O>,
H: Hasher<Out=O>,
{
let handle = match child {
EncodedNodeHandle::Hash(data) => {
let hash = decode_hash::<H>(data)
.ok_or_else(|| Box::new(TrieError::InvalidHash(parent_hash, data.to_vec())))?;
NodeHandle::Hash(hash)
},
EncodedNodeHandle::Inline(data) => {
let child = Node::from_encoded::<C, H>(parent_hash, data, db, storage)?;
NodeHandle::InMemory(storage.alloc(Stored::New(child)))
},
};
Ok(handle)
}
fn from_encoded<'a, 'b, C, H>(
node_hash: H::Out,
data: &'a[u8],
db: &dyn HashDB<H, DBValue>,
storage: &'b mut NodeStorage<H::Out>,
) -> Result<Self, H::Out, C::Error>
where
C: NodeCodec<HashOut = O>, H: Hasher<Out = O>,
{
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(), v.to_vec()),
EncodedNode::Extension(key, cb) => {
Node::Extension(
key.into(),
Self::inline_or_hash::<C, H>(node_hash, cb, db, storage)?
)
},
EncodedNode::Branch(encoded_children, val) => {
let mut child = |i:usize| match encoded_children[i] {
Some(child) => Self::inline_or_hash::<C, H>(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)?,
]);
Node::Branch(children, val.map(|v| v.to_vec()))
},
EncodedNode::NibbledBranch(k, encoded_children, val) => {
let mut child = |i:usize| match encoded_children[i] {
Some(child) => Self::inline_or_hash::<C, H>(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)?,
]);
Node::NibbledBranch(k.into(), children, val.map(|v| v.to_vec()))
},
};
Ok(node)
}
fn into_encoded<F, C, H>(self, mut child_cb: F) -> Vec<u8>
where
C: NodeCodec<HashOut=O>,
F: FnMut(NodeHandle<H::Out>, Option<&NibbleSlice>, Option<u8>) -> ChildReference<H::Out>,
H: Hasher<Out = O>,
{
match self {
Node::Empty => C::empty_node().to_vec(),
Node::Leaf(partial, value) => {
let pr = NibbleSlice::new_offset(&partial.1[..], partial.0);
C::leaf_node(pr.right(), &value)
},
Node::Extension(partial, child) => {
let pr = NibbleSlice::new_offset(&partial.1[..], partial.0);
let it = pr.right_iter();
let c = child_cb(child, Some(&pr), None);
C::extension_node(
it,
pr.len(),
c,
)
},
Node::Branch(mut children, value) => {
C::branch_node(
children.iter_mut()
.map(Option::take)
.enumerate()
.map(|(i, maybe_child)| {
maybe_child.map(|child| child_cb(child, None, Some(i as u8)))
}),
value.as_ref().map(|v| &v[..])
)
},
Node::NibbledBranch(partial, mut children, value) => {
let pr = NibbleSlice::new_offset(&partial.1[..], partial.0);
let it = pr.right_iter();
C::branch_node_nibbled(
it,
pr.len(),
children.iter_mut()
.map(Option::take)
.enumerate()
.map(|(i, maybe_child)| {
maybe_child.map(|child| {
let pr = NibbleSlice::new_offset(&partial.1[..], partial.0);
child_cb(child, Some(&pr), Some(i as u8))
})
}),
value.as_ref().map(|v| &v[..])
)
},
}
}
}
enum Action<H> {
Replace(Node<H>),
Restore(Node<H>),
Delete,
}
enum InsertAction<H> {
Replace(Node<H>),
Restore(Node<H>),
}
impl<H> InsertAction<H> {
fn into_action(self) -> Action<H> {
match self {
InsertAction::Replace(n) => Action::Replace(n),
InsertAction::Restore(n) => Action::Restore(n),
}
}
fn unwrap_node(self) -> Node<H> {
match self {
InsertAction::Replace(n) | InsertAction::Restore(n) => n,
}
}
}
enum Stored<H> {
New(Node<H>),
Cached(Node<H>, H),
}
#[derive(Clone, Copy)]
#[cfg_attr(feature = "std", derive(Debug))]
pub enum ChildReference<HO> {
Hash(HO),
Inline(HO, usize),
}
impl<'a, HO> TryFrom<EncodedNodeHandle<'a>> for ChildReference<HO>
where HO: AsRef<[u8]> + AsMut<[u8]> + Default + Clone + Copy
{
type Error = Vec<u8>;
fn try_from(handle: EncodedNodeHandle<'a>) -> result::Result<Self, Vec<u8>> {
match handle {
EncodedNodeHandle::Hash(data) => {
let mut hash = HO::default();
if data.len() != hash.as_ref().len() {
return Err(data.to_vec());
}
hash.as_mut().copy_from_slice(data);
Ok(ChildReference::Hash(hash))
}
EncodedNodeHandle::Inline(data) => {
let mut hash = HO::default();
if data.len() > hash.as_ref().len() {
return Err(data.to_vec());
}
&mut hash.as_mut()[..data.len()].copy_from_slice(data);
Ok(ChildReference::Inline(hash, data.len()))
}
}
}
}
struct NodeStorage<H> {
nodes: Vec<Stored<H>>,
free_indices: VecDeque<usize>,
}
impl<H> NodeStorage<H> {
fn empty() -> Self {
NodeStorage {
nodes: Vec::new(),
free_indices: VecDeque::new(),
}
}
fn alloc(&mut self, stored: Stored<H>) -> StorageHandle {
if let Some(idx) = self.free_indices.pop_front() {
self.nodes[idx] = stored;
StorageHandle(idx)
} else {
self.nodes.push(stored);
StorageHandle(self.nodes.len() - 1)
}
}
fn destroy(&mut self, handle: StorageHandle) -> Stored<H> {
let idx = handle.0;
self.free_indices.push_back(idx);
mem::replace(&mut self.nodes[idx], Stored::New(Node::Empty))
}
}
impl<'a, H> Index<&'a StorageHandle> for NodeStorage<H> {
type Output = Node<H>;
fn index(&self, handle: &'a StorageHandle) -> &Node<H> {
match self.nodes[handle.0] {
Stored::New(ref node) => node,
Stored::Cached(ref node, _) => node,
}
}
}
pub struct TrieDBMut<'a, L>
where
L: TrieLayout,
{
storage: NodeStorage<TrieHash<L>>,
db: &'a mut dyn HashDB<L::Hash, DBValue>,
root: &'a mut TrieHash<L>,
root_handle: NodeHandle<TrieHash<L>>,
death_row: HashSet<(TrieHash<L>, (BackingByteVec, Option<u8>))>,
hash_count: usize,
}
impl<'a, L> TrieDBMut<'a, L>
where
L: TrieLayout,
{
pub fn new(db: &'a mut dyn HashDB<L::Hash, DBValue>, root: &'a mut TrieHash<L>) -> Self {
*root = L::Codec::hashed_null_node();
let root_handle = NodeHandle::Hash(L::Codec::hashed_null_node());
TrieDBMut {
storage: NodeStorage::empty(),
db,
root,
root_handle,
death_row: HashSet::new(),
hash_count: 0,
}
}
pub fn from_existing(
db: &'a mut dyn HashDB<L::Hash, DBValue>,
root: &'a mut TrieHash<L>,
) -> Result<Self, TrieHash<L>, CError<L>> {
if !db.contains(root, EMPTY_PREFIX) {
return Err(Box::new(TrieError::InvalidStateRoot(*root)));
}
let root_handle = NodeHandle::Hash(*root);
Ok(TrieDBMut {
storage: NodeStorage::empty(),
db,
root,
root_handle,
death_row: HashSet::new(),
hash_count: 0,
})
}
pub fn db(&self) -> &dyn HashDB<L::Hash, DBValue> {
self.db
}
pub fn db_mut(&mut self) -> &mut dyn HashDB<L::Hash, DBValue> {
self.db
}
fn cache(
&mut self,
hash: TrieHash<L>,
key: Prefix,
) -> Result<StorageHandle, TrieHash<L>, CError<L>> {
let node_encoded = self.db.get(&hash, key)
.ok_or_else(|| Box::new(TrieError::IncompleteDatabase(hash)))?;
let node = Node::from_encoded::<L::Codec, L::Hash>(
hash,
&node_encoded,
&*self.db,
&mut self.storage
)?;
Ok(self.storage.alloc(Stored::Cached(node, hash)))
}
fn inspect<F>(
&mut self,
stored: Stored<TrieHash<L>>,
key: &mut NibbleFullKey,
inspector: F,
) -> Result<Option<(Stored<TrieHash<L>>, bool)>, TrieHash<L>, CError<L>>
where
F: FnOnce(
&mut Self,
Node<TrieHash<L>>,
&mut NibbleFullKey,
) -> Result<Action<TrieHash<L>>, TrieHash<L>, CError<L>>,
{
let current_key = key.clone();
Ok(match stored {
Stored::New(node) => match inspector(self, node, key)? {
Action::Restore(node) => Some((Stored::New(node), false)),
Action::Replace(node) => Some((Stored::New(node), true)),
Action::Delete => None,
},
Stored::Cached(node, hash) => match inspector(self, node, key)? {
Action::Restore(node) => Some((Stored::Cached(node, hash), false)),
Action::Replace(node) => {
self.death_row.insert((hash, current_key.left_owned()));
Some((Stored::New(node), true))
}
Action::Delete => {
self.death_row.insert((hash, current_key.left_owned()));
None
}
},
})
}
fn lookup<'x, 'key>(
&'x self,
mut partial: NibbleSlice<'key>,
handle: &NodeHandle<TrieHash<L>>,
) -> Result<Option<DBValue>, TrieHash<L>, CError<L>>
where 'x: 'key
{
let mut handle = handle;
loop {
let (mid, child) = match *handle {
NodeHandle::Hash(ref hash) => return Lookup::<L, _> {
db: &self.db,
query: |v: &[u8]| v.to_vec(),
hash: *hash,
}.look_up(partial),
NodeHandle::InMemory(ref handle) => match self.storage[handle] {
Node::Empty => return Ok(None),
Node::Leaf(ref key, ref value) => {
if NibbleSlice::from_stored(key) == partial {
return Ok(Some(value.to_vec()));
} else {
return Ok(None);
}
},
Node::Extension(ref slice, ref child) => {
let slice = NibbleSlice::from_stored(slice);
if partial.starts_with(&slice) {
(slice.len(), child)
} else {
return Ok(None);
}
},
Node::Branch(ref children, ref value) => {
if partial.is_empty() {
return Ok(value.as_ref().map(|v| v.to_vec()));
} else {
let idx = partial.at(0);
match children[idx as usize].as_ref() {
Some(child) => (1, child),
None => return Ok(None),
}
}
},
Node::NibbledBranch(ref slice, ref children, ref value) => {
let slice = NibbleSlice::from_stored(slice);
if partial.is_empty() {
return Ok(value.as_ref().map(|v| v.to_vec()));
} else if partial.starts_with(&slice) {
let idx = partial.at(0);
match children[idx as usize].as_ref() {
Some(child) => (1 + slice.len(), child),
None => return Ok(None),
}
} else {
return Ok(None)
}
},
}
};
partial = partial.mid(mid);
handle = child;
}
}
fn insert_at(
&mut self,
handle: NodeHandle<TrieHash<L>>,
key: &mut NibbleFullKey,
value: DBValue,
old_val: &mut Option<DBValue>,
) -> Result<(StorageHandle, bool), TrieHash<L>, CError<L>> {
let h = match handle {
NodeHandle::InMemory(h) => h,
NodeHandle::Hash(h) => self.cache(h, key.left())?,
};
let stored = self.storage.destroy(h);
let (new_stored, changed) = self.inspect(stored, key, move |trie, stored, key| {
trie.insert_inspector(stored, key, value, old_val).map(|a| a.into_action())
})?.expect("Insertion never deletes.");
Ok((self.storage.alloc(new_stored), changed))
}
fn insert_inspector(
&mut self,
node: Node<TrieHash<L>>,
key: &mut NibbleFullKey,
value: DBValue,
old_val: &mut Option<DBValue>,
) -> Result<InsertAction<TrieHash<L>>, TrieHash<L>, CError<L>> {
let partial = *key;
#[cfg(feature = "std")]
trace!(target: "trie", "augmented (partial: {:?}, value: {:?})", partial, ToHex(&value));
Ok(match node {
Node::Empty => {
#[cfg(feature = "std")]
trace!(target: "trie", "empty: COMPOSE");
InsertAction::Replace(Node::Leaf(partial.to_stored(), value))
},
Node::Branch(mut children, stored_value) => {
debug_assert!(L::USE_EXTENSION);
#[cfg(feature = "std")]
trace!(target: "trie", "branch: ROUTE,AUGMENT");
if partial.is_empty() {
let unchanged = stored_value.as_ref() == Some(&value);
let branch = Node::Branch(children, Some(value));
*old_val = stored_value;
match unchanged {
true => InsertAction::Restore(branch),
false => InsertAction::Replace(branch),
}
} else {
let idx = partial.at(0) as usize;
key.advance(1);
if let Some(child) = children[idx].take() {
let (new_child, changed) = self.insert_at(child, key, value, old_val)?;
children[idx] = Some(new_child.into());
if !changed {
return Ok(InsertAction::Restore(Node::Branch(children, stored_value)));
}
} else {
let leaf = self.storage.alloc(
Stored::New(Node::Leaf(key.to_stored(), value))
);
children[idx] = Some(leaf.into());
}
InsertAction::Replace(Node::Branch(children, stored_value))
}
},
Node::NibbledBranch(encoded, mut children, stored_value) => {
debug_assert!(!L::USE_EXTENSION);
#[cfg(feature = "std")]
trace!(target: "trie", "branch: ROUTE,AUGMENT");
let existing_key = NibbleSlice::from_stored(&encoded);
let common = partial.common_prefix(&existing_key);
if common == existing_key.len() && common == partial.len() {
let unchanged = stored_value.as_ref() == Some(&value);
let branch = Node::NibbledBranch(
existing_key.to_stored(),
children,
Some(value),
);
*old_val = stored_value;
match unchanged {
true => InsertAction::Restore(branch),
false => InsertAction::Replace(branch),
}
} else if common < existing_key.len() {
#[cfg(feature = "std")]
trace!(
target: "trie",
"partially-shared-prefix (exist={:?}; new={:?}; common={:?}):\
AUGMENT-AT-END",
existing_key.len(),
partial.len(),
common,
);
let nbranch_partial = existing_key.mid(common + 1).to_stored();
let low = Node::NibbledBranch(nbranch_partial, children, stored_value);
let ix = existing_key.at(common);
let mut children = empty_children();
let alloc_storage = self.storage.alloc(Stored::New(low));
children[ix as usize] = Some(alloc_storage.into());
if partial.len() - common == 0 {
InsertAction::Replace(Node::NibbledBranch(
existing_key.to_stored_range(common),
children,
Some(value),
)
)
} else {
let ix = partial.at(common);
let stored_leaf = Node::Leaf(partial.mid(common + 1).to_stored(), value);
let leaf = self.storage.alloc(Stored::New(stored_leaf));
children[ix as usize] = Some(leaf.into());
InsertAction::Replace(Node::NibbledBranch(
existing_key.to_stored_range(common),
children,
None,
)
)
}
} else {
#[cfg(feature = "std")]
trace!(target: "trie", "branch: ROUTE,AUGMENT");
let idx = partial.at(common) as usize;
key.advance(common + 1);
if let Some(child) = children[idx].take() {
let (new_child, changed) = self.insert_at(child, key, value, old_val)?;
children[idx] = Some(new_child.into());
if !changed {
let n_branch = Node::NibbledBranch(
existing_key.to_stored(),
children,
stored_value,
);
return Ok(InsertAction::Restore(n_branch));
}
} else {
let leaf = self.storage.alloc(
Stored::New(Node::Leaf(key.to_stored(), value)),
);
children[idx] = Some(leaf.into());
}
InsertAction::Replace(Node::NibbledBranch(
existing_key.to_stored(),
children,
stored_value,
))
}
},
Node::Leaf(encoded, stored_value) => {
let existing_key = NibbleSlice::from_stored(&encoded);
let common = partial.common_prefix(&existing_key);
if common == existing_key.len() && common == partial.len() {
#[cfg(feature = "std")]
trace!(target: "trie", "equivalent-leaf: REPLACE");
let unchanged = stored_value == value;
*old_val = Some(stored_value);
match unchanged {
true => InsertAction::Restore(Node::Leaf(encoded.clone(), value)),
false => InsertAction::Replace(Node::Leaf(encoded.clone(), value)),
}
} else if (L::USE_EXTENSION && common == 0)
|| (!L::USE_EXTENSION && common < existing_key.len()) {
#[cfg(feature = "std")]
trace!(
target: "trie",
"lesser-common-prefix, not-both-empty (exist={:?}; new={:?}):\
TRANSMUTE,AUGMENT",
existing_key.len(),
partial.len(),
);
let mut children = empty_children();
let branch = if L::USE_EXTENSION && existing_key.is_empty() {
Node::Branch(children, Some(stored_value))
} else {
let idx = existing_key.at(common) as usize;
let new_leaf = Node::Leaf(
existing_key.mid(common + 1).to_stored(),
stored_value,
);
children[idx] = Some(self.storage.alloc(Stored::New(new_leaf)).into());
if L::USE_EXTENSION {
Node::Branch(children, None)
} else {
Node::NibbledBranch(partial.to_stored_range(common), children, None)
}
};
let branch_action = self.insert_inspector(branch, key, value, old_val)?
.unwrap_node();
InsertAction::Replace(branch_action)
} else if !L::USE_EXTENSION {
#[cfg(feature = "std")]
trace!(target: "trie", "complete-prefix (common={:?}): AUGMENT-AT-END", common);
let branch = Node::NibbledBranch(
existing_key.to_stored(),
empty_children(),
Some(stored_value),
);
let branch = self.insert_inspector(branch, key, value, old_val)?
.unwrap_node();
InsertAction::Replace(branch)
} else if common == existing_key.len() {
debug_assert!(L::USE_EXTENSION);
#[cfg(feature = "std")]
trace!(target: "trie", "complete-prefix (common={:?}): AUGMENT-AT-END", common);
let branch = Node::Branch(empty_children(), Some(stored_value));
key.advance(common);
let branch = self.insert_inspector(branch, key, value, old_val)?.unwrap_node();
let branch_handle = self.storage.alloc(Stored::New(branch)).into();
InsertAction::Replace(Node::Extension(existing_key.to_stored(), branch_handle))
} else {
debug_assert!(L::USE_EXTENSION);
#[cfg(feature = "std")]
trace!(
target: "trie",
"partially-shared-prefix (exist={:?}; new={:?}; common={:?}):\
AUGMENT-AT-END",
existing_key.len(),
partial.len(),
common,
);
let low = Node::Leaf(existing_key.mid(common).to_stored(), stored_value);
key.advance(common);
let augmented_low = self.insert_inspector(low, key, value, old_val)?
.unwrap_node();
InsertAction::Replace(Node::Extension(
existing_key.to_stored_range(common),
self.storage.alloc(Stored::New(augmented_low)).into()
))
}
},
Node::Extension(encoded, child_branch) => {
debug_assert!(L::USE_EXTENSION);
let existing_key = NibbleSlice::from_stored(&encoded);
let common = partial.common_prefix(&existing_key);
if common == 0 {
#[cfg(feature = "std")]
trace!(
target: "trie",
"no-common-prefix, not-both-empty (exist={:?}; new={:?}):\
TRANSMUTE,AUGMENT",
existing_key.len(),
partial.len(),
);
assert!(!existing_key.is_empty());
let idx = existing_key.at(0) as usize;
let mut children = empty_children();
children[idx] = if existing_key.len() == 1 {
Some(child_branch)
} else {
let ext = Node::Extension(existing_key.mid(1).to_stored(), child_branch);
Some(self.storage.alloc(Stored::New(ext)).into())
};
let branch_action = self.insert_inspector(
Node::Branch(children, None),
key,
value,
old_val,
)?.unwrap_node();
InsertAction::Replace(branch_action)
} else if common == existing_key.len() {
#[cfg(feature = "std")]
trace!(target: "trie", "complete-prefix (common={:?}): AUGMENT-AT-END", common);
key.advance(common);
let (new_child, changed) = self.insert_at(child_branch, key, value, old_val)?;
let new_ext = Node::Extension(existing_key.to_stored(), new_child.into());
match changed {
true => InsertAction::Replace(new_ext),
false => InsertAction::Restore(new_ext),
}
} else {
#[cfg(feature = "std")]
trace!(
target: "trie",
"partially-shared-prefix (exist={:?}; new={:?}; common={:?}):\
AUGMENT-AT-END",
existing_key.len(),
partial.len(),
common,
);
let low = Node::Extension(existing_key.mid(common).to_stored(), child_branch);
key.advance(common);
let augmented_low = self.insert_inspector(low, key, value, old_val)?
.unwrap_node();
InsertAction::Replace(Node::Extension(
existing_key.to_stored_range(common),
self.storage.alloc(Stored::New(augmented_low)).into()
))
}
},
})
}
fn remove_at(
&mut self,
handle: NodeHandle<TrieHash<L>>,
key: &mut NibbleFullKey,
old_val: &mut Option<DBValue>,
) -> Result<Option<(StorageHandle, bool)>, TrieHash<L>, CError<L>> {
let stored = match handle {
NodeHandle::InMemory(h) => self.storage.destroy(h),
NodeHandle::Hash(h) => {
let handle = self.cache(h, key.left())?;
self.storage.destroy(handle)
}
};
let opt = self.inspect(
stored,
key,
move |trie, node, key| trie.remove_inspector(node, key, old_val),
)?;
Ok(opt.map(|(new, changed)| (self.storage.alloc(new), changed)))
}
fn remove_inspector(
&mut self,
node: Node<TrieHash<L>>,
key: &mut NibbleFullKey,
old_val: &mut Option<DBValue>,
) -> Result<Action<TrieHash<L>>, TrieHash<L>, CError<L>> {
let partial = *key;
Ok(match (node, partial.is_empty()) {
(Node::Empty, _) => Action::Delete,
(Node::Branch(c, None), true) => Action::Restore(Node::Branch(c, None)),
(Node::NibbledBranch(n, c, None), true) =>
Action::Restore(Node::NibbledBranch(n, c, None)),
(Node::Branch(children, Some(val)), true) => {
*old_val = Some(val);
Action::Replace(self.fix(Node::Branch(children, None), key.clone())?)
},
(Node::NibbledBranch(n, children, Some(val)), true) => {
*old_val = Some(val);
Action::Replace(self.fix(Node::NibbledBranch(n, children, None), key.clone())?)
},
(Node::Branch(mut children, value), false) => {
let idx = partial.at(0) as usize;
if let Some(child) = children[idx].take() {
#[cfg(feature = "std")]
trace!(
target: "trie",
"removing value out of branch child, partial={:?}",
partial,
);
let prefix = *key;
key.advance(1);
match self.remove_at(child, key, old_val)? {
Some((new, changed)) => {
children[idx] = Some(new.into());
let branch = Node::Branch(children, value);
match changed {
true => Action::Replace(branch),
false => Action::Restore(branch),
}
}
None => {
#[cfg(feature = "std")]
trace!(target: "trie", "branch child deleted, partial={:?}", partial);
Action::Replace(self.fix(Node::Branch(children, value), prefix)?)
}
}
} else {
Action::Restore(Node::Branch(children, value))
}
},
(Node::NibbledBranch(encoded, mut children, value), false) => {
let (common, existing_length) = {
let existing_key = NibbleSlice::from_stored(&encoded);
(existing_key.common_prefix(&partial), existing_key.len())
};
if common == existing_length && common == partial.len() {
if let Some(val) = value {
*old_val = Some(val);
let f = self.fix(Node::NibbledBranch(encoded, children, None), key.clone());
Action::Replace(f?)
} else {
Action::Restore(Node::NibbledBranch(encoded, children, None))
}
} else if common < existing_length {
Action::Restore(Node::NibbledBranch(encoded, children, value))
} else {
let idx = partial.at(common) as usize;
if let Some(child) = children[idx].take() {
#[cfg(feature = "std")]
trace!(
target: "trie",
"removing value out of branch child, partial={:?}",
partial,
);
let prefix = *key;
key.advance(common + 1);
match self.remove_at(child, key, old_val)? {
Some((new, changed)) => {
children[idx] = Some(new.into());
let branch = Node::NibbledBranch(encoded, children, value);
match changed {
true => Action::Replace(branch),
false => Action::Restore(branch),
}
},
None => {
#[cfg(feature = "std")]
trace!(
target: "trie",
"branch child deleted, partial={:?}",
partial,
);
Action::Replace(
self.fix(Node::NibbledBranch(encoded, children, value), prefix)?
)
},
}
} else {
Action::Restore(Node::NibbledBranch(encoded, children, value))
}
}
},
(Node::Leaf(encoded, value), _) => {
if NibbleSlice::from_stored(&encoded) == partial {
*old_val = Some(value);
Action::Delete
} else {
#[cfg(feature = "std")]
trace!(
target: "trie",
"restoring leaf wrong partial, partial={:?}, existing={:?}",
partial,
NibbleSlice::from_stored(&encoded),
);
Action::Restore(Node::Leaf(encoded, value))
}
},
(Node::Extension(encoded, child_branch), _) => {
let (common, existing_length) = {
let existing_key = NibbleSlice::from_stored(&encoded);
(existing_key.common_prefix(&partial), existing_key.len())
};
if common == existing_length {
#[cfg(feature = "std")]
trace!(target: "trie", "removing from extension child, partial={:?}", partial);
let prefix = *key;
key.advance(common);
match self.remove_at(child_branch, key, old_val)? {
Some((new_child, changed)) => {
let new_child = new_child.into();
match changed {
true => Action::Replace(
self.fix(Node::Extension(encoded, new_child), prefix)?
),
false => Action::Restore(Node::Extension(encoded, new_child)),
}
}
None => {
Action::Delete
}
}
} else {
Action::Restore(Node::Extension(encoded, child_branch))
}
},
})
}
fn fix(
&mut self,
node: Node<TrieHash<L>>,
key: NibbleSlice,
) -> Result<Node<TrieHash<L>>, TrieHash<L>, CError<L>> {
match node {
Node::Branch(mut children, value) => {
#[cfg_attr(feature = "std", derive(Debug))]
enum UsedIndex {
None,
One(u8),
Many,
};
let mut used_index = UsedIndex::None;
for i in 0..16 {
match (children[i].is_none(), &used_index) {
(false, &UsedIndex::None) => used_index = UsedIndex::One(i as u8),
(false, &UsedIndex::One(_)) => {
used_index = UsedIndex::Many;
break;
}
_ => continue,
}
}
match (used_index, value) {
(UsedIndex::None, None) =>
panic!("Branch with no subvalues. Something went wrong."),
(UsedIndex::One(a), None) => {
let new_partial = NibbleSlice::new_offset(&[a], 1).to_stored();
let child = children[a as usize].take()
.expect("used_index only set if occupied; qed");
let new_node = Node::Extension(new_partial, child);
self.fix(new_node, key)
}
(UsedIndex::None, Some(value)) => {
#[cfg(feature = "std")]
trace!(target: "trie", "fixing: branch -> leaf");
Ok(Node::Leaf(NibbleSlice::new(&[]).to_stored(), value))
}
(_, value) => {
#[cfg(feature = "std")]
trace!(target: "trie", "fixing: restoring branch");
Ok(Node::Branch(children, value))
}
}
},
Node::NibbledBranch(enc_nibble, mut children, value) => {
#[cfg_attr(feature = "std", derive(Debug))]
enum UsedIndex {
None,
One(u8),
Many,
};
let mut used_index = UsedIndex::None;
for i in 0..16 {
match (children[i].is_none(), &used_index) {
(false, &UsedIndex::None) => used_index = UsedIndex::One(i as u8),
(false, &UsedIndex::One(_)) => {
used_index = UsedIndex::Many;
break;
}
_ => continue,
}
}
match (used_index, value) {
(UsedIndex::None, None) =>
panic!("Branch with no subvalues. Something went wrong."),
(UsedIndex::One(a), None) => {
let child = children[a as usize].take()
.expect("used_index only set if occupied; qed");
let mut key2 = key.clone();
key2.advance((enc_nibble.1.len() * nibble_ops::NIBBLE_PER_BYTE) - enc_nibble.0);
let (start, alloc_start, prefix_end) = match key2.left() {
(start, None) => (start, None, Some(nibble_ops::push_at_left(0, a, 0))),
(start, Some(v)) => {
let mut so: BackingByteVec = start.into();
so.push(nibble_ops::pad_left(v) | a);
(start, Some(so), None)
},
};
let child_prefix = (alloc_start.as_ref().map(|start| &start[..]).unwrap_or(start), prefix_end);
let stored = match child {
NodeHandle::InMemory(h) => self.storage.destroy(h),
NodeHandle::Hash(h) => {
let handle = self.cache(h, child_prefix)?;
self.storage.destroy(handle)
}
};
let child_node = match stored {
Stored::New(node) => node,
Stored::Cached(node, hash) => {
self.death_row.insert((
hash,
(child_prefix.0[..].into(), child_prefix.1),
));
node
},
};
match child_node {
Node::Leaf(sub_partial, value) => {
let mut enc_nibble = enc_nibble;
combine_key(
&mut enc_nibble,
(nibble_ops::NIBBLE_PER_BYTE - 1, &[a][..]),
);
combine_key(
&mut enc_nibble,
(sub_partial.0, &sub_partial.1[..]),
);
Ok(Node::Leaf(enc_nibble, value))
},
Node::NibbledBranch(sub_partial, ch_children, ch_value) => {
let mut enc_nibble = enc_nibble;
combine_key(
&mut enc_nibble,
(nibble_ops::NIBBLE_PER_BYTE - 1, &[a][..]),
);
combine_key(
&mut enc_nibble,
(sub_partial.0, &sub_partial.1[..]),
);
Ok(Node::NibbledBranch(enc_nibble, ch_children, ch_value))
},
_ => unreachable!(),
}
},
(UsedIndex::None, Some(value)) => {
#[cfg(feature = "std")]
trace!(target: "trie", "fixing: branch -> leaf");
Ok(Node::Leaf(enc_nibble, value))
},
(_, value) => {
#[cfg(feature = "std")]
trace!(target: "trie", "fixing: restoring branch");
Ok(Node::NibbledBranch(enc_nibble, children, value))
},
}
},
Node::Extension(partial, child) => {
let last = partial.1[partial.1.len() - 1] & (255 >> 4);
let mut key2 = key.clone();
key2.advance((partial.1.len() * nibble_ops::NIBBLE_PER_BYTE) - partial.0 - 1);
let (start, alloc_start, prefix_end) = match key2.left() {
(start, None) => (start, None, Some(nibble_ops::push_at_left(0, last, 0))),
(start, Some(v)) => {
let mut so: BackingByteVec = start.into();
so.push(nibble_ops::pad_left(v) | last);
(start, Some(so), None)
},
};
let child_prefix = (alloc_start.as_ref().map(|start| &start[..]).unwrap_or(start), prefix_end);
let stored = match child {
NodeHandle::InMemory(h) => self.storage.destroy(h),
NodeHandle::Hash(h) => {
let handle = self.cache(h, child_prefix)?;
self.storage.destroy(handle)
}
};
let (child_node, maybe_hash) = match stored {
Stored::New(node) => (node, None),
Stored::Cached(node, hash) => (node, Some(hash))
};
match child_node {
Node::Extension(sub_partial, sub_child) => {
if let Some(hash) = maybe_hash {
self.death_row.insert(
(hash, (child_prefix.0[..].into(), child_prefix.1)),
);
}
let mut partial = partial;
combine_key(&mut partial, (sub_partial.0, &sub_partial.1[..]));
#[cfg(feature = "std")]
trace!(
target: "trie",
"fixing: extension combination. new_partial={:?}",
partial,
);
self.fix(Node::Extension(partial, sub_child), key)
}
Node::Leaf(sub_partial, value) => {
if let Some(hash) = maybe_hash {
self.death_row.insert((hash, (child_prefix.0[..].into(), child_prefix.1)));
}
let mut partial = partial;
combine_key(&mut partial, (sub_partial.0, &sub_partial.1[..]));
#[cfg(feature = "std")]
trace!(
target: "trie",
"fixing: extension -> leaf. new_partial={:?}",
partial,
);
Ok(Node::Leaf(partial, value))
}
child_node => {
#[cfg(feature = "std")]
trace!(target: "trie", "fixing: restoring extension");
let stored = if let Some(hash) = maybe_hash {
Stored::Cached(child_node, hash)
} else {
Stored::New(child_node)
};
Ok(Node::Extension(partial, self.storage.alloc(stored).into()))
}
}
},
other => Ok(other),
}
}
pub fn commit(&mut self) {
#[cfg(feature = "std")]
trace!(target: "trie", "Committing trie changes to db.");
#[cfg(feature = "std")]
trace!(target: "trie", "{:?} nodes to remove from db", self.death_row.len());
for (hash, prefix) in self.death_row.drain() {
self.db.remove(&hash, (&prefix.0[..], prefix.1));
}
let handle = match self.root_handle() {
NodeHandle::Hash(_) => return,
NodeHandle::InMemory(h) => h,
};
match self.storage.destroy(handle) {
Stored::New(node) => {
let mut k = NibbleVec::new();
let encoded_root = node.into_encoded::<_, L::Codec, L::Hash>(
|child, o_slice, o_index| {
let mov = k.append_optional_slice_and_nibble(o_slice, o_index);
let cr = self.commit_child(child, &mut k);
k.drop_lasts(mov);
cr
}
);
#[cfg(feature = "std")]
trace!(target: "trie", "encoded root node: {:#x?}", &encoded_root[..]);
*self.root = self.db.insert(EMPTY_PREFIX, &encoded_root[..]);
self.hash_count += 1;
self.root_handle = NodeHandle::Hash(*self.root);
}
Stored::Cached(node, hash) => {
*self.root = hash;
self.root_handle = NodeHandle::InMemory(
self.storage.alloc(Stored::Cached(node, hash)),
);
}
}
}
fn commit_child(
&mut self,
handle: NodeHandle<TrieHash<L>>,
prefix: &mut NibbleVec,
) -> ChildReference<TrieHash<L>> {
match handle {
NodeHandle::Hash(hash) => ChildReference::Hash(hash),
NodeHandle::InMemory(storage_handle) => {
match self.storage.destroy(storage_handle) {
Stored::Cached(_, hash) => ChildReference::Hash(hash),
Stored::New(node) => {
let encoded = {
let commit_child = |
node_handle,
o_slice: Option<&NibbleSlice>,
o_index: Option<u8>
| {
let mov = prefix.append_optional_slice_and_nibble(o_slice, o_index);
let cr = self.commit_child(node_handle, prefix);
prefix.drop_lasts(mov);
cr
};
node.into_encoded::<_, L::Codec, L::Hash>(commit_child)
};
if encoded.len() >= L::Hash::LENGTH {
let hash = self.db.insert(prefix.as_prefix(), &encoded[..]);
self.hash_count +=1;
ChildReference::Hash(hash)
} else {
let mut h = <TrieHash<L>>::default();
let len = encoded.len();
h.as_mut()[..len].copy_from_slice(&encoded[..len]);
ChildReference::Inline(h, len)
}
}
}
}
}
}
fn root_handle(&self) -> NodeHandle<TrieHash<L>> {
match self.root_handle {
NodeHandle::Hash(h) => NodeHandle::Hash(h),
NodeHandle::InMemory(StorageHandle(x)) => NodeHandle::InMemory(StorageHandle(x)),
}
}
}
impl<'a, L> TrieMut<L> for TrieDBMut<'a, L>
where
L: TrieLayout,
{
fn root(&mut self) -> &TrieHash<L> {
self.commit();
self.root
}
fn is_empty(&self) -> bool {
match self.root_handle {
NodeHandle::Hash(h) => h == L::Codec::hashed_null_node(),
NodeHandle::InMemory(ref h) => match self.storage[h] {
Node::Empty => true,
_ => false,
}
}
}
fn get<'x, 'key>(&'x self, key: &'key [u8]) -> Result<Option<DBValue>, TrieHash<L>, CError<L>>
where 'x: 'key
{
self.lookup(NibbleSlice::new(key), &self.root_handle)
}
fn insert(
&mut self,
key: &[u8],
value: &[u8],
) -> Result<Option<DBValue>, TrieHash<L>, CError<L>> {
if !L::ALLOW_EMPTY && value.is_empty() { return self.remove(key) }
let mut old_val = None;
#[cfg(feature = "std")]
trace!(target: "trie", "insert: key={:#x?}, value={:?}", key, ToHex(&value));
let root_handle = self.root_handle();
let (new_handle, _changed) = self.insert_at(
root_handle,
&mut NibbleSlice::new(key),
value.to_vec(),
&mut old_val,
)?;
#[cfg(feature = "std")]
trace!(target: "trie", "insert: altered trie={}", _changed);
self.root_handle = NodeHandle::InMemory(new_handle);
Ok(old_val)
}
fn remove(&mut self, key: &[u8]) -> Result<Option<DBValue>, TrieHash<L>, CError<L>> {
#[cfg(feature = "std")]
trace!(target: "trie", "remove: key={:#x?}", key);
let root_handle = self.root_handle();
let mut key = NibbleSlice::new(key);
let mut old_val = None;
match self.remove_at(root_handle, &mut key, &mut old_val)? {
Some((handle, _changed)) => {
#[cfg(feature = "std")]
trace!(target: "trie", "remove: altered trie={}", _changed);
self.root_handle = NodeHandle::InMemory(handle);
}
None => {
#[cfg(feature = "std")]
trace!(target: "trie", "remove: obliterated trie");
self.root_handle = NodeHandle::Hash(L::Codec::hashed_null_node());
*self.root = L::Codec::hashed_null_node();
}
}
Ok(old_val)
}
}
impl<'a, L> Drop for TrieDBMut<'a, L>
where
L: TrieLayout,
{
fn drop(&mut self) {
self.commit();
}
}
fn combine_key(start: &mut NodeKey, end: (usize, &[u8])) {
debug_assert!(start.0 < nibble_ops::NIBBLE_PER_BYTE);
debug_assert!(end.0 < nibble_ops::NIBBLE_PER_BYTE);
let final_offset = (start.0 + end.0) % nibble_ops::NIBBLE_PER_BYTE;
let _shifted = nibble_ops::shift_key(start, final_offset);
let st = if end.0 > 0 {
let sl = start.1.len();
start.1[sl - 1] |= nibble_ops::pad_right(end.1[0]);
1
} else {
0
};
(st..end.1.len()).for_each(|i| start.1.push(end.1[i]));
}
#[cfg(test)]
mod tests {
use crate::nibble::BackingByteVec;
#[test]
fn combine_test() {
let a: BackingByteVec = [0x12, 0x34][..].into();
let b: &[u8] = [0x56, 0x78][..].into();
let test_comb = |a: (_, &BackingByteVec), b, c| {
let mut a = (a.0, a.1.clone());
super::combine_key(&mut a, b);
assert_eq!((a.0, &a.1[..]), c);
};
test_comb((0, &a), (0, &b), (0, &[0x12, 0x34, 0x56, 0x78][..]));
test_comb((1, &a), (0, &b), (1, &[0x12, 0x34, 0x56, 0x78][..]));
test_comb((0, &a), (1, &b), (1, &[0x01, 0x23, 0x46, 0x78][..]));
test_comb((1, &a), (1, &b), (0, &[0x23, 0x46, 0x78][..]));
}
#[test]
fn nice_debug_for_node() {
use super::Node;
let e: Node<u32> = Node::Leaf((1, vec![1, 2, 3].into()), vec![4, 5, 6]);
assert_eq!(format!("{:?}", e), "Leaf((1, 010203), 040506)");
}
}