1use tetsy_hash_db::Hasher;
21use tetsy_trie_root;
22use codec::Encode;
23use tetcore_std::vec::Vec;
24use crate::trie_constants;
25use crate::node_header::{NodeKind, size_and_prefix_iterator};
26use crate::node_codec::Bitmap;
27
28const BRANCH_NODE_NO_VALUE: u8 = 254;
29const BRANCH_NODE_WITH_VALUE: u8 = 255;
30
31#[derive(Default, Clone)]
32pub struct TrieStream {
34 buffer: Vec<u8>,
35}
36
37impl TrieStream {
38 pub fn as_raw(&self) -> &[u8] { &self.buffer }
40}
41
42fn branch_node_bit_mask(has_children: impl Iterator<Item = bool>) -> (u8, u8) {
43 let mut bitmap: u16 = 0;
44 let mut cursor: u16 = 1;
45 for v in has_children {
46 if v { bitmap |= cursor }
47 cursor <<= 1;
48 }
49 ((bitmap % 256 ) as u8, (bitmap / 256 ) as u8)
50}
51
52
53fn fuse_nibbles_node<'a>(nibbles: &'a [u8], kind: NodeKind) -> impl Iterator<Item = u8> + 'a {
55 let size = tetcore_std::cmp::min(trie_constants::NIBBLE_SIZE_BOUND, nibbles.len());
56
57 let iter_start = match kind {
58 NodeKind::Leaf => size_and_prefix_iterator(size, trie_constants::LEAF_PREFIX_MASK),
59 NodeKind::BranchNoValue => size_and_prefix_iterator(size, trie_constants::BRANCH_WITHOUT_MASK),
60 NodeKind::BranchWithValue => size_and_prefix_iterator(size, trie_constants::BRANCH_WITH_MASK),
61 };
62 iter_start
63 .chain(if nibbles.len() % 2 == 1 { Some(nibbles[0]) } else { None })
64 .chain(nibbles[nibbles.len() % 2..].chunks(2).map(|ch| ch[0] << 4 | ch[1]))
65}
66
67
68impl tetsy_trie_root::TrieStream for TrieStream {
69
70 fn new() -> Self {
71 TrieStream {
72 buffer: Vec::new()
73 }
74 }
75
76 fn append_empty_data(&mut self) {
77 self.buffer.push(trie_constants::EMPTY_TRIE);
78 }
79
80 fn append_leaf(&mut self, key: &[u8], value: &[u8]) {
81 self.buffer.extend(fuse_nibbles_node(key, NodeKind::Leaf));
82 value.encode_to(&mut self.buffer);
83 }
84
85 fn begin_branch(
86 &mut self,
87 maybe_partial: Option<&[u8]>,
88 maybe_value: Option<&[u8]>,
89 has_children: impl Iterator<Item = bool>,
90 ) {
91 if let Some(partial) = maybe_partial {
92 if maybe_value.is_some() {
93 self.buffer.extend(fuse_nibbles_node(partial, NodeKind::BranchWithValue));
94 } else {
95 self.buffer.extend(fuse_nibbles_node(partial, NodeKind::BranchNoValue));
96 }
97 let bm = branch_node_bit_mask(has_children);
98 self.buffer.extend([bm.0,bm.1].iter());
99 } else {
100 debug_assert!(false, "trie stream codec only for no extension trie");
101 self.buffer.extend(&branch_node(maybe_value.is_some(), has_children));
102 }
103 if let Some(value) = maybe_value {
104 value.encode_to(&mut self.buffer);
105 }
106 }
107
108 fn append_extension(&mut self, _key: &[u8]) {
109 debug_assert!(false, "trie stream codec only for no extension trie");
110 }
111
112 fn append_substream<H: Hasher>(&mut self, other: Self) {
113 let data = other.out();
114 match data.len() {
115 0..=31 => data.encode_to(&mut self.buffer),
116 _ => H::hash(&data).as_ref().encode_to(&mut self.buffer),
117 }
118 }
119
120 fn out(self) -> Vec<u8> { self.buffer }
121}
122
123fn branch_node(has_value: bool, has_children: impl Iterator<Item = bool>) -> [u8; 3] {
124 let mut result = [0, 0, 0];
125 branch_node_buffered(has_value, has_children, &mut result[..]);
126 result
127}
128
129fn branch_node_buffered<I>(has_value: bool, has_children: I, output: &mut[u8])
130 where
131 I: Iterator<Item = bool>,
132{
133 let first = if has_value {
134 BRANCH_NODE_WITH_VALUE
135 } else {
136 BRANCH_NODE_NO_VALUE
137 };
138 output[0] = first;
139 Bitmap::encode(has_children, &mut output[1..]);
140}