1use crate::{
16 fatdb::FatDBDoubleEndedIterator,
17 iterator::{TrieDBNodeDoubleEndedIterator, TrieDBRawIterator},
18 lookup::Lookup,
19 nibble::NibbleSlice,
20 node::{decode_hash, NodeHandle, OwnedNode},
21 rstd::boxed::Box,
22 CError, DBValue, MerkleValue, Query, Result, Trie, TrieAccess, TrieCache,
23 TrieDoubleEndedIterator, TrieError, TrieHash, TrieItem, TrieIterator, TrieKeyItem, TrieLayout,
24 TrieRecorder,
25};
26#[cfg(feature = "std")]
27use crate::{
28 nibble::NibbleVec,
29 node::Node,
30 rstd::{fmt, vec::Vec},
31};
32use hash_db::{HashDBRef, Prefix, EMPTY_PREFIX};
33
34pub struct TrieDBBuilder<'db, 'cache, L: TrieLayout> {
36 db: &'db dyn HashDBRef<L::Hash, DBValue>,
37 root: &'db TrieHash<L>,
38 cache: Option<&'cache mut dyn TrieCache<L::Codec>>,
39 recorder: Option<&'cache mut dyn TrieRecorder<TrieHash<L>>>,
40}
41
42impl<'db, 'cache, L: TrieLayout> TrieDBBuilder<'db, 'cache, L> {
43 #[inline]
48 pub fn new(db: &'db dyn HashDBRef<L::Hash, DBValue>, root: &'db TrieHash<L>) -> Self {
49 Self { db, root, cache: None, recorder: None }
50 }
51
52 #[inline]
54 pub fn with_cache(mut self, cache: &'cache mut dyn TrieCache<L::Codec>) -> Self {
55 self.cache = Some(cache);
56 self
57 }
58
59 #[inline]
61 pub fn with_optional_cache<'ocache: 'cache>(
62 mut self,
63 cache: Option<&'ocache mut dyn TrieCache<L::Codec>>,
64 ) -> Self {
65 self.cache = cache.map(|c| c as _);
67 self
68 }
69
70 #[inline]
72 pub fn with_recorder(mut self, recorder: &'cache mut dyn TrieRecorder<TrieHash<L>>) -> Self {
73 self.recorder = Some(recorder);
74 self
75 }
76
77 #[inline]
79 pub fn with_optional_recorder<'recorder: 'cache>(
80 mut self,
81 recorder: Option<&'recorder mut dyn TrieRecorder<TrieHash<L>>>,
82 ) -> Self {
83 self.recorder = recorder.map(|r| r as _);
85 self
86 }
87
88 #[inline]
90 pub fn build(self) -> TrieDB<'db, 'cache, L> {
91 TrieDB {
92 db: self.db,
93 root: self.root,
94 cache: self.cache.map(core::cell::RefCell::new),
95 recorder: self.recorder.map(core::cell::RefCell::new),
96 }
97 }
98}
99
100pub struct TrieDB<'db, 'cache, L>
123where
124 L: TrieLayout,
125{
126 db: &'db dyn HashDBRef<L::Hash, DBValue>,
127 root: &'db TrieHash<L>,
128 cache: Option<core::cell::RefCell<&'cache mut dyn TrieCache<L::Codec>>>,
129 recorder: Option<core::cell::RefCell<&'cache mut dyn TrieRecorder<TrieHash<L>>>>,
130}
131
132impl<'db, 'cache, L> TrieDB<'db, 'cache, L>
133where
134 L: TrieLayout,
135{
136 pub fn db(&'db self) -> &'db dyn HashDBRef<L::Hash, DBValue> {
138 self.db
139 }
140
141 pub fn into_double_ended_iter(
143 &'db self,
144 ) -> Result<TrieDBDoubleEndedIterator<'db, 'cache, L>, TrieHash<L>, CError<L>> {
145 TrieDBDoubleEndedIterator::new(&self)
146 }
147
148 pub fn into_node_double_ended_iter(
150 &'db self,
151 ) -> Result<TrieDBNodeDoubleEndedIterator<'db, 'cache, L>, TrieHash<L>, CError<L>> {
152 TrieDBNodeDoubleEndedIterator::new(&self)
153 }
154
155 pub fn into_key_double_ended_iter(
157 &'db self,
158 ) -> Result<TrieDBKeyDoubleEndedIterator<'db, 'cache, L>, TrieHash<L>, CError<L>> {
159 TrieDBKeyDoubleEndedIterator::new(&self)
160 }
161
162 pub fn into_fat_double_ended_iter(
164 &'db self,
165 ) -> Result<FatDBDoubleEndedIterator<'db, 'cache, L>, TrieHash<L>, CError<L>> {
166 FatDBDoubleEndedIterator::new(&self)
167 }
168
169 pub(crate) fn get_raw_or_lookup(
181 &self,
182 parent_hash: TrieHash<L>,
183 node_handle: NodeHandle,
184 partial_key: Prefix,
185 record_access: bool,
186 ) -> Result<(OwnedNode<DBValue>, Option<TrieHash<L>>), TrieHash<L>, CError<L>> {
187 let (node_hash, node_data) = match node_handle {
188 NodeHandle::Hash(data) => {
189 let node_hash = decode_hash::<L::Hash>(data)
190 .ok_or_else(|| Box::new(TrieError::InvalidHash(parent_hash, data.to_vec())))?;
191 let node_data = self.db.get(&node_hash, partial_key).ok_or_else(|| {
192 if partial_key == EMPTY_PREFIX {
193 Box::new(TrieError::InvalidStateRoot(node_hash))
194 } else {
195 Box::new(TrieError::IncompleteDatabase(node_hash))
196 }
197 })?;
198
199 (Some(node_hash), node_data)
200 },
201 NodeHandle::Inline(data) => (None, data.to_vec()),
202 };
203 let owned_node = OwnedNode::new::<L::Codec>(node_data)
204 .map_err(|e| Box::new(TrieError::DecoderError(node_hash.unwrap_or(parent_hash), e)))?;
205
206 if record_access {
207 if let Some((hash, recorder)) =
208 node_hash.as_ref().and_then(|h| self.recorder.as_ref().map(|r| (h, r)))
209 {
210 recorder.borrow_mut().record(TrieAccess::EncodedNode {
211 hash: *hash,
212 encoded_node: owned_node.data().into(),
213 });
214 }
215 }
216
217 Ok((owned_node, node_hash))
218 }
219
220 pub(crate) fn fetch_value(
222 &self,
223 hash: TrieHash<L>,
224 prefix: Prefix,
225 ) -> Result<DBValue, TrieHash<L>, CError<L>> {
226 let value = self
227 .db
228 .get(&hash, prefix)
229 .ok_or_else(|| Box::new(TrieError::IncompleteDatabase(hash)))?;
230
231 if let Some(recorder) = self.recorder.as_ref() {
232 debug_assert!(prefix.1.is_none(), "A value has never a partial key; qed");
233
234 recorder.borrow_mut().record(TrieAccess::Value {
235 hash,
236 value: value.as_slice().into(),
237 full_key: prefix.0,
238 });
239 }
240
241 Ok(value)
242 }
243}
244
245impl<'db, 'cache, L> Trie<L> for TrieDB<'db, 'cache, L>
246where
247 L: TrieLayout,
248{
249 fn root(&self) -> &TrieHash<L> {
250 self.root
251 }
252
253 fn get_hash(&self, key: &[u8]) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
254 let mut cache = self.cache.as_ref().map(|c| c.borrow_mut());
255 let mut recorder = self.recorder.as_ref().map(|r| r.borrow_mut());
256
257 Lookup::<L, _> {
258 db: self.db,
259 query: |_: &[u8]| (),
260 hash: *self.root,
261 cache: cache.as_mut().map(|c| &mut ***c as &mut dyn TrieCache<L::Codec>),
262 recorder: recorder.as_mut().map(|r| &mut ***r as &mut dyn TrieRecorder<TrieHash<L>>),
263 }
264 .look_up_hash(key, NibbleSlice::new(key))
265 }
266
267 fn get_with<Q: Query<L::Hash>>(
268 &self,
269 key: &[u8],
270 query: Q,
271 ) -> Result<Option<Q::Item>, TrieHash<L>, CError<L>> {
272 let mut cache = self.cache.as_ref().map(|c| c.borrow_mut());
273 let mut recorder = self.recorder.as_ref().map(|r| r.borrow_mut());
274
275 Lookup::<L, Q> {
276 db: self.db,
277 query,
278 hash: *self.root,
279 cache: cache.as_mut().map(|c| &mut ***c as &mut dyn TrieCache<L::Codec>),
280 recorder: recorder.as_mut().map(|r| &mut ***r as &mut dyn TrieRecorder<TrieHash<L>>),
281 }
282 .look_up(key, NibbleSlice::new(key))
283 }
284
285 fn lookup_first_descendant(
286 &self,
287 key: &[u8],
288 ) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>> {
289 let mut cache = self.cache.as_ref().map(|c| c.borrow_mut());
290 let mut recorder = self.recorder.as_ref().map(|r| r.borrow_mut());
291
292 Lookup::<L, _> {
293 db: self.db,
294 query: |_: &[u8]| (),
295 hash: *self.root,
296 cache: cache.as_mut().map(|c| &mut ***c as &mut dyn TrieCache<L::Codec>),
297 recorder: recorder.as_mut().map(|r| &mut ***r as &mut dyn TrieRecorder<TrieHash<L>>),
298 }
299 .lookup_first_descendant(key, NibbleSlice::new(key))
300 }
301
302 fn iter<'a>(
303 &'a self,
304 ) -> Result<
305 Box<dyn TrieIterator<L, Item = TrieItem<TrieHash<L>, CError<L>>> + 'a>,
306 TrieHash<L>,
307 CError<L>,
308 > {
309 TrieDBIterator::new(self).map(|iter| Box::new(iter) as Box<_>)
310 }
311
312 fn key_iter<'a>(
313 &'a self,
314 ) -> Result<
315 Box<dyn TrieIterator<L, Item = TrieKeyItem<TrieHash<L>, CError<L>>> + 'a>,
316 TrieHash<L>,
317 CError<L>,
318 > {
319 TrieDBKeyIterator::new(self).map(|iter| Box::new(iter) as Box<_>)
320 }
321}
322
323#[cfg(feature = "std")]
325struct TrieAwareDebugNode<'db, 'cache, 'a, L>
326where
327 L: TrieLayout,
328{
329 trie: &'db TrieDB<'db, 'cache, L>,
330 node_key: NodeHandle<'a>,
331 partial_key: NibbleVec,
332 index: Option<u8>,
333}
334
335#[cfg(feature = "std")]
336impl<'db, 'cache, 'a, L> fmt::Debug for TrieAwareDebugNode<'db, 'cache, 'a, L>
337where
338 L: TrieLayout,
339{
340 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
341 match self.trie.get_raw_or_lookup(
342 <TrieHash<L>>::default(),
343 self.node_key,
344 self.partial_key.as_prefix(),
345 false,
346 ) {
347 Ok((owned_node, _node_hash)) => match owned_node.node() {
348 Node::Leaf(slice, value) => {
349 let mut disp = f.debug_struct("Node::Leaf");
350 if let Some(i) = self.index {
351 disp.field("index", &i);
352 }
353 disp.field("slice", &slice).field("value", &value);
354 disp.finish()
355 },
356 Node::Extension(slice, item) => {
357 let mut disp = f.debug_struct("Node::Extension");
358 if let Some(i) = self.index {
359 disp.field("index", &i);
360 }
361 disp.field("slice", &slice).field(
362 "item",
363 &TrieAwareDebugNode {
364 trie: self.trie,
365 node_key: item,
366 partial_key: self
367 .partial_key
368 .clone_append_optional_slice_and_nibble(Some(&slice), None),
369 index: None,
370 },
371 );
372 disp.finish()
373 },
374 Node::Branch(ref nodes, ref value) => {
375 let nodes: Vec<TrieAwareDebugNode<L>> = nodes
376 .into_iter()
377 .enumerate()
378 .filter_map(|(i, n)| n.map(|n| (i, n)))
379 .map(|(i, n)| TrieAwareDebugNode {
380 trie: self.trie,
381 index: Some(i as u8),
382 node_key: n,
383 partial_key: self
384 .partial_key
385 .clone_append_optional_slice_and_nibble(None, Some(i as u8)),
386 })
387 .collect();
388 let mut disp = f.debug_struct("Node::Branch");
389 if let Some(i) = self.index {
390 disp.field("index", &i);
391 }
392 disp.field("nodes", &nodes).field("value", &value);
393 disp.finish()
394 },
395 Node::NibbledBranch(slice, nodes, value) => {
396 let nodes: Vec<TrieAwareDebugNode<L>> = nodes
397 .iter()
398 .enumerate()
399 .filter_map(|(i, n)| n.map(|n| (i, n)))
400 .map(|(i, n)| TrieAwareDebugNode {
401 trie: self.trie,
402 index: Some(i as u8),
403 node_key: n,
404 partial_key: self.partial_key.clone_append_optional_slice_and_nibble(
405 Some(&slice),
406 Some(i as u8),
407 ),
408 })
409 .collect();
410 let mut disp = f.debug_struct("Node::NibbledBranch");
411 if let Some(i) = self.index {
412 disp.field("index", &i);
413 }
414 disp.field("slice", &slice).field("nodes", &nodes).field("value", &value);
415 disp.finish()
416 },
417 Node::Empty => {
418 let mut disp = f.debug_struct("Node::Empty");
419 disp.finish()
420 },
421 },
422 Err(e) => f
423 .debug_struct("BROKEN_NODE")
424 .field("index", &self.index)
425 .field("key", &self.node_key)
426 .field("error", &format!("ERROR fetching node: {}", e))
427 .finish(),
428 }
429 }
430}
431
432#[cfg(feature = "std")]
433impl<'db, 'cache, L> fmt::Debug for TrieDB<'db, 'cache, L>
434where
435 L: TrieLayout,
436{
437 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
438 f.debug_struct("TrieDB")
439 .field(
440 "root",
441 &TrieAwareDebugNode {
442 trie: self,
443 node_key: NodeHandle::Hash(self.root().as_ref()),
444 partial_key: NibbleVec::new(),
445 index: None,
446 },
447 )
448 .finish()
449 }
450}
451
452pub struct TrieDBIterator<'a, 'cache, L: TrieLayout> {
454 db: &'a TrieDB<'a, 'cache, L>,
455 raw_iter: TrieDBRawIterator<L>,
456}
457
458pub struct TrieDBKeyIterator<'a, 'cache, L: TrieLayout> {
460 db: &'a TrieDB<'a, 'cache, L>,
461 raw_iter: TrieDBRawIterator<L>,
462}
463
464pub struct TrieDBKeyDoubleEndedIterator<'a, 'cache, L: TrieLayout> {
467 db: &'a TrieDB<'a, 'cache, L>,
468 raw_iter: TrieDBRawIterator<L>,
469 back_raw_iter: TrieDBRawIterator<L>,
470}
471
472impl<'a, 'cache, L: TrieLayout> TrieDBKeyDoubleEndedIterator<'a, 'cache, L> {
473 pub fn new(db: &'a TrieDB<'a, 'cache, L>) -> Result<Self, TrieHash<L>, CError<L>> {
475 Ok(Self {
476 db,
477 raw_iter: TrieDBRawIterator::new(db)?,
478 back_raw_iter: TrieDBRawIterator::new(db)?,
479 })
480 }
481}
482
483pub struct TrieDBDoubleEndedIterator<'a, 'cache, L: TrieLayout> {
485 db: &'a TrieDB<'a, 'cache, L>,
486 raw_iter: TrieDBRawIterator<L>,
487 back_raw_iter: TrieDBRawIterator<L>,
488}
489
490impl<'a, 'cache, L: TrieLayout> TrieDBDoubleEndedIterator<'a, 'cache, L> {
491 pub fn new(db: &'a TrieDB<'a, 'cache, L>) -> Result<Self, TrieHash<L>, CError<L>> {
493 Ok(Self {
494 db,
495 raw_iter: TrieDBRawIterator::new(db)?,
496 back_raw_iter: TrieDBRawIterator::new(db)?,
497 })
498 }
499
500 pub fn new_prefixed(
502 db: &'a TrieDB<'a, 'cache, L>,
503 prefix: &[u8],
504 ) -> Result<Self, TrieHash<L>, CError<L>> {
505 Ok(Self {
506 db,
507 raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)?,
508 back_raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)?,
509 })
510 }
511}
512
513impl<L: TrieLayout> TrieDoubleEndedIterator<L> for TrieDBDoubleEndedIterator<'_, '_, L> {}
514
515impl<'a, 'cache, L: TrieLayout> TrieDBIterator<'a, 'cache, L> {
516 pub fn new(db: &'a TrieDB<'a, 'cache, L>) -> Result<Self, TrieHash<L>, CError<L>> {
518 Ok(Self { db, raw_iter: TrieDBRawIterator::new(db)? })
519 }
520
521 pub fn new_prefixed(
523 db: &'a TrieDB<'a, 'cache, L>,
524 prefix: &[u8],
525 ) -> Result<Self, TrieHash<L>, CError<L>> {
526 Ok(Self { db, raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)? })
527 }
528
529 pub fn new_prefixed_then_seek(
533 db: &'a TrieDB<'a, 'cache, L>,
534 prefix: &[u8],
535 start_at: &[u8],
536 ) -> Result<Self, TrieHash<L>, CError<L>> {
537 Ok(Self { db, raw_iter: TrieDBRawIterator::new_prefixed_then_seek(db, prefix, start_at)? })
538 }
539
540 pub fn from_raw(db: &'a TrieDB<'a, 'cache, L>, raw_iter: TrieDBRawIterator<L>) -> Self {
542 Self { db, raw_iter }
543 }
544
545 pub fn into_raw(self) -> TrieDBRawIterator<L> {
547 self.raw_iter
548 }
549}
550
551impl<'a, 'cache, L: TrieLayout> TrieIterator<L> for TrieDBIterator<'a, 'cache, L> {
552 fn seek(&mut self, key: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
554 self.raw_iter.seek(self.db, key, true).map(|_| ())
555 }
556}
557
558impl<'a, 'cache, L: TrieLayout> Iterator for TrieDBIterator<'a, 'cache, L> {
559 type Item = TrieItem<TrieHash<L>, CError<L>>;
560
561 fn next(&mut self) -> Option<Self::Item> {
562 self.raw_iter.next_item(self.db)
563 }
564}
565
566impl<'a, 'cache, L: TrieLayout> TrieIterator<L> for TrieDBDoubleEndedIterator<'a, 'cache, L> {
567 fn seek(&mut self, key: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
568 self.raw_iter.seek(self.db, key, true).map(|_| ())?;
569 self.back_raw_iter.seek(self.db, key, false).map(|_| ())
570 }
571}
572
573impl<'a, 'cache, L: TrieLayout> Iterator for TrieDBDoubleEndedIterator<'a, 'cache, L> {
574 type Item = TrieItem<TrieHash<L>, CError<L>>;
575
576 fn next(&mut self) -> Option<Self::Item> {
577 self.raw_iter.next_item(self.db)
578 }
579}
580
581impl<'a, 'cache, L: TrieLayout> DoubleEndedIterator for TrieDBDoubleEndedIterator<'a, 'cache, L> {
582 fn next_back(&mut self) -> Option<Self::Item> {
583 self.back_raw_iter.prev_item(self.db)
584 }
585}
586
587impl<'a, 'cache, L: TrieLayout> TrieDBKeyIterator<'a, 'cache, L> {
588 pub fn new(db: &'a TrieDB<'a, 'cache, L>) -> Result<Self, TrieHash<L>, CError<L>> {
590 Ok(Self { db, raw_iter: TrieDBRawIterator::new(db)? })
591 }
592
593 pub fn new_prefixed(
595 db: &'a TrieDB<'a, 'cache, L>,
596 prefix: &[u8],
597 ) -> Result<Self, TrieHash<L>, CError<L>> {
598 Ok(Self { db, raw_iter: TrieDBRawIterator::new_prefixed(db, prefix)? })
599 }
600
601 pub fn new_prefixed_then_seek(
605 db: &'a TrieDB<'a, 'cache, L>,
606 prefix: &[u8],
607 start_at: &[u8],
608 ) -> Result<TrieDBKeyIterator<'a, 'cache, L>, TrieHash<L>, CError<L>> {
609 Ok(Self { db, raw_iter: TrieDBRawIterator::new_prefixed_then_seek(db, prefix, start_at)? })
610 }
611
612 pub fn from_raw(db: &'a TrieDB<'a, 'cache, L>, raw_iter: TrieDBRawIterator<L>) -> Self {
614 Self { db, raw_iter }
615 }
616
617 pub fn into_raw(self) -> TrieDBRawIterator<L> {
619 self.raw_iter
620 }
621}
622
623impl<'a, 'cache, L: TrieLayout> TrieIterator<L> for TrieDBKeyIterator<'a, 'cache, L> {
624 fn seek(&mut self, key: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
626 self.raw_iter.seek(self.db, key, true).map(|_| ())
627 }
628}
629
630impl<'a, 'cache, L: TrieLayout> Iterator for TrieDBKeyIterator<'a, 'cache, L> {
631 type Item = TrieKeyItem<TrieHash<L>, CError<L>>;
632
633 fn next(&mut self) -> Option<Self::Item> {
634 self.raw_iter.next_key(self.db)
635 }
636}
637
638impl<'a, 'cache, L: TrieLayout> TrieIterator<L> for TrieDBKeyDoubleEndedIterator<'a, 'cache, L> {
639 fn seek(&mut self, key: &[u8]) -> Result<(), TrieHash<L>, CError<L>> {
641 self.raw_iter.seek(self.db, key, true).map(|_| ())?;
642 self.back_raw_iter.seek(self.db, key, false).map(|_| ())
643 }
644}
645
646impl<'a, 'cache, L: TrieLayout> Iterator for TrieDBKeyDoubleEndedIterator<'a, 'cache, L> {
647 type Item = TrieKeyItem<TrieHash<L>, CError<L>>;
648
649 fn next(&mut self) -> Option<Self::Item> {
650 self.raw_iter.next_key(self.db)
651 }
652}
653
654impl<'a, 'cache, L: TrieLayout> DoubleEndedIterator
655 for TrieDBKeyDoubleEndedIterator<'a, 'cache, L>
656{
657 fn next_back(&mut self) -> Option<Self::Item> {
658 self.back_raw_iter.prev_key(self.db)
659 }
660}