1use hashbrown::{hash_map::Entry, HashMap};
18use parity_scale_codec::{Compact, Decode, Encode, Error as CodecError, Input, Output};
19use std::{borrow::Borrow, fmt, iter::once, marker::PhantomData, ops::Range};
20use trie_db_fun::{
21 nibble_ops,
22 node::{NibbleSlicePlan, NodeHandlePlan, NodeOwned, NodePlan, Value, ValuePlan},
23 trie_visit,
24 triedbmut::ChildReference,
25 DBValue, NodeCodec, Trie, TrieBuilder, TrieConfiguration, TrieDBBuilder, TrieDBMutBuilder,
26 TrieHash, TrieLayout, TrieMut, TrieRoot,
27};
28pub use trie_root::TrieStream;
29use trie_root::{Hasher, Value as TrieStreamValue};
30
31mod substrate;
32mod substrate_like;
33pub mod node {
34 pub use trie_db_fun::node::Node;
35}
36
37pub use substrate_like::{
38 trie_constants, HashedValueNoExt, HashedValueNoExtThreshold,
39 NodeCodec as ReferenceNodeCodecNoExtMeta, ReferenceTrieStreamNoExt,
40};
41
42pub use paste::paste;
43pub use substrate::{LayoutV0 as SubstrateV0, LayoutV1 as SubstrateV1};
44
45pub type RefHasher = keccak_hasher::KeccakHasher;
47
48#[macro_export]
50macro_rules! test_layouts {
51 ($test:ident, $test_internal:ident) => {
52 #[test]
53 fn $test() {
54 eprintln!("Running with layout `HashedValueNoExtThreshold`");
55 $test_internal::<$crate::HashedValueNoExtThreshold<1>>();
56 eprintln!("Running with layout `HashedValueNoExt`");
57 $test_internal::<$crate::HashedValueNoExt>();
58 eprintln!("Running with layout `NoExtensionLayout`");
59 $test_internal::<$crate::NoExtensionLayout>();
60 eprintln!("Running with layout `ExtensionLayout`");
61 $test_internal::<$crate::ExtensionLayout>();
62 }
63 };
64}
65
66#[macro_export]
67macro_rules! test_layouts_substrate {
68 ($test:ident) => {
69 $crate::paste! {
70 #[test]
71 fn [<$test _substrate_v0>]() {
72 $test::<$crate::SubstrateV0<$crate::RefHasher>>();
73 }
74 #[test]
75 fn [<$test _substrate_v1>]() {
76 $test::<$crate::SubstrateV1<$crate::RefHasher>>();
77 }
78 }
79 };
80}
81
82#[macro_export]
84macro_rules! test_layouts_no_meta {
85 ($test:ident, $test_internal:ident) => {
86 #[test]
87 fn $test() {
88 $test_internal::<$crate::NoExtensionLayout>();
89 $test_internal::<$crate::ExtensionLayout>();
90 }
91 };
92}
93
94#[derive(Default, Clone)]
96pub struct ExtensionLayout;
97
98impl TrieLayout for ExtensionLayout {
99 const USE_EXTENSION: bool = true;
100 const ALLOW_EMPTY: bool = false;
101 const MAX_INLINE_VALUE: Option<u32> = None;
102 type Hash = RefHasher;
103 type Codec = ReferenceNodeCodec<RefHasher>;
104}
105
106impl TrieConfiguration for ExtensionLayout {}
107
108pub struct GenericNoExtensionLayout<H>(PhantomData<H>);
111
112impl<H> Default for GenericNoExtensionLayout<H> {
113 fn default() -> Self {
114 GenericNoExtensionLayout(PhantomData)
115 }
116}
117
118impl<H> Clone for GenericNoExtensionLayout<H> {
119 fn clone(&self) -> Self {
120 GenericNoExtensionLayout(PhantomData)
121 }
122}
123
124impl<H: Hasher> TrieLayout for GenericNoExtensionLayout<H> {
125 const USE_EXTENSION: bool = false;
126 const ALLOW_EMPTY: bool = false;
127 const MAX_INLINE_VALUE: Option<u32> = None;
128 type Hash = H;
129 type Codec = ReferenceNodeCodecNoExt<H>;
130}
131
132#[derive(Default, Clone)]
134pub struct AllowEmptyLayout;
135
136impl TrieLayout for AllowEmptyLayout {
137 const USE_EXTENSION: bool = true;
138 const ALLOW_EMPTY: bool = true;
139 const MAX_INLINE_VALUE: Option<u32> = None;
140 type Hash = RefHasher;
141 type Codec = ReferenceNodeCodec<RefHasher>;
142}
143
144impl<H: Hasher> TrieConfiguration for GenericNoExtensionLayout<H> {}
145
146pub type NoExtensionLayout = GenericNoExtensionLayout<RefHasher>;
148
149pub struct Bitmap(u16);
151
152const BITMAP_LENGTH: usize = 2;
153
154impl Bitmap {
155 fn decode(data: &[u8]) -> Result<Self, CodecError> {
156 Ok(u16::decode(&mut &data[..]).map(|v| Bitmap(v))?)
157 }
158
159 fn value_at(&self, i: usize) -> bool {
160 self.0 & (1u16 << i) != 0
161 }
162
163 fn encode<I: Iterator<Item = bool>>(has_children: I, output: &mut [u8]) {
164 let mut bitmap: u16 = 0;
165 let mut cursor: u16 = 1;
166 for v in has_children {
167 if v {
168 bitmap |= cursor
169 }
170 cursor <<= 1;
171 }
172 output[0] = (bitmap % 256) as u8;
173 output[1] = (bitmap / 256) as u8;
174 }
175}
176
177pub type RefTrieDB<'a, 'cache> = trie_db_fun::TrieDB<'a, 'cache, ExtensionLayout>;
178pub type RefTrieDBBuilder<'a, 'cache> = trie_db_fun::TrieDBBuilder<'a, 'cache, ExtensionLayout>;
179pub type RefTrieDBMut<'a> = trie_db_fun::TrieDBMut<'a, ExtensionLayout>;
180pub type RefTrieDBMutBuilder<'a> = trie_db_fun::TrieDBMutBuilder<'a, ExtensionLayout>;
181pub type RefTrieDBMutNoExt<'a> = trie_db_fun::TrieDBMut<'a, NoExtensionLayout>;
182pub type RefTrieDBMutNoExtBuilder<'a> = trie_db_fun::TrieDBMutBuilder<'a, NoExtensionLayout>;
183pub type RefTrieDBMutAllowEmpty<'a> = trie_db_fun::TrieDBMut<'a, AllowEmptyLayout>;
184pub type RefTrieDBMutAllowEmptyBuilder<'a> = trie_db_fun::TrieDBMutBuilder<'a, AllowEmptyLayout>;
185pub type RefTestTrieDBCache = TestTrieCache<ExtensionLayout>;
186pub type RefTestTrieDBCacheNoExt = TestTrieCache<NoExtensionLayout>;
187pub type RefFatDB<'a, 'cache> = trie_db_fun::FatDB<'a, 'cache, ExtensionLayout>;
188pub type RefFatDBMut<'a> = trie_db_fun::FatDBMut<'a, ExtensionLayout>;
189pub type RefSecTrieDB<'a, 'cache> = trie_db_fun::SecTrieDB<'a, 'cache, ExtensionLayout>;
190pub type RefSecTrieDBMut<'a> = trie_db_fun::SecTrieDBMut<'a, ExtensionLayout>;
191pub type RefLookup<'a, 'cache, Q> = trie_db_fun::Lookup<'a, 'cache, ExtensionLayout, Q>;
192pub type RefLookupNoExt<'a, 'cache, Q> = trie_db_fun::Lookup<'a, 'cache, NoExtensionLayout, Q>;
193
194pub fn reference_trie_root<T: TrieLayout, I, A, B>(input: I) -> <T::Hash as Hasher>::Out
195where
196 I: IntoIterator<Item = (A, B)>,
197 A: AsRef<[u8]> + Ord + fmt::Debug,
198 B: AsRef<[u8]> + fmt::Debug,
199{
200 if T::USE_EXTENSION {
201 trie_root::trie_root::<T::Hash, ReferenceTrieStream, _, _, _>(input, T::MAX_INLINE_VALUE)
202 } else {
203 trie_root::trie_root_no_extension::<T::Hash, ReferenceTrieStreamNoExt, _, _, _>(
204 input,
205 T::MAX_INLINE_VALUE,
206 )
207 }
208}
209
210fn data_sorted_unique<I, A: Ord, B>(input: I) -> Vec<(A, B)>
211where
212 I: IntoIterator<Item = (A, B)>,
213{
214 let mut m = std::collections::BTreeMap::new();
215 for (k, v) in input {
216 let _ = m.insert(k, v); }
218 m.into_iter().collect()
219}
220
221pub fn reference_trie_root_iter_build<T, I, A, B>(input: I) -> <T::Hash as Hasher>::Out
222where
223 T: TrieLayout,
224 I: IntoIterator<Item = (A, B)>,
225 A: AsRef<[u8]> + Ord + fmt::Debug,
226 B: AsRef<[u8]> + fmt::Debug,
227{
228 let mut cb = trie_db_fun::TrieRoot::<T>::default();
229 trie_visit::<T, _, _, _, _>(data_sorted_unique(input), &mut cb);
230 cb.root.unwrap_or_default()
231}
232
233fn reference_trie_root_unhashed<I, A, B>(input: I) -> Vec<u8>
234where
235 I: IntoIterator<Item = (A, B)>,
236 A: AsRef<[u8]> + Ord + fmt::Debug,
237 B: AsRef<[u8]> + fmt::Debug,
238{
239 trie_root::unhashed_trie::<RefHasher, ReferenceTrieStream, _, _, _>(input, Default::default())
240}
241
242fn reference_trie_root_unhashed_no_extension<I, A, B>(input: I) -> Vec<u8>
243where
244 I: IntoIterator<Item = (A, B)>,
245 A: AsRef<[u8]> + Ord + fmt::Debug,
246 B: AsRef<[u8]> + fmt::Debug,
247{
248 trie_root::unhashed_trie_no_extension::<RefHasher, ReferenceTrieStreamNoExt, _, _, _>(
249 input,
250 Default::default(),
251 )
252}
253
254const EMPTY_TRIE: u8 = 0;
255const LEAF_NODE_OFFSET: u8 = 1;
256const EXTENSION_NODE_OFFSET: u8 = 128;
257const BRANCH_NODE_NO_VALUE: u8 = 254;
258const BRANCH_NODE_WITH_VALUE: u8 = 255;
259const LEAF_NODE_OVER: u8 = EXTENSION_NODE_OFFSET - LEAF_NODE_OFFSET;
260const EXTENSION_NODE_OVER: u8 = BRANCH_NODE_NO_VALUE - EXTENSION_NODE_OFFSET;
261const LEAF_NODE_LAST: u8 = EXTENSION_NODE_OFFSET - 1;
262const EXTENSION_NODE_LAST: u8 = BRANCH_NODE_NO_VALUE - 1;
263
264const NIBBLE_SIZE_BOUND_NO_EXT: usize = u16::max_value() as usize;
266const FIRST_PREFIX: u8 = 0b_00 << 6;
267const LEAF_PREFIX_MASK_NO_EXT: u8 = 0b_01 << 6;
268const BRANCH_WITHOUT_MASK_NO_EXT: u8 = 0b_10 << 6;
269const BRANCH_WITH_MASK_NO_EXT: u8 = 0b_11 << 6;
270const EMPTY_TRIE_NO_EXT: u8 = FIRST_PREFIX | 0b_00;
271
272fn fuse_nibbles_node<'a>(nibbles: &'a [u8], leaf: bool) -> impl Iterator<Item = u8> + 'a {
276 debug_assert!(
277 nibbles.len() < LEAF_NODE_OVER.min(EXTENSION_NODE_OVER) as usize,
278 "nibbles length too long. what kind of size of key are you trying to include in the trie!?!"
279 );
280 let first_byte =
281 if leaf { LEAF_NODE_OFFSET } else { EXTENSION_NODE_OFFSET } + nibbles.len() as u8;
282
283 once(first_byte)
284 .chain(if nibbles.len() % 2 == 1 { Some(nibbles[0]) } else { None })
285 .chain(nibbles[nibbles.len() % 2..].chunks(2).map(|ch| ch[0] << 4 | ch[1]))
286}
287
288enum NodeKindNoExt {
289 Leaf,
290 BranchNoValue,
291 BranchWithValue,
292}
293
294fn branch_node(has_value: bool, has_children: impl Iterator<Item = bool>) -> [u8; 3] {
297 let mut result = [0, 0, 0];
298 branch_node_buffered(has_value, has_children, &mut result[..]);
299 result
300}
301
302fn branch_node_buffered<I: Iterator<Item = bool>>(
305 has_value: bool,
306 has_children: I,
307 output: &mut [u8],
308) {
309 let first = if has_value { BRANCH_NODE_WITH_VALUE } else { BRANCH_NODE_NO_VALUE };
310 output[0] = first;
311 Bitmap::encode(has_children, &mut output[1..]);
312}
313
314fn branch_node_bit_mask(has_children: impl Iterator<Item = bool>) -> (u8, u8) {
317 let mut bitmap: u16 = 0;
318 let mut cursor: u16 = 1;
319 for v in has_children {
320 if v {
321 bitmap |= cursor
322 }
323 cursor <<= 1;
324 }
325 ((bitmap % 256) as u8, (bitmap / 256) as u8)
326}
327
328#[derive(Default, Clone)]
330pub struct ReferenceTrieStream {
331 buffer: Vec<u8>,
332}
333
334impl TrieStream for ReferenceTrieStream {
335 fn new() -> Self {
336 ReferenceTrieStream { buffer: Vec::new() }
337 }
338
339 fn append_empty_data(&mut self) {
340 self.buffer.push(EMPTY_TRIE);
341 }
342
343 fn append_leaf(&mut self, key: &[u8], value: TrieStreamValue) {
344 if let TrieStreamValue::Inline(value) = value {
345 self.buffer.extend(fuse_nibbles_node(key, true));
346 value.encode_to(&mut self.buffer);
347 } else {
348 unreachable!("This stream do not allow external value node")
349 }
350 }
351
352 fn begin_branch(
353 &mut self,
354 maybe_key: Option<&[u8]>,
355 maybe_value: Option<TrieStreamValue>,
356 has_children: impl Iterator<Item = bool>,
357 ) {
358 self.buffer.extend(&branch_node(!matches!(maybe_value, None), has_children));
359 if let Some(partial) = maybe_key {
360 self.buffer.extend(fuse_nibbles_node(partial, false));
362 }
363 if let Some(TrieStreamValue::Inline(value)) = maybe_value {
364 value.encode_to(&mut self.buffer);
365 }
366 }
367
368 fn append_extension(&mut self, key: &[u8]) {
369 self.buffer.extend(fuse_nibbles_node(key, false));
370 }
371
372 fn append_substream<H: Hasher>(&mut self, other: Self) {
373 let data = other.out();
374 match data.len() {
375 0..=31 => data.encode_to(&mut self.buffer),
376 _ => H::hash(&data).as_ref().encode_to(&mut self.buffer),
377 }
378 }
379
380 fn out(self) -> Vec<u8> {
381 self.buffer
382 }
383}
384
385#[derive(Copy, Clone, PartialEq, Eq, Debug)]
387enum NodeHeader {
388 Null,
389 Branch(bool),
390 Extension(usize),
391 Leaf(usize),
392}
393
394#[derive(Copy, Clone, PartialEq, Eq, Debug)]
396enum NodeHeaderNoExt {
397 Null,
398 Branch(bool, usize),
399 Leaf(usize),
400}
401
402impl Encode for NodeHeader {
403 fn encode_to<T: Output + ?Sized>(&self, output: &mut T) {
404 match self {
405 NodeHeader::Null => output.push_byte(EMPTY_TRIE),
406 NodeHeader::Branch(true) => output.push_byte(BRANCH_NODE_WITH_VALUE),
407 NodeHeader::Branch(false) => output.push_byte(BRANCH_NODE_NO_VALUE),
408 NodeHeader::Leaf(nibble_count) =>
409 output.push_byte(LEAF_NODE_OFFSET + *nibble_count as u8),
410 NodeHeader::Extension(nibble_count) =>
411 output.push_byte(EXTENSION_NODE_OFFSET + *nibble_count as u8),
412 }
413 }
414}
415
416fn size_and_prefix_iterator(size: usize, prefix: u8) -> impl Iterator<Item = u8> {
419 let size = ::std::cmp::min(NIBBLE_SIZE_BOUND_NO_EXT, size);
420
421 let l1 = std::cmp::min(62, size);
422 let (first_byte, mut rem) =
423 if size == l1 { (once(prefix + l1 as u8), 0) } else { (once(prefix + 63), size - l1) };
424 let next_bytes = move || {
425 if rem > 0 {
426 if rem < 256 {
427 let result = rem - 1;
428 rem = 0;
429 Some(result as u8)
430 } else {
431 rem = rem.saturating_sub(255);
432 Some(255)
433 }
434 } else {
435 None
436 }
437 };
438 first_byte.chain(::std::iter::from_fn(next_bytes))
439}
440
441fn encode_size_and_prefix(size: usize, prefix: u8, out: &mut (impl Output + ?Sized)) {
442 for b in size_and_prefix_iterator(size, prefix) {
443 out.push_byte(b)
444 }
445}
446
447fn decode_size<I: Input>(first: u8, input: &mut I) -> Result<usize, CodecError> {
448 let mut result = (first & 255u8 >> 2) as usize;
449 if result < 63 {
450 return Ok(result)
451 }
452 result -= 1;
453 while result <= NIBBLE_SIZE_BOUND_NO_EXT {
454 let n = input.read_byte()? as usize;
455 if n < 255 {
456 return Ok(result + n + 1)
457 }
458 result += 255;
459 }
460 Err("Size limit reached for a nibble slice".into())
461}
462
463impl Encode for NodeHeaderNoExt {
464 fn encode_to<T: Output + ?Sized>(&self, output: &mut T) {
465 match self {
466 NodeHeaderNoExt::Null => output.push_byte(EMPTY_TRIE_NO_EXT),
467 NodeHeaderNoExt::Branch(true, nibble_count) =>
468 encode_size_and_prefix(*nibble_count, BRANCH_WITH_MASK_NO_EXT, output),
469 NodeHeaderNoExt::Branch(false, nibble_count) =>
470 encode_size_and_prefix(*nibble_count, BRANCH_WITHOUT_MASK_NO_EXT, output),
471 NodeHeaderNoExt::Leaf(nibble_count) =>
472 encode_size_and_prefix(*nibble_count, LEAF_PREFIX_MASK_NO_EXT, output),
473 }
474 }
475}
476
477impl Decode for NodeHeader {
478 fn decode<I: Input>(input: &mut I) -> Result<Self, CodecError> {
479 Ok(match input.read_byte()? {
480 EMPTY_TRIE => NodeHeader::Null,
481 BRANCH_NODE_NO_VALUE => NodeHeader::Branch(false),
482 BRANCH_NODE_WITH_VALUE => NodeHeader::Branch(true),
483 i @ LEAF_NODE_OFFSET..=LEAF_NODE_LAST =>
484 NodeHeader::Leaf((i - LEAF_NODE_OFFSET) as usize),
485 i @ EXTENSION_NODE_OFFSET..=EXTENSION_NODE_LAST =>
486 NodeHeader::Extension((i - EXTENSION_NODE_OFFSET) as usize),
487 })
488 }
489}
490
491impl Decode for NodeHeaderNoExt {
492 fn decode<I: Input>(input: &mut I) -> Result<Self, CodecError> {
493 let i = input.read_byte()?;
494 if i == EMPTY_TRIE_NO_EXT {
495 return Ok(NodeHeaderNoExt::Null)
496 }
497 match i & (0b11 << 6) {
498 LEAF_PREFIX_MASK_NO_EXT => Ok(NodeHeaderNoExt::Leaf(decode_size(i, input)?)),
499 BRANCH_WITHOUT_MASK_NO_EXT =>
500 Ok(NodeHeaderNoExt::Branch(false, decode_size(i, input)?)),
501 BRANCH_WITH_MASK_NO_EXT => Ok(NodeHeaderNoExt::Branch(true, decode_size(i, input)?)),
502 _ => Err("Unknown type of node".into()),
504 }
505 }
506}
507
508#[derive(Default, Clone)]
510pub struct ReferenceNodeCodec<H>(PhantomData<H>);
511
512#[derive(Default, Clone)]
517pub struct ReferenceNodeCodecNoExt<H>(PhantomData<H>);
518
519fn partial_from_iterator_to_key<I: Iterator<Item = u8>>(
520 partial: I,
521 nibble_count: usize,
522 offset: u8,
523 over: u8,
524) -> Vec<u8> {
525 assert!(nibble_count < over as usize);
526 let mut output = Vec::with_capacity(1 + (nibble_count / nibble_ops::NIBBLE_PER_BYTE));
527 output.push(offset + nibble_count as u8);
528 output.extend(partial);
529 output
530}
531
532fn partial_from_iterator_encode<I: Iterator<Item = u8>>(
533 partial: I,
534 nibble_count: usize,
535 node_kind: NodeKindNoExt,
536) -> Vec<u8> {
537 let nibble_count = ::std::cmp::min(NIBBLE_SIZE_BOUND_NO_EXT, nibble_count);
538
539 let mut output = Vec::with_capacity(3 + (nibble_count / nibble_ops::NIBBLE_PER_BYTE));
540 match node_kind {
541 NodeKindNoExt::Leaf => NodeHeaderNoExt::Leaf(nibble_count).encode_to(&mut output),
542 NodeKindNoExt::BranchWithValue =>
543 NodeHeaderNoExt::Branch(true, nibble_count).encode_to(&mut output),
544 NodeKindNoExt::BranchNoValue =>
545 NodeHeaderNoExt::Branch(false, nibble_count).encode_to(&mut output),
546 };
547 output.extend(partial);
548 output
549}
550
551struct ByteSliceInput<'a> {
552 data: &'a [u8],
553 offset: usize,
554}
555
556impl<'a> ByteSliceInput<'a> {
557 fn new(data: &'a [u8]) -> Self {
558 ByteSliceInput { data, offset: 0 }
559 }
560
561 fn take(&mut self, count: usize) -> Result<Range<usize>, CodecError> {
562 if self.offset + count > self.data.len() {
563 return Err("out of data".into())
564 }
565
566 let range = self.offset..(self.offset + count);
567 self.offset += count;
568 Ok(range)
569 }
570}
571
572impl<'a> Input for ByteSliceInput<'a> {
573 fn remaining_len(&mut self) -> Result<Option<usize>, CodecError> {
574 let remaining =
575 if self.offset <= self.data.len() { Some(self.data.len() - self.offset) } else { None };
576 Ok(remaining)
577 }
578
579 fn read(&mut self, into: &mut [u8]) -> Result<(), CodecError> {
580 let range = self.take(into.len())?;
581 into.copy_from_slice(&self.data[range]);
582 Ok(())
583 }
584
585 fn read_byte(&mut self) -> Result<u8, CodecError> {
586 if self.offset + 1 > self.data.len() {
587 return Err("out of data".into())
588 }
589
590 let byte = self.data[self.offset];
591 self.offset += 1;
592 Ok(byte)
593 }
594}
595
596impl<H: Hasher> NodeCodec for ReferenceNodeCodec<H> {
602 type Error = CodecError;
603 type HashOut = H::Out;
604
605 fn hashed_null_node() -> <H as Hasher>::Out {
606 H::hash(<Self as NodeCodec>::empty_node())
607 }
608
609 fn decode_plan(data: &[u8]) -> ::std::result::Result<NodePlan, Self::Error> {
610 let mut input = ByteSliceInput::new(data);
611 match NodeHeader::decode(&mut input)? {
612 NodeHeader::Null => Ok(NodePlan::Empty),
613 NodeHeader::Branch(has_value) => {
614 let bitmap_range = input.take(BITMAP_LENGTH)?;
615 let bitmap = Bitmap::decode(&data[bitmap_range])?;
616
617 let value = if has_value {
618 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
619 Some(ValuePlan::Inline(input.take(count)?))
620 } else {
621 None
622 };
623 let mut children = [
624 None, None, None, None, None, None, None, None, None, None, None, None, None,
625 None, None, None,
626 ];
627 for i in 0..nibble_ops::NIBBLE_LENGTH {
628 if bitmap.value_at(i) {
629 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
630 let range = input.take(count)?;
631 children[i] = Some(if count == H::LENGTH {
632 NodeHandlePlan::Hash(range)
633 } else {
634 NodeHandlePlan::Inline(range)
635 });
636 }
637 }
638 Ok(NodePlan::Branch { value, children })
639 },
640 NodeHeader::Extension(nibble_count) => {
641 let partial = input.take(
642 (nibble_count + (nibble_ops::NIBBLE_PER_BYTE - 1)) /
643 nibble_ops::NIBBLE_PER_BYTE,
644 )?;
645 let partial_padding = nibble_ops::number_padding(nibble_count);
646 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
647 let range = input.take(count)?;
648 let child = if count == H::LENGTH {
649 NodeHandlePlan::Hash(range)
650 } else {
651 NodeHandlePlan::Inline(range)
652 };
653 Ok(NodePlan::Extension {
654 partial: NibbleSlicePlan::new(partial, partial_padding),
655 child,
656 })
657 },
658 NodeHeader::Leaf(nibble_count) => {
659 let partial = input.take(
660 (nibble_count + (nibble_ops::NIBBLE_PER_BYTE - 1)) /
661 nibble_ops::NIBBLE_PER_BYTE,
662 )?;
663 let partial_padding = nibble_ops::number_padding(nibble_count);
664 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
665 let value = input.take(count)?;
666 Ok(NodePlan::Leaf {
667 partial: NibbleSlicePlan::new(partial, partial_padding),
668 value: ValuePlan::Inline(value),
669 })
670 },
671 }
672 }
673
674 fn is_empty_node(data: &[u8]) -> bool {
675 data == <Self as NodeCodec>::empty_node()
676 }
677
678 fn empty_node() -> &'static [u8] {
679 &[EMPTY_TRIE]
680 }
681
682 fn leaf_node(partial: impl Iterator<Item = u8>, number_nibble: usize, value: Value) -> Vec<u8> {
683 let mut output =
684 partial_from_iterator_to_key(partial, number_nibble, LEAF_NODE_OFFSET, LEAF_NODE_OVER);
685 match value {
686 Value::Inline(value) => {
687 Compact(value.len() as u32).encode_to(&mut output);
688 output.extend_from_slice(value);
689 },
690 _ => unimplemented!("unsupported"),
691 }
692 output
693 }
694
695 fn extension_node(
696 partial: impl Iterator<Item = u8>,
697 number_nibble: usize,
698 child: ChildReference<Self::HashOut>,
699 ) -> Vec<u8> {
700 let mut output = partial_from_iterator_to_key(
701 partial,
702 number_nibble,
703 EXTENSION_NODE_OFFSET,
704 EXTENSION_NODE_OVER,
705 );
706 match child {
707 ChildReference::Hash(h) => h.as_ref().encode_to(&mut output),
708 ChildReference::Inline(inline_data, len) =>
709 (&AsRef::<[u8]>::as_ref(&inline_data)[..len]).encode_to(&mut output),
710 };
711 output
712 }
713
714 fn branch_node(
715 children: impl Iterator<Item = impl Borrow<Option<ChildReference<Self::HashOut>>>>,
716 maybe_value: Option<Value>,
717 ) -> Vec<u8> {
718 let mut output = vec![0; BITMAP_LENGTH + 1];
719 let mut prefix: [u8; 3] = [0; 3];
720 let have_value = match maybe_value {
721 Some(Value::Inline(value)) => {
722 Compact(value.len() as u32).encode_to(&mut output);
723 output.extend_from_slice(value);
724 true
725 },
726 None => false,
727 _ => unimplemented!("unsupported"),
728 };
729 let has_children = children.map(|maybe_child| match maybe_child.borrow() {
730 Some(ChildReference::Hash(h)) => {
731 h.as_ref().encode_to(&mut output);
732 true
733 },
734 &Some(ChildReference::Inline(inline_data, len)) => {
735 inline_data.as_ref()[..len].encode_to(&mut output);
736 true
737 },
738 None => false,
739 });
740 branch_node_buffered(have_value, has_children, prefix.as_mut());
741 output[0..BITMAP_LENGTH + 1].copy_from_slice(prefix.as_ref());
742 output
743 }
744
745 fn branch_node_nibbled(
746 _partial: impl Iterator<Item = u8>,
747 _number_nibble: usize,
748 _children: impl Iterator<Item = impl Borrow<Option<ChildReference<Self::HashOut>>>>,
749 _maybe_value: Option<Value>,
750 ) -> Vec<u8> {
751 unreachable!("codec with extension branch")
752 }
753}
754
755impl<H: Hasher> NodeCodec for ReferenceNodeCodecNoExt<H> {
756 type Error = CodecError;
757 type HashOut = <H as Hasher>::Out;
758
759 fn hashed_null_node() -> <H as Hasher>::Out {
760 H::hash(<Self as NodeCodec>::empty_node())
761 }
762
763 fn decode_plan(data: &[u8]) -> Result<NodePlan, Self::Error> {
764 if data.len() < 1 {
765 return Err(CodecError::from("Empty encoded node."))
766 }
767 let mut input = ByteSliceInput::new(data);
768
769 Ok(match NodeHeaderNoExt::decode(&mut input)? {
770 NodeHeaderNoExt::Null => NodePlan::Empty,
771 NodeHeaderNoExt::Branch(has_value, nibble_count) => {
772 let padding = nibble_count % nibble_ops::NIBBLE_PER_BYTE != 0;
773 if padding && nibble_ops::pad_left(data[input.offset]) != 0 {
775 return Err(CodecError::from("Bad format"))
776 }
777 let partial = input.take(
778 (nibble_count + (nibble_ops::NIBBLE_PER_BYTE - 1)) /
779 nibble_ops::NIBBLE_PER_BYTE,
780 )?;
781 let partial_padding = nibble_ops::number_padding(nibble_count);
782 let bitmap_range = input.take(BITMAP_LENGTH)?;
783 let bitmap = Bitmap::decode(&data[bitmap_range])?;
784 let value = if has_value {
785 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
786 Some(ValuePlan::Inline(input.take(count)?))
787 } else {
788 None
789 };
790 let mut children = [
791 None, None, None, None, None, None, None, None, None, None, None, None, None,
792 None, None, None,
793 ];
794 for i in 0..nibble_ops::NIBBLE_LENGTH {
795 if bitmap.value_at(i) {
796 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
797 let range = input.take(count)?;
798 children[i] = Some(if count == H::LENGTH {
799 NodeHandlePlan::Hash(range)
800 } else {
801 NodeHandlePlan::Inline(range)
802 });
803 }
804 }
805 NodePlan::NibbledBranch {
806 partial: NibbleSlicePlan::new(partial, partial_padding),
807 value,
808 children,
809 }
810 },
811 NodeHeaderNoExt::Leaf(nibble_count) => {
812 let padding = nibble_count % nibble_ops::NIBBLE_PER_BYTE != 0;
813 if padding && nibble_ops::pad_left(data[input.offset]) != 0 {
815 return Err(CodecError::from("Bad format"))
816 }
817 let partial = input.take(
818 (nibble_count + (nibble_ops::NIBBLE_PER_BYTE - 1)) /
819 nibble_ops::NIBBLE_PER_BYTE,
820 )?;
821 let partial_padding = nibble_ops::number_padding(nibble_count);
822 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
823 let value = ValuePlan::Inline(input.take(count)?);
824
825 NodePlan::Leaf { partial: NibbleSlicePlan::new(partial, partial_padding), value }
826 },
827 })
828 }
829
830 fn is_empty_node(data: &[u8]) -> bool {
831 data == <Self as NodeCodec>::empty_node()
832 }
833
834 fn empty_node() -> &'static [u8] {
835 &[EMPTY_TRIE_NO_EXT]
836 }
837
838 fn leaf_node(partial: impl Iterator<Item = u8>, number_nibble: usize, value: Value) -> Vec<u8> {
839 let mut output = partial_from_iterator_encode(partial, number_nibble, NodeKindNoExt::Leaf);
840 match value {
841 Value::Inline(value) => {
842 Compact(value.len() as u32).encode_to(&mut output);
843 output.extend_from_slice(value);
844 },
845 Value::Node(..) => unimplemented!("No support for inner hashed value"),
846 }
847 output
848 }
849
850 fn extension_node(
851 _partial: impl Iterator<Item = u8>,
852 _nbnibble: usize,
853 _child: ChildReference<<H as Hasher>::Out>,
854 ) -> Vec<u8> {
855 unreachable!("no extension codec")
856 }
857
858 fn branch_node(
859 _children: impl Iterator<Item = impl Borrow<Option<ChildReference<<H as Hasher>::Out>>>>,
860 _maybe_value: Option<Value>,
861 ) -> Vec<u8> {
862 unreachable!("no extension codec")
863 }
864
865 fn branch_node_nibbled(
866 partial: impl Iterator<Item = u8>,
867 number_nibble: usize,
868 children: impl Iterator<Item = impl Borrow<Option<ChildReference<Self::HashOut>>>>,
869 maybe_value: Option<Value>,
870 ) -> Vec<u8> {
871 let mut output = if maybe_value.is_none() {
872 partial_from_iterator_encode(partial, number_nibble, NodeKindNoExt::BranchNoValue)
873 } else {
874 partial_from_iterator_encode(partial, number_nibble, NodeKindNoExt::BranchWithValue)
875 };
876 let bitmap_index = output.len();
877 let mut bitmap: [u8; BITMAP_LENGTH] = [0; BITMAP_LENGTH];
878 (0..BITMAP_LENGTH).for_each(|_| output.push(0));
879 match maybe_value {
880 Some(Value::Inline(value)) => {
881 Compact(value.len() as u32).encode_to(&mut output);
882 output.extend_from_slice(value);
883 },
884 Some(Value::Node(..)) => unimplemented!("No support for inner hashed value"),
885 None => (),
886 }
887
888 Bitmap::encode(
889 children.map(|maybe_child| match maybe_child.borrow() {
890 Some(ChildReference::Hash(h)) => {
891 h.as_ref().encode_to(&mut output);
892 true
893 },
894 &Some(ChildReference::Inline(inline_data, len)) => {
895 inline_data.as_ref()[..len].encode_to(&mut output);
896 true
897 },
898 None => false,
899 }),
900 bitmap.as_mut(),
901 );
902 output[bitmap_index..bitmap_index + BITMAP_LENGTH]
903 .copy_from_slice(&bitmap.as_ref()[..BITMAP_LENGTH]);
904 output
905 }
906}
907
908pub fn compare_implementations<T, DB>(data: Vec<(Vec<u8>, Vec<u8>)>, mut memdb: DB, mut hashdb: DB)
910where
911 T: TrieLayout,
912 DB: hash_db::HashDB<T::Hash, DBValue> + Eq,
913{
914 let root_new = calc_root_build::<T, _, _, _, _>(data.clone(), &mut hashdb);
915 let root = {
916 let mut root = Default::default();
917 let mut t = TrieDBMutBuilder::<T>::new(&mut memdb, &mut root).build();
918 for i in 0..data.len() {
919 t.insert(&data[i].0[..], &data[i].1[..]).unwrap();
920 }
921 t.commit();
922 *t.root()
923 };
924 if root_new != root {
925 {
926 let db: &dyn hash_db::HashDB<_, _> = &hashdb;
927 let t = TrieDBBuilder::<T>::new(&db, &root_new).build();
928 println!("{:?}", t);
929 for a in t.iter().unwrap() {
930 println!("a:{:x?}", a);
931 }
932 }
933 {
934 let db: &dyn hash_db::HashDB<_, _> = &memdb;
935 let t = TrieDBBuilder::<T>::new(&db, &root).build();
936 println!("{:?}", t);
937 for a in t.iter().unwrap() {
938 println!("a:{:x?}", a);
939 }
940 }
941 }
942
943 assert_eq!(root, root_new);
944 assert!(memdb == hashdb);
946}
947
948pub fn compare_root<T: TrieLayout, DB: hash_db::HashDB<T::Hash, DBValue>>(
950 data: Vec<(Vec<u8>, Vec<u8>)>,
951 mut memdb: DB,
952) {
953 let root_new = reference_trie_root_iter_build::<T, _, _, _>(data.clone());
954 let root = {
955 let mut root = Default::default();
956 let mut t = TrieDBMutBuilder::<T>::new(&mut memdb, &mut root).build();
957 for i in 0..data.len() {
958 t.insert(&data[i].0[..], &data[i].1[..]).unwrap();
959 }
960 *t.root()
961 };
962
963 assert_eq!(root, root_new);
964}
965
966pub fn compare_unhashed(data: Vec<(Vec<u8>, Vec<u8>)>) {
968 let root_new = {
969 let mut cb = trie_db_fun::TrieRootUnhashed::<ExtensionLayout>::default();
970 trie_visit::<ExtensionLayout, _, _, _, _>(data.clone().into_iter(), &mut cb);
971 cb.root.unwrap_or(Default::default())
972 };
973 let root = reference_trie_root_unhashed(data);
974
975 assert_eq!(root, root_new);
976}
977
978pub fn compare_unhashed_no_extension(data: Vec<(Vec<u8>, Vec<u8>)>) {
981 let root_new = {
982 let mut cb = trie_db_fun::TrieRootUnhashed::<NoExtensionLayout>::default();
983 trie_visit::<NoExtensionLayout, _, _, _, _>(data.clone().into_iter(), &mut cb);
984 cb.root.unwrap_or(Default::default())
985 };
986 let root = reference_trie_root_unhashed_no_extension(data);
987
988 assert_eq!(root, root_new);
989}
990
991pub fn calc_root<T, I, A, B>(data: I) -> <T::Hash as Hasher>::Out
993where
994 T: TrieLayout,
995 I: IntoIterator<Item = (A, B)>,
996 A: AsRef<[u8]> + Ord + fmt::Debug,
997 B: AsRef<[u8]> + fmt::Debug,
998{
999 let mut cb = TrieRoot::<T>::default();
1000 trie_visit::<T, _, _, _, _>(data.into_iter(), &mut cb);
1001 cb.root.unwrap_or_default()
1002}
1003
1004pub fn calc_root_build<T, I, A, B, DB>(data: I, hashdb: &mut DB) -> <T::Hash as Hasher>::Out
1006where
1007 T: TrieLayout,
1008 I: IntoIterator<Item = (A, B)>,
1009 A: AsRef<[u8]> + Ord + fmt::Debug,
1010 B: AsRef<[u8]> + fmt::Debug,
1011 DB: hash_db::HashDB<T::Hash, DBValue>,
1012{
1013 let mut cb = TrieBuilder::<T, DB>::new(hashdb);
1014 trie_visit::<T, _, _, _, _>(data.into_iter(), &mut cb);
1015 cb.root.unwrap_or_default()
1016}
1017
1018pub fn compare_implementations_unordered<T, DB>(
1021 data: Vec<(Vec<u8>, Vec<u8>)>,
1022 mut memdb: DB,
1023 mut hashdb: DB,
1024) where
1025 T: TrieLayout,
1026 DB: hash_db::HashDB<T::Hash, DBValue> + Eq,
1027{
1028 let mut b_map = std::collections::btree_map::BTreeMap::new();
1029 let root = {
1030 let mut root = Default::default();
1031 let mut t = TrieDBMutBuilder::<T>::new(&mut memdb, &mut root).build();
1032 for i in 0..data.len() {
1033 t.insert(&data[i].0[..], &data[i].1[..]).unwrap();
1034 b_map.insert(data[i].0.clone(), data[i].1.clone());
1035 }
1036 *t.root()
1037 };
1038 let root_new = {
1039 let mut cb = TrieBuilder::<T, DB>::new(&mut hashdb);
1040 trie_visit::<T, _, _, _, _>(b_map.into_iter(), &mut cb);
1041 cb.root.unwrap_or_default()
1042 };
1043
1044 if root != root_new {
1045 {
1046 let db: &dyn hash_db::HashDB<_, _> = &memdb;
1047 let t = TrieDBBuilder::<T>::new(&db, &root).build();
1048 println!("{:?}", t);
1049 for a in t.iter().unwrap() {
1050 println!("a:{:?}", a);
1051 }
1052 }
1053 {
1054 let db: &dyn hash_db::HashDB<_, _> = &hashdb;
1055 let t = TrieDBBuilder::<T>::new(&db, &root_new).build();
1056 println!("{:?}", t);
1057 for a in t.iter().unwrap() {
1058 println!("a:{:?}", a);
1059 }
1060 }
1061 }
1062
1063 assert_eq!(root, root_new);
1064}
1065
1066pub fn compare_insert_remove<T, DB: hash_db::HashDB<T::Hash, DBValue>>(
1069 data: Vec<(bool, Vec<u8>, Vec<u8>)>,
1070 mut memdb: DB,
1071) where
1072 T: TrieLayout,
1073 DB: hash_db::HashDB<T::Hash, DBValue> + Eq,
1074{
1075 let mut data2 = std::collections::BTreeMap::new();
1076 let mut root = Default::default();
1077 let mut a = 0;
1078 {
1079 let mut t = TrieDBMutBuilder::<T>::new(&mut memdb, &mut root).build();
1080 t.commit();
1081 }
1082 while a < data.len() {
1083 root = {
1085 let mut t = TrieDBMutBuilder::<T>::from_existing(&mut memdb, &mut root).build();
1086 for _ in 0..3 {
1087 if data[a].0 {
1088 t.remove(&data[a].1[..]).unwrap();
1090 data2.remove(&data[a].1[..]);
1091 } else {
1092 t.insert(&data[a].1[..], &data[a].2[..]).unwrap();
1094 data2.insert(&data[a].1[..], &data[a].2[..]);
1095 }
1096
1097 a += 1;
1098 if a == data.len() {
1099 break
1100 }
1101 }
1102 t.commit();
1103 *t.root()
1104 };
1105 }
1106 let mut t = TrieDBMutBuilder::<T>::from_existing(&mut memdb, &mut root).build();
1107 assert_eq!(*t.root(), calc_root::<T, _, _, _>(data2));
1110}
1111
1112pub struct TestTrieCache<L: TrieLayout> {
1116 value_cache: HashMap<Vec<u8>, trie_db_fun::CachedValue<TrieHash<L>>>,
1118 node_cache: HashMap<TrieHash<L>, NodeOwned<TrieHash<L>>>,
1119}
1120
1121impl<L: TrieLayout> TestTrieCache<L> {
1122 pub fn clear_value_cache(&mut self) {
1124 self.value_cache.clear();
1125 }
1126
1127 pub fn clear_node_cache(&mut self) {
1129 self.node_cache.clear();
1130 }
1131}
1132
1133impl<L: TrieLayout> Default for TestTrieCache<L> {
1134 fn default() -> Self {
1135 Self { value_cache: Default::default(), node_cache: Default::default() }
1136 }
1137}
1138
1139impl<L: TrieLayout> trie_db_fun::TrieCache<L::Codec> for TestTrieCache<L> {
1140 fn lookup_value_for_key(&mut self, key: &[u8]) -> Option<&trie_db_fun::CachedValue<TrieHash<L>>> {
1141 self.value_cache.get(key)
1142 }
1143
1144 fn cache_value_for_key(&mut self, key: &[u8], value: trie_db_fun::CachedValue<TrieHash<L>>) {
1145 self.value_cache.insert(key.to_vec(), value);
1146 }
1147
1148 fn get_or_insert_node(
1149 &mut self,
1150 hash: TrieHash<L>,
1151 fetch_node: &mut dyn FnMut() -> trie_db_fun::Result<
1152 NodeOwned<TrieHash<L>>,
1153 TrieHash<L>,
1154 trie_db_fun::CError<L>,
1155 >,
1156 ) -> trie_db_fun::Result<&NodeOwned<TrieHash<L>>, TrieHash<L>, trie_db_fun::CError<L>> {
1157 match self.node_cache.entry(hash) {
1158 Entry::Occupied(e) => Ok(e.into_mut()),
1159 Entry::Vacant(e) => {
1160 let node = (*fetch_node)()?;
1161 Ok(e.insert(node))
1162 },
1163 }
1164 }
1165
1166 fn get_node(&mut self, hash: &TrieHash<L>) -> Option<&NodeOwned<TrieHash<L>>> {
1167 self.node_cache.get(hash)
1168 }
1169}
1170
1171#[cfg(test)]
1172mod tests {
1173 use super::*;
1174 use trie_db_fun::{nibble_ops::NIBBLE_PER_BYTE, node::Node};
1175
1176 const _: fn() -> () = || {
1177 struct AssertTrieDBRawIteratorIsSendAndSync
1178 where
1179 trie_db_fun::TrieDBRawIterator<NoExtensionLayout>: Send + Sync;
1180 };
1181
1182 #[test]
1183 fn test_encoding_simple_trie() {
1184 for prefix in
1185 [LEAF_PREFIX_MASK_NO_EXT, BRANCH_WITHOUT_MASK_NO_EXT, BRANCH_WITH_MASK_NO_EXT].iter()
1186 {
1187 for i in (0..1000).chain(NIBBLE_SIZE_BOUND_NO_EXT - 2..NIBBLE_SIZE_BOUND_NO_EXT + 2) {
1188 let mut output = Vec::new();
1189 encode_size_and_prefix(i, *prefix, &mut output);
1190 let input = &mut &output[..];
1191 let first = input.read_byte().unwrap();
1192 assert_eq!(first & (0b11 << 6), *prefix);
1193 let v = decode_size(first, input);
1194 assert_eq!(Ok(std::cmp::min(i, NIBBLE_SIZE_BOUND_NO_EXT)), v);
1195 }
1196 }
1197 }
1198
1199 #[test]
1200 fn too_big_nibble_length() {
1201 let input = vec![0u8; (NIBBLE_SIZE_BOUND_NO_EXT as usize + 1) / 2 + 1];
1203 let enc = <ReferenceNodeCodecNoExt<RefHasher> as NodeCodec>::leaf_node(
1204 input.iter().cloned(),
1205 input.len() * NIBBLE_PER_BYTE,
1206 Value::Inline(&[1]),
1207 );
1208 let dec = <ReferenceNodeCodecNoExt<RefHasher> as NodeCodec>::decode(&enc).unwrap();
1209 let o_sl = if let Node::Leaf(sl, _) = dec { Some(sl) } else { None };
1210 assert!(o_sl.is_some());
1211 }
1212
1213 #[test]
1214 fn size_encode_limit_values() {
1215 let sizes = [0, 1, 62, 63, 64, 317, 318, 319, 572, 573, 574];
1216 let encs = [
1217 vec![0],
1218 vec![1],
1219 vec![0x3e],
1220 vec![0x3f, 0],
1221 vec![0x3f, 1],
1222 vec![0x3f, 0xfe],
1223 vec![0x3f, 0xff, 0],
1224 vec![0x3f, 0xff, 1],
1225 vec![0x3f, 0xff, 0xfe],
1226 vec![0x3f, 0xff, 0xff, 0],
1227 vec![0x3f, 0xff, 0xff, 1],
1228 ];
1229 for i in 0..sizes.len() {
1230 let mut enc = Vec::new();
1231 encode_size_and_prefix(sizes[i], 0, &mut enc);
1232 assert_eq!(enc, encs[i]);
1233 let s_dec = decode_size(encs[i][0], &mut &encs[i][1..]);
1234 assert_eq!(s_dec, Ok(sizes[i]));
1235 }
1236 }
1237}