1use crate::{
40 nibble_ops::NIBBLE_LENGTH,
41 node::{decode_hash, Node, NodeHandle, NodeHandlePlan, NodePlan, OwnedNode, ValuePlan},
42 rstd::{
43 boxed::Box, convert::TryInto, marker::PhantomData, result, sync::Arc, vec, vec::Vec,
44 BTreeSet,
45 },
46 CError, ChildReference, DBValue, NibbleVec, NodeCodec, Result, TrieDB, TrieDBRawIterator,
47 TrieError, TrieHash, TrieLayout,
48};
49use hash_db::{HashDB, Prefix};
50
51const OMIT_VALUE_HASH: crate::node::Value<'static> = crate::node::Value::Inline(&[]);
52
53struct EncoderStackEntry<C: NodeCodec> {
54 prefix: NibbleVec,
56 node: Arc<OwnedNode<DBValue>>,
58 child_index: usize,
61 omit_children: Vec<bool>,
63 omit_value: bool,
65 output_index: usize,
68 _marker: PhantomData<C>,
69}
70
71impl<C: NodeCodec> EncoderStackEntry<C> {
72 fn advance_child_index(
81 &mut self,
82 child_prefix: &NibbleVec,
83 ) -> result::Result<(), &'static str> {
84 match self.node.node_plan() {
85 NodePlan::Empty | NodePlan::Leaf { .. } =>
86 return Err("empty and leaf nodes have no children"),
87 NodePlan::Extension { .. } =>
88 if self.child_index != 0 {
89 return Err("extension node cannot have multiple children")
90 },
91 NodePlan::Branch { .. } => {
92 if child_prefix.len() <= self.prefix.len() {
93 return Err("child_prefix does not contain prefix")
94 }
95 let child_index = child_prefix.at(self.prefix.len()) as usize;
96 if child_index < self.child_index {
97 return Err("iterator returned children in non-ascending order by prefix")
98 }
99 self.child_index = child_index;
100 },
101 NodePlan::NibbledBranch { partial, .. } => {
102 if child_prefix.len() <= self.prefix.len() + partial.len() {
103 return Err("child_prefix does not contain prefix and node partial")
104 }
105 let child_index = child_prefix.at(self.prefix.len() + partial.len()) as usize;
106 if child_index < self.child_index {
107 return Err("iterator returned children in non-ascending order by prefix")
108 }
109 self.child_index = child_index;
110 },
111 }
112 Ok(())
113 }
114
115 fn encode_node(&mut self) -> Result<Vec<u8>, C::HashOut, C::Error> {
117 let node_data = self.node.data();
118 let node_plan = self.node.node_plan();
119 let mut encoded = match node_plan {
120 NodePlan::Empty => node_data.to_vec(),
121 NodePlan::Leaf { partial, value: _ } =>
122 if self.omit_value {
123 let partial = partial.build(node_data);
124 C::leaf_node(partial.right_iter(), partial.len(), OMIT_VALUE_HASH)
125 } else {
126 node_data.to_vec()
127 },
128 NodePlan::Extension { partial, child: _ } =>
129 if !self.omit_children[0] {
130 node_data.to_vec()
131 } else {
132 let partial = partial.build(node_data);
133 let empty_child = ChildReference::Inline(C::HashOut::default(), 0);
134 C::extension_node(partial.right_iter(), partial.len(), empty_child)
135 },
136 NodePlan::Branch { value, children } => {
137 let value = if self.omit_value {
138 value.is_some().then_some(OMIT_VALUE_HASH)
139 } else {
140 value.as_ref().map(|v| v.build(node_data))
141 };
142 C::branch_node(
143 Self::branch_children(node_data, &children, &self.omit_children)?.iter(),
144 value,
145 )
146 },
147 NodePlan::NibbledBranch { partial, value, children } => {
148 let partial = partial.build(node_data);
149 let value = if self.omit_value {
150 value.is_some().then_some(OMIT_VALUE_HASH)
151 } else {
152 value.as_ref().map(|v| v.build(node_data))
153 };
154 C::branch_node_nibbled(
155 partial.right_iter(),
156 partial.len(),
157 Self::branch_children(node_data, &children, &self.omit_children)?.iter(),
158 value,
159 )
160 },
161 };
162
163 if self.omit_value {
164 if let Some(header) = C::ESCAPE_HEADER {
165 encoded.insert(0, header);
166 } else {
167 return Err(Box::new(TrieError::InvalidStateRoot(Default::default())))
168 }
169 }
170 Ok(encoded)
171 }
172
173 fn branch_children(
179 node_data: &[u8],
180 child_handles: &[Option<NodeHandlePlan>; NIBBLE_LENGTH],
181 omit_children: &[bool],
182 ) -> Result<[Option<ChildReference<C::HashOut>>; NIBBLE_LENGTH], C::HashOut, C::Error> {
183 let empty_child = ChildReference::Inline(C::HashOut::default(), 0);
184 let mut children = [None; NIBBLE_LENGTH];
185 for i in 0..NIBBLE_LENGTH {
186 children[i] = if omit_children[i] {
187 Some(empty_child)
188 } else if let Some(child_plan) = &child_handles[i] {
189 let child_ref = child_plan.build(node_data).try_into().map_err(|hash| {
190 Box::new(TrieError::InvalidHash(C::HashOut::default(), hash))
191 })?;
192 Some(child_ref)
193 } else {
194 None
195 };
196 }
197 Ok(children)
198 }
199}
200
201pub struct SeenHashes<L: TrieLayout> {
208 nodes: BTreeSet<TrieHash<L>>,
209 values: BTreeSet<TrieHash<L>>,
210}
211
212impl<L: TrieLayout> Default for SeenHashes<L> {
214 fn default() -> Self {
215 SeenHashes { nodes: BTreeSet::new(), values: BTreeSet::new() }
216 }
217}
218
219impl<L: TrieLayout> Clone for SeenHashes<L> {
220 fn clone(&self) -> Self {
221 SeenHashes { nodes: self.nodes.clone(), values: self.values.clone() }
222 }
223}
224
225fn detached_value<L: TrieLayout>(
232 db: &TrieDB<L>,
233 value: &ValuePlan,
234 node_data: &[u8],
235 node_prefix: Prefix,
236 seen: Option<&mut SeenHashes<L>>,
237) -> Option<Vec<u8>> {
238 let hash_plan = match value {
239 ValuePlan::Node(hash_plan) => hash_plan,
240 _ => return None,
241 };
242 let value_hash = &node_data[hash_plan.clone()];
243
244 let dedup_key = seen.as_ref().and_then(|_| decode_hash::<L::Hash>(value_hash));
245 if let (Some(seen), Some(key)) = (&seen, &dedup_key) {
246 if seen.values.contains(key) {
248 return None
249 }
250 }
251 let fetched = TrieDBRawIterator::fetch_value(db, value_hash, node_prefix).ok()?;
252 if let (Some(seen), Some(key)) = (seen, dedup_key) {
253 seen.values.insert(key);
254 }
255 Some(fetched)
256}
257
258pub fn encode_compact<L>(db: &TrieDB<L>) -> Result<Vec<Vec<u8>>, TrieHash<L>, CError<L>>
269where
270 L: TrieLayout,
271{
272 encode_compact_inner(db, None)
273}
274
275pub fn encode_compact_skip_duplicates<L>(
296 db: &TrieDB<L>,
297 seen: &mut SeenHashes<L>,
298) -> Result<Vec<Vec<u8>>, TrieHash<L>, CError<L>>
299where
300 L: TrieLayout,
301{
302 encode_compact_inner(db, Some(seen))
303}
304
305fn encode_compact_inner<L>(
306 db: &TrieDB<L>,
307 mut seen: Option<&mut SeenHashes<L>>,
308) -> Result<Vec<Vec<u8>>, TrieHash<L>, CError<L>>
309where
310 L: TrieLayout,
311{
312 let mut output = Vec::new();
313
314 let mut stack: Vec<EncoderStackEntry<L::Codec>> = Vec::new();
317
318 let mut iter = TrieDBRawIterator::new(db)?;
323
324 while let Some(item) = iter.next_raw_item(db, true) {
329 match item {
330 Ok((prefix, node_hash, node)) => {
331 let Some(node_hash) = node_hash else { continue };
334
335 if let Some(seen) = seen.as_deref_mut() {
336 let is_root = stack.is_empty();
337 if !is_root && seen.nodes.contains(node_hash) {
344 iter.skip_current_subtree();
345 continue
346 }
347 seen.nodes.insert(*node_hash);
348 }
349
350 while let Some(mut last_entry) = stack.pop() {
355 if prefix.starts_with(&last_entry.prefix) {
356 last_entry.advance_child_index(&prefix).expect(
359 "all errors from advance_child_index indicate bugs with \
360 TrieDBRawIterator or this function",
361 );
362 last_entry.omit_children[last_entry.child_index] = true;
363 last_entry.child_index += 1;
364 stack.push(last_entry);
365 break
366 } else {
367 output[last_entry.output_index] = last_entry.encode_node()?;
368 }
369 }
370
371 let (children_len, detached_value) = match node.node_plan() {
372 NodePlan::Empty => (0, None),
373 NodePlan::Leaf { value, .. } => (
374 0,
375 detached_value(
376 db,
377 value,
378 node.data(),
379 prefix.as_prefix(),
380 seen.as_deref_mut(),
381 ),
382 ),
383 NodePlan::Extension { .. } => (1, None),
384 NodePlan::NibbledBranch { value: Some(value), .. } |
385 NodePlan::Branch { value: Some(value), .. } => (
386 NIBBLE_LENGTH,
387 detached_value(
388 db,
389 value,
390 node.data(),
391 prefix.as_prefix(),
392 seen.as_deref_mut(),
393 ),
394 ),
395 NodePlan::NibbledBranch { value: None, .. } |
396 NodePlan::Branch { value: None, .. } => (NIBBLE_LENGTH, None),
397 };
398
399 stack.push(EncoderStackEntry {
400 prefix: prefix.clone(),
401 node: node.clone(),
402 child_index: 0,
403 omit_children: vec![false; children_len],
404 omit_value: detached_value.is_some(),
405 output_index: output.len(),
406 _marker: PhantomData::default(),
407 });
408 output.push(Vec::new());
411 if let Some(value) = detached_value {
412 output.push(value);
413 }
414 },
415 Err(err) => match *err {
416 TrieError::IncompleteDatabase(_) => {},
420 _ => return Err(err),
421 },
422 }
423 }
424
425 while let Some(mut entry) = stack.pop() {
426 output[entry.output_index] = entry.encode_node()?;
427 }
428
429 Ok(output)
430}
431
432struct DecoderStackEntry<'a, C: NodeCodec> {
433 node: Node<'a>,
434 child_index: usize,
437 children: Vec<Option<ChildReference<C::HashOut>>>,
439 attached_value: Option<&'a [u8]>,
441 _marker: PhantomData<C>,
442}
443
444impl<'a, C: NodeCodec> DecoderStackEntry<'a, C> {
445 fn advance_child_index(&mut self) -> Result<bool, C::HashOut, C::Error> {
454 match self.node {
455 Node::Extension(_, child) if self.child_index == 0 => {
456 match child {
457 NodeHandle::Inline(data) if data.is_empty() => return Ok(false),
458 _ => {
459 let child_ref = child.try_into().map_err(|hash| {
460 Box::new(TrieError::InvalidHash(C::HashOut::default(), hash))
461 })?;
462 self.children[self.child_index] = Some(child_ref);
463 },
464 }
465 self.child_index += 1;
466 },
467 Node::Branch(children, _) | Node::NibbledBranch(_, children, _) => {
468 while self.child_index < NIBBLE_LENGTH {
469 match children[self.child_index] {
470 Some(NodeHandle::Inline(data)) if data.is_empty() => return Ok(false),
471 Some(child) => {
472 let child_ref = child.try_into().map_err(|hash| {
473 Box::new(TrieError::InvalidHash(C::HashOut::default(), hash))
474 })?;
475 self.children[self.child_index] = Some(child_ref);
476 },
477 None => {},
478 }
479 self.child_index += 1;
480 }
481 },
482 _ => {},
483 }
484 Ok(true)
485 }
486
487 fn push_to_prefix(&self, prefix: &mut NibbleVec) {
490 match self.node {
491 Node::Empty => {},
492 Node::Leaf(partial, _) | Node::Extension(partial, _) => {
493 prefix.append_partial(partial.right());
494 },
495 Node::Branch(_, _) => {
496 prefix.push(self.child_index as u8);
497 },
498 Node::NibbledBranch(partial, _, _) => {
499 prefix.append_partial(partial.right());
500 prefix.push(self.child_index as u8);
501 },
502 }
503 }
504
505 fn pop_from_prefix(&self, prefix: &mut NibbleVec) {
508 match self.node {
509 Node::Empty => {},
510 Node::Leaf(partial, _) | Node::Extension(partial, _) => {
511 prefix.drop_lasts(partial.len());
512 },
513 Node::Branch(_, _) => {
514 prefix.pop();
515 },
516 Node::NibbledBranch(partial, _, _) => {
517 prefix.pop();
518 prefix.drop_lasts(partial.len());
519 },
520 }
521 }
522
523 fn encode_node(self, attached_hash: Option<&[u8]>) -> Vec<u8> {
528 let attached_hash = attached_hash.map(|h| crate::node::Value::Node(h));
529 match self.node {
530 Node::Empty => C::empty_node().to_vec(),
531 Node::Leaf(partial, value) =>
532 C::leaf_node(partial.right_iter(), partial.len(), attached_hash.unwrap_or(value)),
533 Node::Extension(partial, _) => C::extension_node(
534 partial.right_iter(),
535 partial.len(),
536 self.children[0].expect("required by method precondition; qed"),
537 ),
538 Node::Branch(_, value) => C::branch_node(
539 self.children.into_iter(),
540 if attached_hash.is_some() { attached_hash } else { value },
541 ),
542 Node::NibbledBranch(partial, _, value) => C::branch_node_nibbled(
543 partial.right_iter(),
544 partial.len(),
545 self.children.iter(),
546 if attached_hash.is_some() { attached_hash } else { value },
547 ),
548 }
549 }
550}
551
552pub fn decode_compact<L, DB>(
566 db: &mut DB,
567 encoded: &[Vec<u8>],
568) -> Result<(TrieHash<L>, usize), TrieHash<L>, CError<L>>
569where
570 L: TrieLayout,
571 DB: HashDB<L::Hash, DBValue>,
572{
573 decode_compact_from_iter::<L, DB, _>(db, encoded.iter().map(Vec::as_slice))
574}
575
576pub fn decode_compact_from_iter<'a, L, DB, I>(
578 db: &mut DB,
579 encoded: I,
580) -> Result<(TrieHash<L>, usize), TrieHash<L>, CError<L>>
581where
582 L: TrieLayout,
583 DB: HashDB<L::Hash, DBValue>,
584 I: IntoIterator<Item = &'a [u8]>,
585{
586 let mut stack: Vec<DecoderStackEntry<L::Codec>> = Vec::new();
589
590 let mut prefix = NibbleVec::new();
592
593 let mut iter = encoded.into_iter().enumerate();
594 while let Some((i, encoded_node)) = iter.next() {
595 let mut attached_node = 0;
596 if let Some(header) = L::Codec::ESCAPE_HEADER {
597 if encoded_node.starts_with(&[header]) {
598 attached_node = 1;
599 }
600 }
601 let node = L::Codec::decode(&encoded_node[attached_node..])
602 .map_err(|err| Box::new(TrieError::DecoderError(<TrieHash<L>>::default(), err)))?;
603
604 let children_len = match node {
605 Node::Empty | Node::Leaf(..) => 0,
606 Node::Extension(..) => 1,
607 Node::Branch(..) | Node::NibbledBranch(..) => NIBBLE_LENGTH,
608 };
609 let mut last_entry = DecoderStackEntry {
610 node,
611 child_index: 0,
612 children: vec![None; children_len],
613 attached_value: None,
614 _marker: PhantomData::default(),
615 };
616
617 if attached_node > 0 {
618 if let Some((_, fetched_value)) = iter.next() {
620 last_entry.attached_value = Some(fetched_value);
621 } else {
622 return Err(Box::new(TrieError::IncompleteDatabase(<TrieHash<L>>::default())))
623 }
624 }
625
626 loop {
627 if !last_entry.advance_child_index()? {
628 last_entry.push_to_prefix(&mut prefix);
629 stack.push(last_entry);
630 break
631 }
632
633 let hash = last_entry.attached_value.as_ref().map(|value| {
636 let partial_prefix_len = match &last_entry.node {
637 Node::Leaf(partial, _) | Node::NibbledBranch(partial, _, _) => {
638 prefix.append_partial(partial.right());
639 partial.len()
640 },
641 _ => 0,
642 };
643 let hash = db.insert(prefix.as_prefix(), value);
644 prefix.drop_lasts(partial_prefix_len);
645 hash
646 });
647 let node_data = last_entry.encode_node(hash.as_ref().map(|h| h.as_ref()));
648 let node_hash = db.insert(prefix.as_prefix(), node_data.as_ref());
649
650 if let Some(entry) = stack.pop() {
651 last_entry = entry;
652 last_entry.pop_from_prefix(&mut prefix);
653 last_entry.children[last_entry.child_index] = Some(ChildReference::Hash(node_hash));
654 last_entry.child_index += 1;
655 } else {
656 return Ok((node_hash, i + 1 + attached_node))
657 }
658 }
659 }
660
661 Err(Box::new(TrieError::IncompleteDatabase(<TrieHash<L>>::default())))
662}