trie_db/
recorder.rs

1// Copyright 2017, 2021 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//! Trie query recorder.
16
17use crate::{
18	rstd::{vec::Vec, BTreeMap},
19	RecordedForKey, TrieAccess, TrieHash, TrieLayout, TrieRecorder,
20};
21
22/// The record of a visited node.
23#[cfg_attr(feature = "std", derive(Debug))]
24#[derive(PartialEq, Eq, Clone)]
25pub struct Record<HO> {
26	/// The hash of the node.
27	pub hash: HO,
28	/// The data representing the node.
29	pub data: Vec<u8>,
30}
31
32/// Records trie nodes as they pass it.
33#[cfg_attr(feature = "std", derive(Debug))]
34pub struct Recorder<L: TrieLayout> {
35	nodes: Vec<Record<TrieHash<L>>>,
36	recorded_keys: BTreeMap<Vec<u8>, RecordedForKey>,
37}
38
39impl<L: TrieLayout> Default for Recorder<L> {
40	fn default() -> Self {
41		Recorder::new()
42	}
43}
44
45impl<L: TrieLayout> Recorder<L> {
46	/// Create a new `Recorder` which records all given nodes.
47	pub fn new() -> Self {
48		Self { nodes: Default::default(), recorded_keys: Default::default() }
49	}
50
51	/// Drain all visited records.
52	pub fn drain(&mut self) -> Vec<Record<TrieHash<L>>> {
53		self.recorded_keys.clear();
54		crate::rstd::mem::take(&mut self.nodes)
55	}
56}
57
58impl<L: TrieLayout> TrieRecorder<TrieHash<L>> for Recorder<L> {
59	fn record<'a>(&mut self, access: TrieAccess<'a, TrieHash<L>>) {
60		match access {
61			TrieAccess::EncodedNode { hash, encoded_node, .. } => {
62				self.nodes.push(Record { hash, data: encoded_node.to_vec() });
63			},
64			TrieAccess::NodeOwned { hash, node_owned, .. } => {
65				self.nodes.push(Record { hash, data: node_owned.to_encoded::<L::Codec>() });
66			},
67			TrieAccess::Value { hash, value, full_key } => {
68				self.nodes.push(Record { hash, data: value.to_vec() });
69				self.recorded_keys.insert(full_key.to_vec(), RecordedForKey::Value);
70			},
71			TrieAccess::Hash { full_key } => {
72				self.recorded_keys.entry(full_key.to_vec()).or_insert(RecordedForKey::Hash);
73			},
74			TrieAccess::NonExisting { full_key } => {
75				// We handle the non existing value/hash like having recorded the value.
76				self.recorded_keys.insert(full_key.to_vec(), RecordedForKey::Value);
77			},
78			TrieAccess::InlineValue { full_key } => {
79				self.recorded_keys.insert(full_key.to_vec(), RecordedForKey::Value);
80			},
81		}
82	}
83
84	fn trie_nodes_recorded_for_key(&self, key: &[u8]) -> RecordedForKey {
85		self.recorded_keys.get(key).copied().unwrap_or(RecordedForKey::None)
86	}
87}
trie_db/recorder.rs

trie_db/
recorder.rs
