tetsy_trie_db/
triedbmut.rs

1// Copyright 2017, 2020 Parity Technologies
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! In-memory trie representation.
16
17use super::{DBValue, node::NodeKey};
18use super::{Result, TrieError, TrieMut, TrieLayout, TrieHash, CError};
19use super::lookup::Lookup;
20use super::node::{NodeHandle as EncodedNodeHandle, Node as EncodedNode, decode_hash};
21
22use tetsy_hash_db::{HashDB, Hasher, Prefix, EMPTY_PREFIX};
23use hashbrown::HashSet;
24
25use crate::node_codec::NodeCodec;
26use crate::nibble::{NibbleVec, NibbleSlice, nibble_ops, BackingByteVec};
27use crate::rstd::{
28	boxed::Box, convert::TryFrom, hash::Hash, mem, ops::Index, result, vec::Vec, VecDeque,
29};
30
31#[cfg(feature = "std")]
32use log::trace;
33
34#[cfg(feature = "std")]
35use crate::rstd::fmt::{self, Debug};
36
37
38// For lookups into the Node storage buffer.
39// This is deliberately non-copyable.
40#[cfg_attr(feature = "std", derive(Debug))]
41struct StorageHandle(usize);
42
43// Handles to nodes in the trie.
44#[cfg_attr(feature = "std", derive(Debug))]
45enum NodeHandle<H> {
46	/// Loaded into memory.
47	InMemory(StorageHandle),
48	/// Either a hash or an inline node
49	Hash(H),
50}
51
52impl<H> From<StorageHandle> for NodeHandle<H> {
53	fn from(handle: StorageHandle) -> Self {
54		NodeHandle::InMemory(handle)
55	}
56}
57
58fn empty_children<H>() -> Box<[Option<NodeHandle<H>>; 16]> {
59	Box::new([
60		None, None, None, None, None, None, None, None,
61		None, None, None, None, None, None, None, None,
62	])
63}
64
65/// Type alias to indicate the nible covers a full key,
66/// therefore its left side is a full prefix.
67type NibbleFullKey<'key> = NibbleSlice<'key>;
68
69/// Node types in the Trie.
70enum Node<H> {
71	/// Empty node.
72	Empty,
73	/// A leaf node contains the end of a key and a value.
74	/// This key is encoded from a `NibbleSlice`, meaning it contains
75	/// a flag indicating it is a leaf.
76	Leaf(NodeKey, DBValue),
77	/// An extension contains a shared portion of a key and a child node.
78	/// The shared portion is encoded from a `NibbleSlice` meaning it contains
79	/// a flag indicating it is an extension.
80	/// The child node is always a branch.
81	Extension(NodeKey, NodeHandle<H>),
82	/// A branch has up to 16 children and an optional value.
83	Branch(Box<[Option<NodeHandle<H>>; 16]>, Option<DBValue>),
84	/// Branch node with support for a nibble (to avoid extension node).
85	NibbledBranch(NodeKey, Box<[Option<NodeHandle<H>>; 16]>, Option<DBValue>),
86}
87
88#[cfg(feature = "std")]
89struct ToHex<'a>(&'a [u8]);
90#[cfg(feature = "std")]
91impl<'a> Debug for ToHex<'a> {
92	fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
93		let hex = rustc_hex::ToHexIter::new(self.0.iter());
94		for b in hex {
95			write!(fmt, "{}", b)?;
96		}
97		Ok(())
98	}
99}
100
101#[cfg(feature = "std")]
102impl<H: Debug> Debug for Node<H> {
103	fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
104		match *self {
105			Self::Empty => write!(fmt, "Empty"),
106			Self::Leaf((ref a, ref b), ref c) =>
107				write!(fmt, "Leaf({:?}, {:?})", (a, ToHex(&*b)), ToHex(&*c)),
108			Self::Extension((ref a, ref b), ref c) =>
109				write!(fmt, "Extension({:?}, {:?})", (a, ToHex(&*b)), c),
110			Self::Branch(ref a, ref b) =>
111				write!(fmt, "Branch({:?}, {:?}", a, b.as_ref().map(Vec::as_slice).map(ToHex)),
112			Self::NibbledBranch((ref a, ref b), ref c, ref d) =>
113				write!(fmt, "NibbledBranch({:?}, {:?}, {:?})", (a, ToHex(&*b)), c, d.as_ref().map(Vec::as_slice).map(ToHex)),
114		}
115	}
116}
117
118impl<O> Node<O>
119where
120	O: AsRef<[u8]> + AsMut<[u8]> + Default + crate::MaybeDebug
121		+ PartialEq + Eq + Hash + Send + Sync + Clone + Copy
122{
123	// load an inline node into memory or get the hash to do the lookup later.
124	fn inline_or_hash<C, H>(
125		parent_hash: H::Out,
126		child: EncodedNodeHandle,
127		db: &dyn HashDB<H, DBValue>,
128		storage: &mut NodeStorage<H::Out>
129	) -> Result<NodeHandle<H::Out>, H::Out, C::Error>
130	where
131		C: NodeCodec<HashOut=O>,
132		H: Hasher<Out=O>,
133	{
134		let handle = match child {
135			EncodedNodeHandle::Hash(data) => {
136				let hash = decode_hash::<H>(data)
137					.ok_or_else(|| Box::new(TrieError::InvalidHash(parent_hash, data.to_vec())))?;
138				NodeHandle::Hash(hash)
139			},
140			EncodedNodeHandle::Inline(data) => {
141				let child = Node::from_encoded::<C, H>(parent_hash, data, db, storage)?;
142				NodeHandle::InMemory(storage.alloc(Stored::New(child)))
143			},
144		};
145		Ok(handle)
146	}
147
148	// Decode a node from encoded bytes.
149	fn from_encoded<'a, 'b, C, H>(
150		node_hash: H::Out,
151		data: &'a[u8],
152		db: &dyn HashDB<H, DBValue>,
153		storage: &'b mut NodeStorage<H::Out>,
154	) -> Result<Self, H::Out, C::Error>
155		where
156			C: NodeCodec<HashOut = O>, H: Hasher<Out = O>,
157	{
158		let encoded_node = C::decode(data)
159			.map_err(|e| Box::new(TrieError::DecoderError(node_hash, e)))?;
160		let node = match encoded_node {
161			EncodedNode::Empty => Node::Empty,
162			EncodedNode::Leaf(k, v) => Node::Leaf(k.into(), v.to_vec()),
163			EncodedNode::Extension(key, cb) => {
164				Node::Extension(
165					key.into(),
166					Self::inline_or_hash::<C, H>(node_hash, cb, db, storage)?
167				)
168			},
169			EncodedNode::Branch(encoded_children, val) => {
170				let mut child = |i:usize| match encoded_children[i] {
171					Some(child) => Self::inline_or_hash::<C, H>(node_hash, child, db, storage)
172						.map(Some),
173					None => Ok(None),
174				};
175
176				let children = Box::new([
177					child(0)?, child(1)?, child(2)?, child(3)?,
178					child(4)?, child(5)?, child(6)?, child(7)?,
179					child(8)?, child(9)?, child(10)?, child(11)?,
180					child(12)?, child(13)?, child(14)?, child(15)?,
181				]);
182
183				Node::Branch(children, val.map(|v| v.to_vec()))
184			},
185			EncodedNode::NibbledBranch(k, encoded_children, val) => {
186				let mut child = |i:usize| match encoded_children[i] {
187					Some(child) => Self::inline_or_hash::<C, H>(node_hash, child, db, storage)
188						.map(Some),
189					None => Ok(None),
190				};
191
192				let children = Box::new([
193					child(0)?, child(1)?, child(2)?, child(3)?,
194					child(4)?, child(5)?, child(6)?, child(7)?,
195					child(8)?, child(9)?, child(10)?, child(11)?,
196					child(12)?, child(13)?, child(14)?, child(15)?,
197				]);
198
199				Node::NibbledBranch(k.into(), children, val.map(|v| v.to_vec()))
200			},
201		};
202		Ok(node)
203	}
204
205	// TODO: parallelize
206	fn into_encoded<F, C, H>(self, mut child_cb: F) -> Vec<u8>
207	where
208		C: NodeCodec<HashOut=O>,
209		F: FnMut(NodeHandle<H::Out>, Option<&NibbleSlice>, Option<u8>) -> ChildReference<H::Out>,
210		H: Hasher<Out = O>,
211	{
212		match self {
213			Node::Empty => C::empty_node().to_vec(),
214			Node::Leaf(partial, value) => {
215				let pr = NibbleSlice::new_offset(&partial.1[..], partial.0);
216				C::leaf_node(pr.right(), &value)
217			},
218			Node::Extension(partial, child) => {
219				let pr = NibbleSlice::new_offset(&partial.1[..], partial.0);
220				let it = pr.right_iter();
221				let c = child_cb(child, Some(&pr), None);
222				C::extension_node(
223					it,
224					pr.len(),
225					c,
226				)
227			},
228			Node::Branch(mut children, value) => {
229				C::branch_node(
230					// map the `NodeHandle`s from the Branch to `ChildReferences`
231					children.iter_mut()
232						.map(Option::take)
233						.enumerate()
234						.map(|(i, maybe_child)| {
235							maybe_child.map(|child| child_cb(child, None, Some(i as u8)))
236						}),
237					value.as_ref().map(|v| &v[..])
238				)
239			},
240			Node::NibbledBranch(partial, mut children, value) => {
241				let pr = NibbleSlice::new_offset(&partial.1[..], partial.0);
242				let it = pr.right_iter();
243				C::branch_node_nibbled(
244					it,
245					pr.len(),
246					// map the `NodeHandle`s from the Branch to `ChildReferences`
247					children.iter_mut()
248						.map(Option::take)
249						.enumerate()
250						.map(|(i, maybe_child)| {
251							//let branch_index = [i as u8];
252							maybe_child.map(|child| {
253								let pr = NibbleSlice::new_offset(&partial.1[..], partial.0);
254								child_cb(child, Some(&pr), Some(i as u8))
255							})
256						}),
257					value.as_ref().map(|v| &v[..])
258				)
259			},
260		}
261	}
262}
263
264// post-inspect action.
265enum Action<H> {
266	// Replace a node with a new one.
267	Replace(Node<H>),
268	// Restore the original node. This trusts that the node is actually the original.
269	Restore(Node<H>),
270	// if it is a new node, just clears the storage.
271	Delete,
272}
273
274// post-insert action. Same as action without delete
275enum InsertAction<H> {
276	// Replace a node with a new one.
277	Replace(Node<H>),
278	// Restore the original node.
279	Restore(Node<H>),
280}
281
282impl<H> InsertAction<H> {
283	fn into_action(self) -> Action<H> {
284		match self {
285			InsertAction::Replace(n) => Action::Replace(n),
286			InsertAction::Restore(n) => Action::Restore(n),
287		}
288	}
289
290	// unwrap the node, disregarding replace or restore state.
291	fn unwrap_node(self) -> Node<H> {
292		match self {
293			InsertAction::Replace(n) | InsertAction::Restore(n) => n,
294		}
295	}
296}
297
298// What kind of node is stored here.
299enum Stored<H> {
300	// A new node.
301	New(Node<H>),
302	// A cached node, loaded from the DB.
303	Cached(Node<H>, H),
304}
305
306/// Used to build a collection of child nodes from a collection of `NodeHandle`s
307#[derive(Clone, Copy)]
308#[cfg_attr(feature = "std", derive(Debug))]
309pub enum ChildReference<HO> { // `HO` is e.g. `H256`, i.e. the output of a `Hasher`
310	Hash(HO),
311	Inline(HO, usize), // usize is the length of the node data we store in the `H::Out`
312}
313
314impl<'a, HO> TryFrom<EncodedNodeHandle<'a>> for ChildReference<HO>
315	where HO: AsRef<[u8]> + AsMut<[u8]> + Default + Clone + Copy
316{
317	type Error = Vec<u8>;
318
319	fn try_from(handle: EncodedNodeHandle<'a>) -> result::Result<Self, Vec<u8>> {
320		match handle {
321			EncodedNodeHandle::Hash(data) => {
322				let mut hash = HO::default();
323				if data.len() != hash.as_ref().len() {
324					return Err(data.to_vec());
325				}
326				hash.as_mut().copy_from_slice(data);
327				Ok(ChildReference::Hash(hash))
328			}
329			EncodedNodeHandle::Inline(data) => {
330				let mut hash = HO::default();
331				if data.len() > hash.as_ref().len() {
332					return Err(data.to_vec());
333				}
334				&mut hash.as_mut()[..data.len()].copy_from_slice(data);
335				Ok(ChildReference::Inline(hash, data.len()))
336			}
337		}
338	}
339}
340
341/// Compact and cache-friendly storage for Trie nodes.
342struct NodeStorage<H> {
343	nodes: Vec<Stored<H>>,
344	free_indices: VecDeque<usize>,
345}
346
347impl<H> NodeStorage<H> {
348	/// Create a new storage.
349	fn empty() -> Self {
350		NodeStorage {
351			nodes: Vec::new(),
352			free_indices: VecDeque::new(),
353		}
354	}
355
356	/// Allocate a new node in the storage.
357	fn alloc(&mut self, stored: Stored<H>) -> StorageHandle {
358		if let Some(idx) = self.free_indices.pop_front() {
359			self.nodes[idx] = stored;
360			StorageHandle(idx)
361		} else {
362			self.nodes.push(stored);
363			StorageHandle(self.nodes.len() - 1)
364		}
365	}
366
367	/// Remove a node from the storage, consuming the handle and returning the node.
368	fn destroy(&mut self, handle: StorageHandle) -> Stored<H> {
369		let idx = handle.0;
370
371		self.free_indices.push_back(idx);
372		mem::replace(&mut self.nodes[idx], Stored::New(Node::Empty))
373	}
374}
375
376impl<'a, H> Index<&'a StorageHandle> for NodeStorage<H> {
377	type Output = Node<H>;
378
379	fn index(&self, handle: &'a StorageHandle) -> &Node<H> {
380		match self.nodes[handle.0] {
381			Stored::New(ref node) => node,
382			Stored::Cached(ref node, _) => node,
383		}
384	}
385}
386
387/// A `Trie` implementation using a generic `HashDB` backing database.
388///
389/// Use it as a `TrieMut` trait object. You can use `db()` to get the backing database object.
390/// Note that changes are not committed to the database until `commit` is called.
391///
392/// Querying the root or dropping the trie will commit automatically.
393///
394///
395/// # Example
396/// ```ignore
397/// use tetsy_hash_db::Hasher;
398/// use tetsy_reference_trie::{RefTrieDBMut, TrieMut};
399/// use tetsy_trie_db::DBValue;
400/// use tetsy_keccak_hasher::KeccakHasher;
401/// use tetsy_memory_db::*;
402///
403/// let mut memdb = MemoryDB::<KeccakHasher, HashKey<_>, DBValue>::default();
404/// let mut root = Default::default();
405/// let mut t = RefTrieDBMut::new(&mut memdb, &mut root);
406/// assert!(t.is_empty());
407/// assert_eq!(*t.root(), KeccakHasher::hash(&[0u8][..]));
408/// t.insert(b"foo", b"bar").unwrap();
409/// assert!(t.contains(b"foo").unwrap());
410/// assert_eq!(t.get(b"foo").unwrap().unwrap(), b"bar".to_vec());
411/// t.remove(b"foo").unwrap();
412/// assert!(!t.contains(b"foo").unwrap());
413/// ```
414pub struct TrieDBMut<'a, L>
415where
416	L: TrieLayout,
417{
418	storage: NodeStorage<TrieHash<L>>,
419	db: &'a mut dyn HashDB<L::Hash, DBValue>,
420	root: &'a mut TrieHash<L>,
421	root_handle: NodeHandle<TrieHash<L>>,
422	death_row: HashSet<(TrieHash<L>, (BackingByteVec, Option<u8>))>,
423	/// The number of hash operations this trie has performed.
424	/// Note that none are performed until changes are committed.
425	hash_count: usize,
426}
427
428impl<'a, L> TrieDBMut<'a, L>
429where
430	L: TrieLayout,
431{
432	/// Create a new trie with backing database `db` and empty `root`.
433	pub fn new(db: &'a mut dyn HashDB<L::Hash, DBValue>, root: &'a mut TrieHash<L>) -> Self {
434		*root = L::Codec::hashed_null_node();
435		let root_handle = NodeHandle::Hash(L::Codec::hashed_null_node());
436
437		TrieDBMut {
438			storage: NodeStorage::empty(),
439			db,
440			root,
441			root_handle,
442			death_row: HashSet::new(),
443			hash_count: 0,
444		}
445	}
446
447	/// Create a new trie with the backing database `db` and `root.
448	/// Returns an error if `root` does not exist.
449	pub fn from_existing(
450		db: &'a mut dyn HashDB<L::Hash, DBValue>,
451		root: &'a mut TrieHash<L>,
452	) -> Result<Self, TrieHash<L>, CError<L>> {
453		if !db.contains(root, EMPTY_PREFIX) {
454			return Err(Box::new(TrieError::InvalidStateRoot(*root)));
455		}
456
457		let root_handle = NodeHandle::Hash(*root);
458		Ok(TrieDBMut {
459			storage: NodeStorage::empty(),
460			db,
461			root,
462			root_handle,
463			death_row: HashSet::new(),
464			hash_count: 0,
465		})
466	}
467	/// Get the backing database.
468	pub fn db(&self) -> &dyn HashDB<L::Hash, DBValue> {
469		self.db
470	}
471
472	/// Get the backing database mutably.
473	pub fn db_mut(&mut self) -> &mut dyn HashDB<L::Hash, DBValue> {
474		self.db
475	}
476
477	// Cache a node by hash.
478	fn cache(
479		&mut self,
480		hash: TrieHash<L>,
481		key: Prefix,
482	) -> Result<StorageHandle, TrieHash<L>, CError<L>> {
483		let node_encoded = self.db.get(&hash, key)
484			.ok_or_else(|| Box::new(TrieError::IncompleteDatabase(hash)))?;
485		let node = Node::from_encoded::<L::Codec, L::Hash>(
486			hash,
487			&node_encoded,
488			&*self.db,
489			&mut self.storage
490		)?;
491		Ok(self.storage.alloc(Stored::Cached(node, hash)))
492	}
493
494	// Inspect a node, choosing either to replace, restore, or delete it.
495	// If restored or replaced, returns the new node along with a flag of whether it was changed.
496	fn inspect<F>(
497		&mut self,
498		stored: Stored<TrieHash<L>>,
499		key: &mut NibbleFullKey,
500		inspector: F,
501	) -> Result<Option<(Stored<TrieHash<L>>, bool)>, TrieHash<L>, CError<L>>
502		where
503			F: FnOnce(
504				&mut Self,
505				Node<TrieHash<L>>,
506				&mut NibbleFullKey,
507			) -> Result<Action<TrieHash<L>>, TrieHash<L>, CError<L>>,
508	{
509		let current_key = key.clone();
510		Ok(match stored {
511			Stored::New(node) => match inspector(self, node, key)? {
512				Action::Restore(node) => Some((Stored::New(node), false)),
513				Action::Replace(node) => Some((Stored::New(node), true)),
514				Action::Delete => None,
515			},
516			Stored::Cached(node, hash) => match inspector(self, node, key)? {
517				Action::Restore(node) => Some((Stored::Cached(node, hash), false)),
518				Action::Replace(node) => {
519					self.death_row.insert((hash, current_key.left_owned()));
520					Some((Stored::New(node), true))
521				}
522				Action::Delete => {
523					self.death_row.insert((hash, current_key.left_owned()));
524					None
525				}
526			},
527		})
528	}
529
530	// Walk the trie, attempting to find the key's node.
531	fn lookup<'x, 'key>(
532		&'x self,
533		mut partial: NibbleSlice<'key>,
534		handle: &NodeHandle<TrieHash<L>>,
535	) -> Result<Option<DBValue>, TrieHash<L>, CError<L>>
536		where 'x: 'key
537	{
538		let mut handle = handle;
539		loop {
540			let (mid, child) = match *handle {
541				NodeHandle::Hash(ref hash) => return Lookup::<L, _> {
542					db: &self.db,
543					query: |v: &[u8]| v.to_vec(),
544					hash: *hash,
545				}.look_up(partial),
546				NodeHandle::InMemory(ref handle) => match self.storage[handle] {
547					Node::Empty => return Ok(None),
548					Node::Leaf(ref key, ref value) => {
549						if NibbleSlice::from_stored(key) == partial {
550							return Ok(Some(value.to_vec()));
551						} else {
552							return Ok(None);
553						}
554					},
555					Node::Extension(ref slice, ref child) => {
556						let slice = NibbleSlice::from_stored(slice);
557						if partial.starts_with(&slice) {
558							(slice.len(), child)
559						} else {
560							return Ok(None);
561						}
562					},
563					Node::Branch(ref children, ref value) => {
564						if partial.is_empty() {
565							return Ok(value.as_ref().map(|v| v.to_vec()));
566						} else {
567							let idx = partial.at(0);
568							match children[idx as usize].as_ref() {
569								Some(child) => (1, child),
570								None => return Ok(None),
571							}
572						}
573					},
574					Node::NibbledBranch(ref slice, ref children, ref value) => {
575						let slice = NibbleSlice::from_stored(slice);
576						if partial.is_empty() {
577							return Ok(value.as_ref().map(|v| v.to_vec()));
578						} else if partial.starts_with(&slice) {
579							let idx = partial.at(0);
580							match children[idx as usize].as_ref() {
581								Some(child) => (1 + slice.len(), child),
582								None => return Ok(None),
583							}
584						} else {
585							return Ok(None)
586						}
587					},
588				}
589			};
590
591			partial = partial.mid(mid);
592			handle = child;
593		}
594	}
595
596	/// Insert a key-value pair into the trie, creating new nodes if necessary.
597	fn insert_at(
598		&mut self,
599		handle: NodeHandle<TrieHash<L>>,
600		key: &mut NibbleFullKey,
601		value: DBValue,
602		old_val: &mut Option<DBValue>,
603	) -> Result<(StorageHandle, bool), TrieHash<L>, CError<L>> {
604		let h = match handle {
605			NodeHandle::InMemory(h) => h,
606			NodeHandle::Hash(h) => self.cache(h, key.left())?,
607		};
608		// cache then destroy for hash handle (handle being root in most case)
609		let stored = self.storage.destroy(h);
610		let (new_stored, changed) = self.inspect(stored, key, move |trie, stored, key| {
611			trie.insert_inspector(stored, key, value, old_val).map(|a| a.into_action())
612		})?.expect("Insertion never deletes.");
613
614		Ok((self.storage.alloc(new_stored), changed))
615	}
616
617	/// The insertion inspector.
618	fn insert_inspector(
619		&mut self,
620		node: Node<TrieHash<L>>,
621		key: &mut NibbleFullKey,
622		value: DBValue,
623		old_val: &mut Option<DBValue>,
624	) -> Result<InsertAction<TrieHash<L>>, TrieHash<L>, CError<L>> {
625		let partial = *key;
626
627		#[cfg(feature = "std")]
628		trace!(target: "trie", "augmented (partial: {:?}, value: {:?})", partial, ToHex(&value));
629
630		Ok(match node {
631			Node::Empty => {
632				#[cfg(feature = "std")]
633				trace!(target: "trie", "empty: COMPOSE");
634				InsertAction::Replace(Node::Leaf(partial.to_stored(), value))
635			},
636			Node::Branch(mut children, stored_value) => {
637				debug_assert!(L::USE_EXTENSION);
638				#[cfg(feature = "std")]
639				trace!(target: "trie", "branch: ROUTE,AUGMENT");
640
641				if partial.is_empty() {
642					let unchanged = stored_value.as_ref() == Some(&value);
643					let branch = Node::Branch(children, Some(value));
644					*old_val = stored_value;
645
646					match unchanged {
647						true => InsertAction::Restore(branch),
648						false => InsertAction::Replace(branch),
649					}
650				} else {
651					let idx = partial.at(0) as usize;
652					key.advance(1);
653					if let Some(child) = children[idx].take() {
654						// Original had something there. recurse down into it.
655						let (new_child, changed) = self.insert_at(child, key, value, old_val)?;
656						children[idx] = Some(new_child.into());
657						if !changed {
658							// The new node we composed didn't change.
659							// It means our branch is untouched too.
660							return Ok(InsertAction::Restore(Node::Branch(children, stored_value)));
661						}
662					} else {
663						// Original had nothing there. compose a leaf.
664						let leaf = self.storage.alloc(
665							Stored::New(Node::Leaf(key.to_stored(), value))
666						);
667						children[idx] = Some(leaf.into());
668					}
669
670					InsertAction::Replace(Node::Branch(children, stored_value))
671				}
672			},
673			Node::NibbledBranch(encoded, mut children, stored_value) => {
674				debug_assert!(!L::USE_EXTENSION);
675				#[cfg(feature = "std")]
676				trace!(target: "trie", "branch: ROUTE,AUGMENT");
677				let existing_key = NibbleSlice::from_stored(&encoded);
678
679				let common = partial.common_prefix(&existing_key);
680				if common == existing_key.len() && common == partial.len() {
681					let unchanged = stored_value.as_ref() == Some(&value);
682					let branch = Node::NibbledBranch(
683						existing_key.to_stored(),
684						children,
685						Some(value),
686					);
687					*old_val = stored_value;
688
689					match unchanged {
690						true => InsertAction::Restore(branch),
691						false => InsertAction::Replace(branch),
692					}
693				} else if common < existing_key.len() {
694					// insert a branch value in between
695					#[cfg(feature = "std")]
696					trace!(
697						target: "trie",
698						"partially-shared-prefix (exist={:?}; new={:?}; common={:?}):\
699							 AUGMENT-AT-END",
700						existing_key.len(),
701						partial.len(),
702						common,
703					);
704					let nbranch_partial = existing_key.mid(common + 1).to_stored();
705					let low = Node::NibbledBranch(nbranch_partial, children, stored_value);
706					let ix = existing_key.at(common);
707					let mut children = empty_children();
708					let alloc_storage = self.storage.alloc(Stored::New(low));
709
710
711					children[ix as usize] = Some(alloc_storage.into());
712
713					if partial.len() - common == 0 {
714						InsertAction::Replace(Node::NibbledBranch(
715							existing_key.to_stored_range(common),
716							children,
717							Some(value),
718							)
719						)
720					} else {
721						let ix = partial.at(common);
722						let stored_leaf = Node::Leaf(partial.mid(common + 1).to_stored(), value);
723						let leaf = self.storage.alloc(Stored::New(stored_leaf));
724
725						children[ix as usize] = Some(leaf.into());
726						InsertAction::Replace(Node::NibbledBranch(
727							existing_key.to_stored_range(common),
728							children,
729							None,
730							)
731						)
732
733					}
734
735				} else {
736					// Append after common == existing_key and partial > common
737					#[cfg(feature = "std")]
738					trace!(target: "trie", "branch: ROUTE,AUGMENT");
739					let idx = partial.at(common) as usize;
740					key.advance(common + 1);
741					if let Some(child) = children[idx].take() {
742						// Original had something there. recurse down into it.
743						let (new_child, changed) = self.insert_at(child, key, value, old_val)?;
744						children[idx] = Some(new_child.into());
745						if !changed {
746							// The new node we composed didn't change.
747							// It means our branch is untouched too.
748							let n_branch = Node::NibbledBranch(
749								existing_key.to_stored(),
750								children,
751								stored_value,
752							);
753							return Ok(InsertAction::Restore(n_branch));
754						}
755					} else {
756						// Original had nothing there. compose a leaf.
757						let leaf = self.storage.alloc(
758							Stored::New(Node::Leaf(key.to_stored(), value)),
759						);
760						children[idx] = Some(leaf.into());
761					}
762					InsertAction::Replace(Node::NibbledBranch(
763						existing_key.to_stored(),
764						children,
765						stored_value,
766					))
767				}
768			},
769			Node::Leaf(encoded, stored_value) => {
770				let existing_key = NibbleSlice::from_stored(&encoded);
771				let common = partial.common_prefix(&existing_key);
772				if common == existing_key.len() && common == partial.len() {
773					#[cfg(feature = "std")]
774					trace!(target: "trie", "equivalent-leaf: REPLACE");
775					// equivalent leaf.
776					let unchanged = stored_value == value;
777					*old_val = Some(stored_value);
778
779					match unchanged {
780						// unchanged. restore
781						true => InsertAction::Restore(Node::Leaf(encoded.clone(), value)),
782						false => InsertAction::Replace(Node::Leaf(encoded.clone(), value)),
783					}
784				} else if (L::USE_EXTENSION && common == 0)
785					|| (!L::USE_EXTENSION && common < existing_key.len()) {
786					#[cfg(feature = "std")]
787					trace!(
788						target: "trie",
789						"lesser-common-prefix, not-both-empty (exist={:?}; new={:?}):\
790							TRANSMUTE,AUGMENT",
791						existing_key.len(),
792						partial.len(),
793					);
794
795					// one of us isn't empty: transmute to branch here
796					let mut children = empty_children();
797					let branch = if L::USE_EXTENSION && existing_key.is_empty() {
798						// always replace since branch isn't leaf.
799						Node::Branch(children, Some(stored_value))
800					} else {
801						let idx = existing_key.at(common) as usize;
802						let new_leaf = Node::Leaf(
803							existing_key.mid(common + 1).to_stored(),
804							stored_value,
805						);
806						children[idx] = Some(self.storage.alloc(Stored::New(new_leaf)).into());
807
808						if L::USE_EXTENSION {
809							Node::Branch(children, None)
810						} else {
811							Node::NibbledBranch(partial.to_stored_range(common), children, None)
812						}
813					};
814
815					// always replace because whatever we get out here
816					// is not the branch we started with.
817					let branch_action = self.insert_inspector(branch, key, value, old_val)?
818						.unwrap_node();
819					InsertAction::Replace(branch_action)
820				} else if !L::USE_EXTENSION {
821					#[cfg(feature = "std")]
822					trace!(target: "trie", "complete-prefix (common={:?}): AUGMENT-AT-END", common);
823
824					// fully-shared prefix for an extension.
825					// make a stub branch
826					let branch = Node::NibbledBranch(
827						existing_key.to_stored(),
828						empty_children(),
829						Some(stored_value),
830					);
831					// augment the new branch.
832					let branch = self.insert_inspector(branch, key, value, old_val)?
833						.unwrap_node();
834
835					InsertAction::Replace(branch)
836
837				} else if common == existing_key.len() {
838					debug_assert!(L::USE_EXTENSION);
839					#[cfg(feature = "std")]
840					trace!(target: "trie", "complete-prefix (common={:?}): AUGMENT-AT-END", common);
841
842					// fully-shared prefix for an extension.
843					// make a stub branch and an extension.
844					let branch = Node::Branch(empty_children(), Some(stored_value));
845					// augment the new branch.
846					key.advance(common);
847					let branch = self.insert_inspector(branch, key, value, old_val)?.unwrap_node();
848
849					// always replace since we took a leaf and made an extension.
850					let branch_handle = self.storage.alloc(Stored::New(branch)).into();
851					InsertAction::Replace(Node::Extension(existing_key.to_stored(), branch_handle))
852				} else {
853					debug_assert!(L::USE_EXTENSION);
854					#[cfg(feature = "std")]
855					trace!(
856						target: "trie",
857						"partially-shared-prefix (exist={:?}; new={:?}; common={:?}):\
858							 AUGMENT-AT-END",
859						existing_key.len(),
860						partial.len(),
861						common,
862					);
863
864					// partially-shared prefix for an extension.
865					// start by making a leaf.
866					let low = Node::Leaf(existing_key.mid(common).to_stored(), stored_value);
867
868					// augment it. this will result in the Leaf -> common == 0 routine,
869					// which creates a branch.
870					key.advance(common);
871					let augmented_low = self.insert_inspector(low, key, value, old_val)?
872						.unwrap_node();
873					// make an extension using it. this is a replacement.
874					InsertAction::Replace(Node::Extension(
875						existing_key.to_stored_range(common),
876						self.storage.alloc(Stored::New(augmented_low)).into()
877					))
878				}
879			},
880			Node::Extension(encoded, child_branch) => {
881				debug_assert!(L::USE_EXTENSION);
882				let existing_key = NibbleSlice::from_stored(&encoded);
883				let common = partial.common_prefix(&existing_key);
884				if common == 0 {
885					#[cfg(feature = "std")]
886					trace!(
887						target: "trie",
888						"no-common-prefix, not-both-empty (exist={:?}; new={:?}):\
889							 TRANSMUTE,AUGMENT",
890						existing_key.len(),
891						partial.len(),
892					);
893
894					// partial isn't empty: make a branch here
895					// extensions may not have empty partial keys.
896					assert!(!existing_key.is_empty());
897					let idx = existing_key.at(0) as usize;
898
899					let mut children = empty_children();
900					children[idx] = if existing_key.len() == 1 {
901						// direct extension, just replace.
902						Some(child_branch)
903					} else {
904						// more work required after branching.
905						let ext = Node::Extension(existing_key.mid(1).to_stored(), child_branch);
906						Some(self.storage.alloc(Stored::New(ext)).into())
907					};
908
909					// continue inserting.
910					let branch_action = self.insert_inspector(
911						Node::Branch(children, None),
912						key,
913						value,
914						old_val,
915					)?.unwrap_node();
916					InsertAction::Replace(branch_action)
917				} else if common == existing_key.len() {
918					#[cfg(feature = "std")]
919					trace!(target: "trie", "complete-prefix (common={:?}): AUGMENT-AT-END", common);
920
921					// fully-shared prefix.
922
923					// insert into the child node.
924					key.advance(common);
925					let (new_child, changed) = self.insert_at(child_branch, key, value, old_val)?;
926					let new_ext = Node::Extension(existing_key.to_stored(), new_child.into());
927
928					// if the child branch wasn't changed, meaning this extension remains the same.
929					match changed {
930						true => InsertAction::Replace(new_ext),
931						false => InsertAction::Restore(new_ext),
932					}
933				} else {
934					#[cfg(feature = "std")]
935					trace!(
936						target: "trie",
937						"partially-shared-prefix (exist={:?}; new={:?}; common={:?}):\
938							 AUGMENT-AT-END",
939						existing_key.len(),
940						partial.len(),
941						common,
942					);
943
944					// partially-shared.
945					let low = Node::Extension(existing_key.mid(common).to_stored(), child_branch);
946					// augment the extension. this will take the common == 0 path,
947					// creating a branch.
948					key.advance(common);
949					let augmented_low = self.insert_inspector(low, key, value, old_val)?
950						.unwrap_node();
951
952					// always replace, since this extension is not the one we started with.
953					// this is known because the partial key is only the common prefix.
954					InsertAction::Replace(Node::Extension(
955						existing_key.to_stored_range(common),
956						self.storage.alloc(Stored::New(augmented_low)).into()
957					))
958				}
959			},
960		})
961	}
962
963	/// Removes a node from the trie based on key.
964	fn remove_at(
965		&mut self,
966		handle: NodeHandle<TrieHash<L>>,
967		key: &mut NibbleFullKey,
968		old_val: &mut Option<DBValue>,
969	) -> Result<Option<(StorageHandle, bool)>, TrieHash<L>, CError<L>> {
970		let stored = match handle {
971			NodeHandle::InMemory(h) => self.storage.destroy(h),
972			NodeHandle::Hash(h) => {
973				let handle = self.cache(h, key.left())?;
974				self.storage.destroy(handle)
975			}
976		};
977
978		let opt = self.inspect(
979			stored,
980			key,
981			move |trie, node, key| trie.remove_inspector(node, key, old_val),
982		)?;
983
984		Ok(opt.map(|(new, changed)| (self.storage.alloc(new), changed)))
985	}
986
987	/// The removal inspector.
988	fn remove_inspector(
989		&mut self,
990		node: Node<TrieHash<L>>,
991		key: &mut NibbleFullKey,
992		old_val: &mut Option<DBValue>,
993	) -> Result<Action<TrieHash<L>>, TrieHash<L>, CError<L>> {
994		let partial = *key;
995		Ok(match (node, partial.is_empty()) {
996			(Node::Empty, _) => Action::Delete,
997			(Node::Branch(c, None), true) => Action::Restore(Node::Branch(c, None)),
998			(Node::NibbledBranch(n, c, None), true) =>
999				Action::Restore(Node::NibbledBranch(n, c, None)),
1000			(Node::Branch(children, Some(val)), true) => {
1001				*old_val = Some(val);
1002				// always replace since we took the value out.
1003				Action::Replace(self.fix(Node::Branch(children, None), key.clone())?)
1004			},
1005			(Node::NibbledBranch(n, children, Some(val)), true) => {
1006				*old_val = Some(val);
1007				// always replace since we took the value out.
1008				Action::Replace(self.fix(Node::NibbledBranch(n, children, None), key.clone())?)
1009			},
1010			(Node::Branch(mut children, value), false) => {
1011				let idx = partial.at(0) as usize;
1012				if let Some(child) = children[idx].take() {
1013					#[cfg(feature = "std")]
1014					trace!(
1015						target: "trie",
1016						"removing value out of branch child, partial={:?}",
1017						partial,
1018					);
1019					let prefix = *key;
1020					key.advance(1);
1021					match self.remove_at(child, key, old_val)? {
1022						Some((new, changed)) => {
1023							children[idx] = Some(new.into());
1024							let branch = Node::Branch(children, value);
1025							match changed {
1026								// child was changed, so we were too.
1027								true => Action::Replace(branch),
1028								// unchanged, so we are too.
1029								false => Action::Restore(branch),
1030							}
1031						}
1032						None => {
1033							// the child we took was deleted.
1034							// the node may need fixing.
1035							#[cfg(feature = "std")]
1036							trace!(target: "trie", "branch child deleted, partial={:?}", partial);
1037							Action::Replace(self.fix(Node::Branch(children, value), prefix)?)
1038						}
1039					}
1040				} else {
1041					// no change needed.
1042					Action::Restore(Node::Branch(children, value))
1043				}
1044			},
1045			(Node::NibbledBranch(encoded, mut children, value), false) => {
1046				let (common, existing_length) = {
1047					let existing_key = NibbleSlice::from_stored(&encoded);
1048					(existing_key.common_prefix(&partial), existing_key.len())
1049				};
1050				if common == existing_length && common == partial.len() {
1051
1052					// replace val
1053					if let Some(val) = value {
1054						*old_val = Some(val);
1055
1056						let f = self.fix(Node::NibbledBranch(encoded, children, None), key.clone());
1057						Action::Replace(f?)
1058					} else {
1059						Action::Restore(Node::NibbledBranch(encoded, children, None))
1060					}
1061				} else if common < existing_length {
1062					// partway through an extension -- nothing to do here.
1063					Action::Restore(Node::NibbledBranch(encoded, children, value))
1064				} else {
1065					// common == existing_length && common < partial.len() : check children
1066					let idx = partial.at(common) as usize;
1067
1068					if let Some(child) = children[idx].take() {
1069						#[cfg(feature = "std")]
1070						trace!(
1071							target: "trie",
1072							"removing value out of branch child, partial={:?}",
1073							partial,
1074						);
1075						let prefix = *key;
1076						key.advance(common + 1);
1077						match self.remove_at(child, key, old_val)? {
1078							Some((new, changed)) => {
1079								children[idx] = Some(new.into());
1080								let branch = Node::NibbledBranch(encoded, children, value);
1081								match changed {
1082									// child was changed, so we were too.
1083									true => Action::Replace(branch),
1084									// unchanged, so we are too.
1085									false => Action::Restore(branch),
1086								}
1087							},
1088							None => {
1089								// the child we took was deleted.
1090								// the node may need fixing.
1091								#[cfg(feature = "std")]
1092								trace!(
1093									target: "trie",
1094									"branch child deleted, partial={:?}",
1095									partial,
1096								);
1097								Action::Replace(
1098									self.fix(Node::NibbledBranch(encoded, children, value), prefix)?
1099								)
1100							},
1101						}
1102					} else {
1103						// no change needed.
1104						Action::Restore(Node::NibbledBranch(encoded, children, value))
1105					}
1106				}
1107			},
1108			(Node::Leaf(encoded, value), _) => {
1109				if NibbleSlice::from_stored(&encoded) == partial {
1110					// this is the node we were looking for. Let's delete it.
1111					*old_val = Some(value);
1112					Action::Delete
1113				} else {
1114					// leaf the node alone.
1115					#[cfg(feature = "std")]
1116					trace!(
1117						target: "trie",
1118						"restoring leaf wrong partial, partial={:?}, existing={:?}",
1119						partial,
1120						NibbleSlice::from_stored(&encoded),
1121					);
1122					Action::Restore(Node::Leaf(encoded, value))
1123				}
1124			},
1125			(Node::Extension(encoded, child_branch), _) => {
1126				let (common, existing_length) = {
1127					let existing_key = NibbleSlice::from_stored(&encoded);
1128					(existing_key.common_prefix(&partial), existing_key.len())
1129				};
1130				if common == existing_length {
1131					// try to remove from the child branch.
1132					#[cfg(feature = "std")]
1133					trace!(target: "trie", "removing from extension child, partial={:?}", partial);
1134					let prefix = *key;
1135					key.advance(common);
1136					match self.remove_at(child_branch, key, old_val)? {
1137						Some((new_child, changed)) => {
1138							let new_child = new_child.into();
1139
1140							// if the child branch was unchanged, then the extension is too.
1141							// otherwise, this extension may need fixing.
1142							match changed {
1143								true => Action::Replace(
1144									self.fix(Node::Extension(encoded, new_child), prefix)?
1145								),
1146								false => Action::Restore(Node::Extension(encoded, new_child)),
1147							}
1148						}
1149						None => {
1150							// the whole branch got deleted.
1151							// that means that this extension is useless.
1152							Action::Delete
1153						}
1154					}
1155				} else {
1156					// partway through an extension -- nothing to do here.
1157					Action::Restore(Node::Extension(encoded, child_branch))
1158				}
1159			},
1160		})
1161	}
1162
1163	/// Given a node which may be in an _invalid state_, fix it such that it is then in a valid
1164	/// state.
1165	///
1166	/// _invalid state_ means:
1167	/// - Branch node where there is only a single entry;
1168	/// - Extension node followed by anything other than a Branch node.
1169	fn fix(
1170		&mut self,
1171		node: Node<TrieHash<L>>,
1172		key: NibbleSlice,
1173	) -> Result<Node<TrieHash<L>>, TrieHash<L>, CError<L>> {
1174		match node {
1175			Node::Branch(mut children, value) => {
1176				// if only a single value, transmute to leaf/extension and feed through fixed.
1177				#[cfg_attr(feature = "std", derive(Debug))]
1178				enum UsedIndex {
1179					None,
1180					One(u8),
1181					Many,
1182				};
1183				let mut used_index = UsedIndex::None;
1184				for i in 0..16 {
1185					match (children[i].is_none(), &used_index) {
1186						(false, &UsedIndex::None) => used_index = UsedIndex::One(i as u8),
1187						(false, &UsedIndex::One(_)) => {
1188							used_index = UsedIndex::Many;
1189							break;
1190						}
1191						_ => continue,
1192					}
1193				}
1194
1195				match (used_index, value) {
1196					(UsedIndex::None, None) =>
1197						panic!("Branch with no subvalues. Something went wrong."),
1198					(UsedIndex::One(a), None) => {
1199						// only one onward node. make an extension.
1200
1201						let new_partial = NibbleSlice::new_offset(&[a], 1).to_stored();
1202						let child = children[a as usize].take()
1203							.expect("used_index only set if occupied; qed");
1204						let new_node = Node::Extension(new_partial, child);
1205						self.fix(new_node, key)
1206					}
1207					(UsedIndex::None, Some(value)) => {
1208						// make a leaf.
1209						#[cfg(feature = "std")]
1210						trace!(target: "trie", "fixing: branch -> leaf");
1211						Ok(Node::Leaf(NibbleSlice::new(&[]).to_stored(), value))
1212					}
1213					(_, value) => {
1214						// all is well.
1215						#[cfg(feature = "std")]
1216						trace!(target: "trie", "fixing: restoring branch");
1217						Ok(Node::Branch(children, value))
1218					}
1219				}
1220			},
1221			Node::NibbledBranch(enc_nibble, mut children, value) => {
1222				// if only a single value, transmute to leaf/extension and feed through fixed.
1223				#[cfg_attr(feature = "std", derive(Debug))]
1224				enum UsedIndex {
1225					None,
1226					One(u8),
1227					Many,
1228				};
1229				let mut used_index = UsedIndex::None;
1230				for i in 0..16 {
1231					match (children[i].is_none(), &used_index) {
1232						(false, &UsedIndex::None) => used_index = UsedIndex::One(i as u8),
1233						(false, &UsedIndex::One(_)) => {
1234							used_index = UsedIndex::Many;
1235							break;
1236						}
1237						_ => continue,
1238					}
1239				}
1240
1241				match (used_index, value) {
1242					(UsedIndex::None, None) =>
1243						panic!("Branch with no subvalues. Something went wrong."),
1244					(UsedIndex::One(a), None) => {
1245						// only one onward node. use child instead
1246						let child = children[a as usize].take()
1247							.expect("used_index only set if occupied; qed");
1248						let mut key2 = key.clone();
1249						key2.advance((enc_nibble.1.len() * nibble_ops::NIBBLE_PER_BYTE) - enc_nibble.0);
1250						let (start, alloc_start, prefix_end) = match key2.left() {
1251							(start, None) => (start, None, Some(nibble_ops::push_at_left(0, a, 0))),
1252							(start, Some(v)) => {
1253								let mut so: BackingByteVec = start.into();
1254								so.push(nibble_ops::pad_left(v) | a);
1255								(start, Some(so), None)
1256							},
1257						};
1258						let child_prefix = (alloc_start.as_ref().map(|start| &start[..]).unwrap_or(start), prefix_end);
1259						let stored = match child {
1260							NodeHandle::InMemory(h) => self.storage.destroy(h),
1261							NodeHandle::Hash(h) => {
1262								let handle = self.cache(h, child_prefix)?;
1263								self.storage.destroy(handle)
1264							}
1265						};
1266						let child_node = match stored {
1267							Stored::New(node) => node,
1268							Stored::Cached(node, hash) => {
1269								self.death_row.insert((
1270									hash,
1271									(child_prefix.0[..].into(), child_prefix.1),
1272								));
1273								node
1274							},
1275						};
1276						match child_node {
1277							Node::Leaf(sub_partial, value) => {
1278								let mut enc_nibble = enc_nibble;
1279								combine_key(
1280									&mut enc_nibble,
1281									(nibble_ops::NIBBLE_PER_BYTE - 1, &[a][..]),
1282								);
1283								combine_key(
1284									&mut enc_nibble,
1285									(sub_partial.0, &sub_partial.1[..]),
1286								);
1287								Ok(Node::Leaf(enc_nibble, value))
1288							},
1289							Node::NibbledBranch(sub_partial, ch_children, ch_value) => {
1290								let mut enc_nibble = enc_nibble;
1291								combine_key(
1292									&mut enc_nibble,
1293									(nibble_ops::NIBBLE_PER_BYTE - 1, &[a][..]),
1294								);
1295								combine_key(
1296									&mut enc_nibble,
1297									(sub_partial.0, &sub_partial.1[..]),
1298								);
1299								Ok(Node::NibbledBranch(enc_nibble, ch_children, ch_value))
1300							},
1301							_ => unreachable!(),
1302						}
1303					},
1304					(UsedIndex::None, Some(value)) => {
1305						// make a leaf.
1306						#[cfg(feature = "std")]
1307						trace!(target: "trie", "fixing: branch -> leaf");
1308						Ok(Node::Leaf(enc_nibble, value))
1309					},
1310					(_, value) => {
1311						// all is well.
1312						#[cfg(feature = "std")]
1313						trace!(target: "trie", "fixing: restoring branch");
1314						Ok(Node::NibbledBranch(enc_nibble, children, value))
1315					},
1316				}
1317			},
1318			Node::Extension(partial, child) => {
1319				// We could advance key, but this code can also be called
1320				// recursively, so there might be some prefix from branch.
1321				let last = partial.1[partial.1.len() - 1] & (255 >> 4);
1322				let mut key2 = key.clone();
1323				key2.advance((partial.1.len() * nibble_ops::NIBBLE_PER_BYTE) - partial.0 - 1);
1324				let (start, alloc_start, prefix_end) = match key2.left() {
1325					(start, None) => (start, None, Some(nibble_ops::push_at_left(0, last, 0))),
1326					(start, Some(v)) => {
1327						let mut so: BackingByteVec = start.into();
1328						// Complete last byte with `last`.
1329						so.push(nibble_ops::pad_left(v) | last);
1330						(start, Some(so), None)
1331					},
1332				};
1333				let child_prefix = (alloc_start.as_ref().map(|start| &start[..]).unwrap_or(start), prefix_end);
1334
1335				let stored = match child {
1336					NodeHandle::InMemory(h) => self.storage.destroy(h),
1337					NodeHandle::Hash(h) => {
1338						let handle = self.cache(h, child_prefix)?;
1339						self.storage.destroy(handle)
1340					}
1341				};
1342
1343				let (child_node, maybe_hash) = match stored {
1344					Stored::New(node) => (node, None),
1345					Stored::Cached(node, hash) => (node, Some(hash))
1346				};
1347
1348				match child_node {
1349					Node::Extension(sub_partial, sub_child) => {
1350						// combine with node below.
1351						if let Some(hash) = maybe_hash {
1352							// delete the cached child since we are going to replace it.
1353							self.death_row.insert(
1354								(hash, (child_prefix.0[..].into(), child_prefix.1)),
1355							);
1356						}
1357						// subpartial
1358						let mut partial = partial;
1359						combine_key(&mut partial, (sub_partial.0, &sub_partial.1[..]));
1360						#[cfg(feature = "std")]
1361						trace!(
1362							target: "trie",
1363							"fixing: extension combination. new_partial={:?}",
1364							partial,
1365						);
1366						self.fix(Node::Extension(partial, sub_child), key)
1367					}
1368					Node::Leaf(sub_partial, value) => {
1369						// combine with node below.
1370						if let Some(hash) = maybe_hash {
1371							// delete the cached child since we are going to replace it.
1372							self.death_row.insert((hash, (child_prefix.0[..].into(), child_prefix.1)));
1373						}
1374						// subpartial oly
1375						let mut partial = partial;
1376						combine_key(&mut partial, (sub_partial.0, &sub_partial.1[..]));
1377						#[cfg(feature = "std")]
1378						trace!(
1379							target: "trie",
1380							"fixing: extension -> leaf. new_partial={:?}",
1381							partial,
1382						);
1383						Ok(Node::Leaf(partial, value))
1384					}
1385					child_node => {
1386						#[cfg(feature = "std")]
1387						trace!(target: "trie", "fixing: restoring extension");
1388
1389						// reallocate the child node.
1390						let stored = if let Some(hash) = maybe_hash {
1391							Stored::Cached(child_node, hash)
1392						} else {
1393							Stored::New(child_node)
1394						};
1395
1396						Ok(Node::Extension(partial, self.storage.alloc(stored).into()))
1397					}
1398				}
1399			},
1400			other => Ok(other), // only ext and branch need fixing.
1401		}
1402	}
1403
1404	/// Commit the in-memory changes to disk, freeing their storage and
1405	/// updating the state root.
1406	pub fn commit(&mut self) {
1407		#[cfg(feature = "std")]
1408		trace!(target: "trie", "Committing trie changes to db.");
1409
1410		// always kill all the nodes on death row.
1411		#[cfg(feature = "std")]
1412		trace!(target: "trie", "{:?} nodes to remove from db", self.death_row.len());
1413		for (hash, prefix) in self.death_row.drain() {
1414			self.db.remove(&hash, (&prefix.0[..], prefix.1));
1415		}
1416
1417		let handle = match self.root_handle() {
1418			NodeHandle::Hash(_) => return, // no changes necessary.
1419			NodeHandle::InMemory(h) => h,
1420		};
1421
1422		match self.storage.destroy(handle) {
1423			Stored::New(node) => {
1424				let mut k = NibbleVec::new();
1425				let encoded_root = node.into_encoded::<_, L::Codec, L::Hash>(
1426					|child, o_slice, o_index| {
1427						let mov = k.append_optional_slice_and_nibble(o_slice, o_index);
1428						let cr = self.commit_child(child, &mut k);
1429						k.drop_lasts(mov);
1430						cr
1431					}
1432				);
1433				#[cfg(feature = "std")]
1434				trace!(target: "trie", "encoded root node: {:#x?}", &encoded_root[..]);
1435				*self.root = self.db.insert(EMPTY_PREFIX, &encoded_root[..]);
1436				self.hash_count += 1;
1437
1438				self.root_handle = NodeHandle::Hash(*self.root);
1439			}
1440			Stored::Cached(node, hash) => {
1441				// probably won't happen, but update the root and move on.
1442				*self.root = hash;
1443				self.root_handle = NodeHandle::InMemory(
1444					self.storage.alloc(Stored::Cached(node, hash)),
1445				);
1446			}
1447		}
1448	}
1449
1450	/// Commit a node by hashing it and writing it to the db. Returns a
1451	/// `ChildReference` which in most cases carries a normal hash but for the
1452	/// case where we can fit the actual data in the `Hasher`s output type, we
1453	/// store the data inline. This function is used as the callback to the
1454	/// `into_encoded` method of `Node`.
1455	fn commit_child(
1456		&mut self,
1457		handle: NodeHandle<TrieHash<L>>,
1458		prefix: &mut NibbleVec,
1459	) -> ChildReference<TrieHash<L>> {
1460		match handle {
1461			NodeHandle::Hash(hash) => ChildReference::Hash(hash),
1462			NodeHandle::InMemory(storage_handle) => {
1463				match self.storage.destroy(storage_handle) {
1464					Stored::Cached(_, hash) => ChildReference::Hash(hash),
1465					Stored::New(node) => {
1466						let encoded = {
1467							let commit_child = |
1468								node_handle,
1469								o_slice: Option<&NibbleSlice>,
1470								o_index: Option<u8>
1471							| {
1472								let mov = prefix.append_optional_slice_and_nibble(o_slice, o_index);
1473								let cr = self.commit_child(node_handle, prefix);
1474								prefix.drop_lasts(mov);
1475								cr
1476							};
1477							node.into_encoded::<_, L::Codec, L::Hash>(commit_child)
1478						};
1479						if encoded.len() >= L::Hash::LENGTH {
1480							let hash = self.db.insert(prefix.as_prefix(), &encoded[..]);
1481							self.hash_count +=1;
1482							ChildReference::Hash(hash)
1483						} else {
1484							// it's a small value, so we cram it into a `TrieHash<L>`
1485							// and tag with length
1486							let mut h = <TrieHash<L>>::default();
1487							let len = encoded.len();
1488							h.as_mut()[..len].copy_from_slice(&encoded[..len]);
1489							ChildReference::Inline(h, len)
1490						}
1491					}
1492				}
1493			}
1494		}
1495	}
1496
1497	// a hack to get the root node's handle
1498	fn root_handle(&self) -> NodeHandle<TrieHash<L>> {
1499		match self.root_handle {
1500			NodeHandle::Hash(h) => NodeHandle::Hash(h),
1501			NodeHandle::InMemory(StorageHandle(x)) => NodeHandle::InMemory(StorageHandle(x)),
1502		}
1503	}
1504}
1505
1506
1507impl<'a, L> TrieMut<L> for TrieDBMut<'a, L>
1508where
1509	L: TrieLayout,
1510{
1511	fn root(&mut self) -> &TrieHash<L> {
1512		self.commit();
1513		self.root
1514	}
1515
1516	fn is_empty(&self) -> bool {
1517		match self.root_handle {
1518			NodeHandle::Hash(h) => h == L::Codec::hashed_null_node(),
1519			NodeHandle::InMemory(ref h) => match self.storage[h] {
1520				Node::Empty => true,
1521				_ => false,
1522			}
1523		}
1524	}
1525
1526	fn get<'x, 'key>(&'x self, key: &'key [u8]) -> Result<Option<DBValue>, TrieHash<L>, CError<L>>
1527		where 'x: 'key
1528	{
1529		self.lookup(NibbleSlice::new(key), &self.root_handle)
1530	}
1531
1532	fn insert(
1533		&mut self,
1534		key: &[u8],
1535		value: &[u8],
1536	) -> Result<Option<DBValue>, TrieHash<L>, CError<L>> {
1537		if !L::ALLOW_EMPTY && value.is_empty() { return self.remove(key) }
1538
1539		let mut old_val = None;
1540
1541		#[cfg(feature = "std")]
1542		trace!(target: "trie", "insert: key={:#x?}, value={:?}", key, ToHex(&value));
1543
1544		let root_handle = self.root_handle();
1545		let (new_handle, _changed) = self.insert_at(
1546			root_handle,
1547			&mut NibbleSlice::new(key),
1548			value.to_vec(),
1549			&mut old_val,
1550		)?;
1551
1552		#[cfg(feature = "std")]
1553		trace!(target: "trie", "insert: altered trie={}", _changed);
1554		self.root_handle = NodeHandle::InMemory(new_handle);
1555
1556		Ok(old_val)
1557	}
1558
1559	fn remove(&mut self, key: &[u8]) -> Result<Option<DBValue>, TrieHash<L>, CError<L>> {
1560		#[cfg(feature = "std")]
1561		trace!(target: "trie", "remove: key={:#x?}", key);
1562
1563		let root_handle = self.root_handle();
1564		let mut key = NibbleSlice::new(key);
1565		let mut old_val = None;
1566
1567		match self.remove_at(root_handle, &mut key, &mut old_val)? {
1568			Some((handle, _changed)) => {
1569				#[cfg(feature = "std")]
1570				trace!(target: "trie", "remove: altered trie={}", _changed);
1571				self.root_handle = NodeHandle::InMemory(handle);
1572			}
1573			None => {
1574				#[cfg(feature = "std")]
1575				trace!(target: "trie", "remove: obliterated trie");
1576				self.root_handle = NodeHandle::Hash(L::Codec::hashed_null_node());
1577				*self.root = L::Codec::hashed_null_node();
1578			}
1579		}
1580
1581		Ok(old_val)
1582	}
1583}
1584
1585impl<'a, L> Drop for TrieDBMut<'a, L>
1586where
1587	L: TrieLayout,
1588{
1589	fn drop(&mut self) {
1590		self.commit();
1591	}
1592}
1593
1594/// combine two NodeKeys
1595fn combine_key(start: &mut NodeKey, end: (usize, &[u8])) {
1596	debug_assert!(start.0 < nibble_ops::NIBBLE_PER_BYTE);
1597	debug_assert!(end.0 < nibble_ops::NIBBLE_PER_BYTE);
1598	let final_offset = (start.0 + end.0) % nibble_ops::NIBBLE_PER_BYTE;
1599	let _shifted = nibble_ops::shift_key(start, final_offset);
1600	let st = if end.0 > 0 {
1601		let sl = start.1.len();
1602		start.1[sl - 1] |= nibble_ops::pad_right(end.1[0]);
1603		1
1604	} else {
1605		0
1606	};
1607	(st..end.1.len()).for_each(|i| start.1.push(end.1[i]));
1608}
1609
1610#[cfg(test)]
1611mod tests {
1612	use crate::nibble::BackingByteVec;
1613
1614	#[test]
1615	fn combine_test() {
1616		let a: BackingByteVec = [0x12, 0x34][..].into();
1617		let b: &[u8] = [0x56, 0x78][..].into();
1618		let test_comb = |a: (_, &BackingByteVec), b, c| {
1619			let mut a = (a.0, a.1.clone());
1620			super::combine_key(&mut a, b);
1621			assert_eq!((a.0, &a.1[..]), c);
1622		};
1623		test_comb((0, &a), (0, &b), (0, &[0x12, 0x34, 0x56, 0x78][..]));
1624		test_comb((1, &a), (0, &b), (1, &[0x12, 0x34, 0x56, 0x78][..]));
1625		test_comb((0, &a), (1, &b), (1, &[0x01, 0x23, 0x46, 0x78][..]));
1626		test_comb((1, &a), (1, &b), (0, &[0x23, 0x46, 0x78][..]));
1627	}
1628
1629	#[test]
1630	fn nice_debug_for_node() {
1631		use super::Node;
1632		let e: Node<u32> = Node::Leaf((1, vec![1, 2, 3].into()), vec![4, 5, 6]);
1633		assert_eq!(format!("{:?}", e), "Leaf((1, 010203), 040506)");
1634	}
1635}