1use crate::{
18 nibble::NibbleSlice,
19 node::{decode_hash, Node, NodeHandle, NodeHandleOwned, NodeOwned, Value, ValueOwned},
20 node_codec::NodeCodec,
21 rstd::{boxed::Box, vec::Vec},
22 Bytes, CError, CachedValue, DBValue, MerkleValue, Query, RecordedForKey, Result, TrieAccess,
23 TrieCache, TrieError, TrieHash, TrieLayout, TrieRecorder,
24};
25use hash_db::{HashDBRef, Hasher, Prefix};
26
27pub struct Lookup<'a, 'cache, L: TrieLayout, Q: Query<L::Hash>> {
29 pub db: &'a dyn HashDBRef<L::Hash, DBValue>,
31 pub query: Q,
33 pub hash: TrieHash<L>,
35 pub cache: Option<&'cache mut dyn TrieCache<L::Codec>>,
37 pub recorder: Option<&'cache mut dyn TrieRecorder<TrieHash<L>>>,
39}
40
41impl<'a, 'cache, L, Q> Lookup<'a, 'cache, L, Q>
42where
43 L: TrieLayout,
44 Q: Query<L::Hash>,
45{
46 fn load_value(
53 v: Value,
54 prefix: Prefix,
55 full_key: &[u8],
56 db: &dyn HashDBRef<L::Hash, DBValue>,
57 recorder: &mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
58 query: Q,
59 ) -> Result<Q::Item, TrieHash<L>, CError<L>> {
60 match v {
61 Value::Inline(value) => {
62 if let Some(recorder) = recorder {
63 recorder.record(TrieAccess::InlineValue { full_key });
64 }
65
66 Ok(query.decode(&value))
67 },
68 Value::Node(hash) => {
69 let mut res = TrieHash::<L>::default();
70 res.as_mut().copy_from_slice(hash);
71 if let Some(value) = db.get(&res, prefix) {
72 if let Some(recorder) = recorder {
73 recorder.record(TrieAccess::Value {
74 hash: res,
75 value: value.as_slice().into(),
76 full_key,
77 });
78 }
79
80 Ok(query.decode(&value))
81 } else {
82 Err(Box::new(TrieError::IncompleteDatabase(res)))
83 }
84 },
85 }
86 }
87
88 fn load_owned_value(
95 v: ValueOwned<TrieHash<L>>,
96 prefix: Prefix,
97 full_key: &[u8],
98 cache: &mut dyn crate::TrieCache<L::Codec>,
99 db: &dyn HashDBRef<L::Hash, DBValue>,
100 recorder: &mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
101 ) -> Result<(Bytes, TrieHash<L>), TrieHash<L>, CError<L>> {
102 match v {
103 ValueOwned::Inline(value, hash) => {
104 if let Some(recorder) = recorder {
105 recorder.record(TrieAccess::InlineValue { full_key });
106 }
107
108 Ok((value.clone(), hash))
109 },
110 ValueOwned::Node(hash) => {
111 let node = cache.get_or_insert_node(hash, &mut || {
112 let value = db
113 .get(&hash, prefix)
114 .ok_or_else(|| Box::new(TrieError::IncompleteDatabase(hash)))?;
115
116 Ok(NodeOwned::Value(value.into(), hash))
117 })?;
118
119 let value = node
120 .data()
121 .expect(
122 "We are caching a `NodeOwned::Value` for a value node \
123 hash and this cached node has always data attached; qed",
124 )
125 .clone();
126
127 if let Some(recorder) = recorder {
128 recorder.record(TrieAccess::Value {
129 hash,
130 value: value.as_ref().into(),
131 full_key,
132 });
133 }
134
135 Ok((value, hash))
136 },
137 }
138 }
139
140 fn record<'b>(&mut self, get_access: impl FnOnce() -> TrieAccess<'b, TrieHash<L>>)
141 where
142 TrieHash<L>: 'b,
143 {
144 if let Some(recorder) = self.recorder.as_mut() {
145 recorder.record(get_access());
146 }
147 }
148 pub fn lookup_first_descendant(
155 mut self,
156 full_key: &[u8],
157 nibble_key: NibbleSlice,
158 ) -> Result<Option<MerkleValue<TrieHash<L>>>, TrieHash<L>, CError<L>> {
159 let mut partial = nibble_key;
160 let mut hash = self.hash;
161 let mut key_nibbles = 0;
162
163 let mut cache = self.cache.take();
164
165 for depth in 0.. {
167 let mut _owned_node = NodeOwned::Empty;
170
171 let mut node_data = Vec::new();
176
177 let mut get_owned_node = |depth: i32| {
179 let data = match self.db.get(&hash, nibble_key.mid(key_nibbles).left()) {
180 Some(value) => value,
181 None =>
182 return Err(Box::new(match depth {
183 0 => TrieError::InvalidStateRoot(hash),
184 _ => TrieError::IncompleteDatabase(hash),
185 })),
186 };
187
188 let decoded = match L::Codec::decode(&data[..]) {
189 Ok(node) => node,
190 Err(e) => return Err(Box::new(TrieError::DecoderError(hash, e))),
191 };
192
193 let owned = decoded.to_owned_node::<L>()?;
194 node_data = data;
195 Ok(owned)
196 };
197
198 let mut node = if let Some(cache) = &mut cache {
199 let node = cache.get_or_insert_node(hash, &mut || get_owned_node(depth))?;
200
201 self.record(|| TrieAccess::NodeOwned { hash, node_owned: node });
202
203 node
204 } else {
205 _owned_node = get_owned_node(depth)?;
206
207 self.record(|| TrieAccess::EncodedNode {
208 hash,
209 encoded_node: node_data.as_slice().into(),
210 });
211
212 &_owned_node
213 };
214
215 let mut is_inline = false;
218 loop {
219 let next_node = match node {
220 NodeOwned::Leaf(slice, _) => {
221 if !slice.starts_with_slice(&partial) {
224 self.record(|| TrieAccess::NonExisting { full_key });
225 return Ok(None)
226 }
227
228 if partial.len() != slice.len() {
229 self.record(|| TrieAccess::NonExisting { full_key });
230 }
231
232 let res = is_inline
233 .then(|| MerkleValue::Node(node_data))
234 .unwrap_or_else(|| MerkleValue::Hash(hash));
235 return Ok(Some(res))
236 },
237 NodeOwned::Extension(slice, item) => {
238 if partial.len() < slice.len() {
239 self.record(|| TrieAccess::NonExisting { full_key });
240
241 return if slice.starts_with_slice(&partial) {
245 let res = is_inline
246 .then(|| MerkleValue::Node(node_data))
247 .unwrap_or_else(|| MerkleValue::Hash(hash));
248 Ok(Some(res))
249 } else {
250 Ok(None)
251 }
252 }
253
254 if partial.starts_with_vec(&slice) {
258 partial = partial.mid(slice.len());
260 key_nibbles += slice.len();
261 item
262 } else {
263 self.record(|| TrieAccess::NonExisting { full_key });
264
265 return Ok(None)
266 }
267 },
268 NodeOwned::Branch(children, value) =>
269 if partial.is_empty() {
270 if value.is_none() {
271 self.record(|| TrieAccess::NonExisting { full_key });
272 }
273 let res = is_inline
274 .then(|| MerkleValue::Node(node_data))
275 .unwrap_or_else(|| MerkleValue::Hash(hash));
276 return Ok(Some(res))
277 } else {
278 match children.get(partial.at(0) as usize) {
279 Some(x) => {
280 partial = partial.mid(1);
281 key_nibbles += 1;
282 x
283 },
284 None => {
285 self.record(|| TrieAccess::NonExisting { full_key });
286
287 return Ok(None)
288 },
289 }
290 },
291 NodeOwned::NibbledBranch(slice, children, value) => {
292 if partial.len() < slice.len() {
294 self.record(|| TrieAccess::NonExisting { full_key });
295
296 return if slice.starts_with_slice(&partial) {
299 let res = is_inline
300 .then(|| MerkleValue::Node(node_data))
301 .unwrap_or_else(|| MerkleValue::Hash(hash));
302 Ok(Some(res))
303 } else {
304 Ok(None)
305 }
306 }
307
308 if !partial.starts_with_vec(&slice) {
311 self.record(|| TrieAccess::NonExisting { full_key });
312 return Ok(None)
313 }
314
315 if partial.len() == slice.len() {
317 if value.is_none() {
318 self.record(|| TrieAccess::NonExisting { full_key });
319 }
320
321 let res = is_inline
322 .then(|| MerkleValue::Node(node_data))
323 .unwrap_or_else(|| MerkleValue::Hash(hash));
324 return Ok(Some(res))
325 } else {
326 match children.get(partial.at(slice.len()) as usize) {
327 Some(x) => {
328 partial = partial.mid(slice.len() + 1);
329 key_nibbles += slice.len() + 1;
330 x
331 },
332 None => {
333 self.record(|| TrieAccess::NonExisting { full_key });
334
335 return Ok(None)
336 },
337 }
338 }
339 },
340 NodeOwned::Empty => {
341 self.record(|| TrieAccess::NonExisting { full_key });
342
343 return Ok(None)
344 },
345 NodeOwned::Value(_, _) => {
346 unreachable!(
347 "`NodeOwned::Value` can not be reached by using the hash of a node. \
348 `NodeOwned::Value` is only constructed when loading a value into memory, \
349 which needs to have a different hash than any node; qed",
350 )
351 },
352 };
353
354 match next_node {
356 NodeHandleOwned::Hash(new_hash) => {
357 hash = *new_hash;
358 break
359 },
360 NodeHandleOwned::Inline(inline_node) => {
361 node = &inline_node;
362 is_inline = true;
363 },
364 }
365 }
366 }
367 Ok(None)
368 }
369
370 pub fn look_up(
377 mut self,
378 full_key: &[u8],
379 nibble_key: NibbleSlice,
380 ) -> Result<Option<Q::Item>, TrieHash<L>, CError<L>> {
381 match self.cache.take() {
382 Some(cache) => self.look_up_with_cache(full_key, nibble_key, cache),
383 None => self.look_up_without_cache(nibble_key, full_key, Self::load_value),
384 }
385 }
386
387 pub fn look_up_hash(
392 mut self,
393 full_key: &[u8],
394 nibble_key: NibbleSlice,
395 ) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
396 match self.cache.take() {
397 Some(cache) => self.look_up_hash_with_cache(full_key, nibble_key, cache),
398 None => self.look_up_without_cache(
399 nibble_key,
400 full_key,
401 |v, _, full_key, _, recorder, _| {
402 Ok(match v {
403 Value::Inline(v) => {
404 if let Some(recoder) = recorder.as_mut() {
405 recoder.record(TrieAccess::InlineValue { full_key });
409 }
410
411 L::Hash::hash(&v)
412 },
413 Value::Node(hash_bytes) => {
414 if let Some(recoder) = recorder.as_mut() {
415 recoder.record(TrieAccess::Hash { full_key });
416 }
417
418 let mut hash = TrieHash::<L>::default();
419 hash.as_mut().copy_from_slice(hash_bytes);
420 hash
421 },
422 })
423 },
424 ),
425 }
426 }
427
428 fn look_up_hash_with_cache(
432 mut self,
433 full_key: &[u8],
434 nibble_key: NibbleSlice,
435 cache: &mut dyn crate::TrieCache<L::Codec>,
436 ) -> Result<Option<TrieHash<L>>, TrieHash<L>, CError<L>> {
437 let value_cache_allowed = self
438 .recorder
439 .as_ref()
440 .map(|r| !r.trie_nodes_recorded_for_key(full_key).is_none())
442 .unwrap_or(true);
444
445 let res = if let Some(hash) = value_cache_allowed
446 .then(|| cache.lookup_value_for_key(full_key).map(|v| v.hash()))
447 .flatten()
448 {
449 hash
450 } else {
451 let hash_and_value = self.look_up_with_cache_internal(
452 nibble_key,
453 full_key,
454 cache,
455 |value, _, full_key, _, _, recorder| match value {
456 ValueOwned::Inline(value, hash) => {
457 if let Some(recoder) = recorder.as_mut() {
458 recoder.record(TrieAccess::InlineValue { full_key });
462 }
463
464 Ok((hash, Some(value.clone())))
465 },
466 ValueOwned::Node(hash) => {
467 if let Some(recoder) = recorder.as_mut() {
468 recoder.record(TrieAccess::Hash { full_key });
469 }
470
471 Ok((hash, None))
472 },
473 },
474 )?;
475
476 match &hash_and_value {
477 Some((hash, Some(value))) =>
478 cache.cache_value_for_key(full_key, (value.clone(), *hash).into()),
479 Some((hash, None)) => cache.cache_value_for_key(full_key, (*hash).into()),
480 None => cache.cache_value_for_key(full_key, CachedValue::NonExisting),
481 }
482
483 hash_and_value.map(|v| v.0)
484 };
485
486 Ok(res)
487 }
488
489 fn look_up_with_cache(
494 mut self,
495 full_key: &[u8],
496 nibble_key: NibbleSlice,
497 cache: &mut dyn crate::TrieCache<L::Codec>,
498 ) -> Result<Option<Q::Item>, TrieHash<L>, CError<L>> {
499 let trie_nodes_recorded =
500 self.recorder.as_ref().map(|r| r.trie_nodes_recorded_for_key(full_key));
501
502 let (value_cache_allowed, value_recording_required) = match trie_nodes_recorded {
503 Some(RecordedForKey::Value) | None => (true, false),
506 Some(RecordedForKey::Hash) => (true, true),
509 Some(RecordedForKey::None) => (false, true),
511 };
512
513 let lookup_data = |lookup: &mut Self,
514 cache: &mut dyn crate::TrieCache<L::Codec>|
515 -> Result<Option<Bytes>, TrieHash<L>, CError<L>> {
516 let data = lookup.look_up_with_cache_internal(
517 nibble_key,
518 full_key,
519 cache,
520 Self::load_owned_value,
521 )?;
522
523 cache.cache_value_for_key(full_key, data.clone().into());
524
525 Ok(data.map(|d| d.0))
526 };
527
528 let res = match value_cache_allowed.then(|| cache.lookup_value_for_key(full_key)).flatten()
529 {
530 Some(CachedValue::NonExisting) => None,
531 Some(CachedValue::ExistingHash(hash)) => {
532 let data = Self::load_owned_value(
533 ValueOwned::Node(*hash),
536 nibble_key.original_data_as_prefix(),
537 full_key,
538 cache,
539 self.db,
540 &mut self.recorder,
541 )?;
542
543 cache.cache_value_for_key(full_key, data.clone().into());
544
545 Some(data.0)
546 },
547 Some(CachedValue::Existing { data, hash, .. }) =>
548 if let Some(data) = data.upgrade() {
549 let is_inline =
552 L::MAX_INLINE_VALUE.map_or(true, |max| max as usize > data.as_ref().len());
553 if value_recording_required && !is_inline {
554 self.record(|| TrieAccess::Value {
556 hash: *hash,
557 value: data.as_ref().into(),
558 full_key,
559 });
560 }
561
562 Some(data)
563 } else {
564 lookup_data(&mut self, cache)?
565 },
566 None => lookup_data(&mut self, cache)?,
567 };
568
569 Ok(res.map(|v| self.query.decode(&v)))
570 }
571
572 fn look_up_with_cache_internal<R>(
575 &mut self,
576 nibble_key: NibbleSlice,
577 full_key: &[u8],
578 cache: &mut dyn crate::TrieCache<L::Codec>,
579 load_value_owned: impl Fn(
580 ValueOwned<TrieHash<L>>,
581 Prefix,
582 &[u8],
583 &mut dyn crate::TrieCache<L::Codec>,
584 &dyn HashDBRef<L::Hash, DBValue>,
585 &mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
586 ) -> Result<R, TrieHash<L>, CError<L>>,
587 ) -> Result<Option<R>, TrieHash<L>, CError<L>> {
588 let mut partial = nibble_key;
589 let mut hash = self.hash;
590 let mut key_nibbles = 0;
591
592 for depth in 0.. {
594 let mut node = cache.get_or_insert_node(hash, &mut || {
595 let node_data = match self.db.get(&hash, nibble_key.mid(key_nibbles).left()) {
596 Some(value) => value,
597 None =>
598 return Err(Box::new(match depth {
599 0 => TrieError::InvalidStateRoot(hash),
600 _ => TrieError::IncompleteDatabase(hash),
601 })),
602 };
603
604 let decoded = match L::Codec::decode(&node_data[..]) {
605 Ok(node) => node,
606 Err(e) => return Err(Box::new(TrieError::DecoderError(hash, e))),
607 };
608
609 decoded.to_owned_node::<L>()
610 })?;
611
612 self.record(|| TrieAccess::NodeOwned { hash, node_owned: node });
613
614 loop {
617 let next_node = match node {
618 NodeOwned::Leaf(slice, value) =>
619 return if partial == *slice {
620 let value = (*value).clone();
621 load_value_owned(
622 value,
623 nibble_key.original_data_as_prefix(),
624 full_key,
625 cache,
626 self.db,
627 &mut self.recorder,
628 )
629 .map(Some)
630 } else {
631 self.record(|| TrieAccess::NonExisting { full_key });
632
633 Ok(None)
634 },
635 NodeOwned::Extension(slice, item) =>
636 if partial.starts_with_vec(&slice) {
637 partial = partial.mid(slice.len());
638 key_nibbles += slice.len();
639 item
640 } else {
641 self.record(|| TrieAccess::NonExisting { full_key });
642
643 return Ok(None)
644 },
645 NodeOwned::Branch(children, value) =>
646 if partial.is_empty() {
647 return if let Some(value) = value.clone() {
648 load_value_owned(
649 value,
650 nibble_key.original_data_as_prefix(),
651 full_key,
652 cache,
653 self.db,
654 &mut self.recorder,
655 )
656 .map(Some)
657 } else {
658 self.record(|| TrieAccess::NonExisting { full_key });
659
660 Ok(None)
661 }
662 } else {
663 match children.get(partial.at(0) as usize) {
664 Some(x) => {
665 partial = partial.mid(1);
666 key_nibbles += 1;
667 x
668 },
669 None => {
670 self.record(|| TrieAccess::NonExisting { full_key });
671
672 return Ok(None)
673 },
674 }
675 },
676 NodeOwned::NibbledBranch(slice, children, value) => {
677 if !partial.starts_with_vec(&slice) {
678 self.record(|| TrieAccess::NonExisting { full_key });
679
680 return Ok(None)
681 }
682
683 if partial.len() == slice.len() {
684 return if let Some(value) = value.clone() {
685 load_value_owned(
686 value,
687 nibble_key.original_data_as_prefix(),
688 full_key,
689 cache,
690 self.db,
691 &mut self.recorder,
692 )
693 .map(Some)
694 } else {
695 self.record(|| TrieAccess::NonExisting { full_key });
696
697 Ok(None)
698 }
699 } else {
700 match children.get(partial.at(slice.len()) as usize) {
701 Some(x) => {
702 partial = partial.mid(slice.len() + 1);
703 key_nibbles += slice.len() + 1;
704 x
705 },
706 None => {
707 self.record(|| TrieAccess::NonExisting { full_key });
708
709 return Ok(None)
710 },
711 }
712 }
713 },
714 NodeOwned::Empty => {
715 self.record(|| TrieAccess::NonExisting { full_key });
716
717 return Ok(None)
718 },
719 NodeOwned::Value(_, _) => {
720 unreachable!(
721 "`NodeOwned::Value` can not be reached by using the hash of a node. \
722 `NodeOwned::Value` is only constructed when loading a value into memory, \
723 which needs to have a different hash than any node; qed",
724 )
725 },
726 };
727
728 match next_node {
730 NodeHandleOwned::Hash(new_hash) => {
731 hash = *new_hash;
732 break
733 },
734 NodeHandleOwned::Inline(inline_node) => {
735 node = &inline_node;
736 },
737 }
738 }
739 }
740
741 Ok(None)
742 }
743
744 fn look_up_without_cache<R>(
750 mut self,
751 nibble_key: NibbleSlice,
752 full_key: &[u8],
753 load_value: impl Fn(
754 Value,
755 Prefix,
756 &[u8],
757 &dyn HashDBRef<L::Hash, DBValue>,
758 &mut Option<&mut dyn TrieRecorder<TrieHash<L>>>,
759 Q,
760 ) -> Result<R, TrieHash<L>, CError<L>>,
761 ) -> Result<Option<R>, TrieHash<L>, CError<L>> {
762 let mut partial = nibble_key;
763 let mut hash = self.hash;
764 let mut key_nibbles = 0;
765
766 for depth in 0.. {
768 let node_data = match self.db.get(&hash, nibble_key.mid(key_nibbles).left()) {
769 Some(value) => value,
770 None =>
771 return Err(Box::new(match depth {
772 0 => TrieError::InvalidStateRoot(hash),
773 _ => TrieError::IncompleteDatabase(hash),
774 })),
775 };
776
777 self.record(|| TrieAccess::EncodedNode {
778 hash,
779 encoded_node: node_data.as_slice().into(),
780 });
781
782 let mut node_data = &node_data[..];
785 loop {
786 let decoded = match L::Codec::decode(node_data) {
787 Ok(node) => node,
788 Err(e) => return Err(Box::new(TrieError::DecoderError(hash, e))),
789 };
790
791 let next_node = match decoded {
792 Node::Leaf(slice, value) =>
793 return if slice == partial {
794 load_value(
795 value,
796 nibble_key.original_data_as_prefix(),
797 full_key,
798 self.db,
799 &mut self.recorder,
800 self.query,
801 )
802 .map(Some)
803 } else {
804 self.record(|| TrieAccess::NonExisting { full_key });
805
806 Ok(None)
807 },
808 Node::Extension(slice, item) =>
809 if partial.starts_with(&slice) {
810 partial = partial.mid(slice.len());
811 key_nibbles += slice.len();
812 item
813 } else {
814 self.record(|| TrieAccess::NonExisting { full_key });
815
816 return Ok(None)
817 },
818 Node::Branch(children, value) =>
819 if partial.is_empty() {
820 return if let Some(val) = value {
821 load_value(
822 val,
823 nibble_key.original_data_as_prefix(),
824 full_key,
825 self.db,
826 &mut self.recorder,
827 self.query,
828 )
829 .map(Some)
830 } else {
831 self.record(|| TrieAccess::NonExisting { full_key });
832
833 Ok(None)
834 }
835 } else {
836 match children[partial.at(0) as usize] {
837 Some(x) => {
838 partial = partial.mid(1);
839 key_nibbles += 1;
840 x
841 },
842 None => {
843 self.record(|| TrieAccess::NonExisting { full_key });
844
845 return Ok(None)
846 },
847 }
848 },
849 Node::NibbledBranch(slice, children, value) => {
850 if !partial.starts_with(&slice) {
851 self.record(|| TrieAccess::NonExisting { full_key });
852
853 return Ok(None)
854 }
855
856 if partial.len() == slice.len() {
857 return if let Some(val) = value {
858 load_value(
859 val,
860 nibble_key.original_data_as_prefix(),
861 full_key,
862 self.db,
863 &mut self.recorder,
864 self.query,
865 )
866 .map(Some)
867 } else {
868 self.record(|| TrieAccess::NonExisting { full_key });
869
870 Ok(None)
871 }
872 } else {
873 match children[partial.at(slice.len()) as usize] {
874 Some(x) => {
875 partial = partial.mid(slice.len() + 1);
876 key_nibbles += slice.len() + 1;
877 x
878 },
879 None => {
880 self.record(|| TrieAccess::NonExisting { full_key });
881
882 return Ok(None)
883 },
884 }
885 }
886 },
887 Node::Empty => {
888 self.record(|| TrieAccess::NonExisting { full_key });
889
890 return Ok(None)
891 },
892 };
893
894 match next_node {
896 NodeHandle::Hash(data) => {
897 hash = decode_hash::<L::Hash>(data)
898 .ok_or_else(|| Box::new(TrieError::InvalidHash(hash, data.to_vec())))?;
899 break
900 },
901 NodeHandle::Inline(data) => {
902 node_data = data;
903 },
904 }
905 }
906 }
907 Ok(None)
908 }
909}