1use super::node_header::{NodeHeader, NodeKind};
21use crate::{error::Error, trie_constants};
22use alloc::{borrow::Borrow, vec::Vec};
23use codec::{Compact, Decode, Encode, Input};
24use core::{marker::PhantomData, ops::Range};
25use hash_db::Hasher;
26use trie_db::{
27 nibble_ops,
28 node::{NibbleSlicePlan, NodeHandlePlan, NodePlan, Value, ValuePlan},
29 ChildReference, NodeCodec as NodeCodecT,
30};
31
32struct ByteSliceInput<'a> {
36 data: &'a [u8],
37 offset: usize,
38}
39
40impl<'a> ByteSliceInput<'a> {
41 fn new(data: &'a [u8]) -> Self {
42 ByteSliceInput { data, offset: 0 }
43 }
44
45 fn take(&mut self, count: usize) -> Result<Range<usize>, codec::Error> {
46 if self.offset + count > self.data.len() {
47 return Err("out of data".into())
48 }
49
50 let range = self.offset..(self.offset + count);
51 self.offset += count;
52 Ok(range)
53 }
54}
55
56impl<'a> Input for ByteSliceInput<'a> {
57 fn remaining_len(&mut self) -> Result<Option<usize>, codec::Error> {
58 Ok(Some(self.data.len().saturating_sub(self.offset)))
59 }
60
61 fn read(&mut self, into: &mut [u8]) -> Result<(), codec::Error> {
62 let range = self.take(into.len())?;
63 into.copy_from_slice(&self.data[range]);
64 Ok(())
65 }
66
67 fn read_byte(&mut self) -> Result<u8, codec::Error> {
68 if self.offset + 1 > self.data.len() {
69 return Err("out of data".into())
70 }
71
72 let byte = self.data[self.offset];
73 self.offset += 1;
74 Ok(byte)
75 }
76}
77
78#[derive(Default, Clone)]
82pub struct NodeCodec<H>(PhantomData<H>);
83
84impl<H> NodeCodecT for NodeCodec<H>
85where
86 H: Hasher,
87{
88 const ESCAPE_HEADER: Option<u8> = Some(trie_constants::ESCAPE_COMPACT_HEADER);
89 type Error = Error<H::Out>;
90 type HashOut = H::Out;
91
92 fn hashed_null_node() -> <H as Hasher>::Out {
93 H::hash(<Self as NodeCodecT>::empty_node())
94 }
95
96 fn decode_plan(data: &[u8]) -> Result<NodePlan, Self::Error> {
97 let mut input = ByteSliceInput::new(data);
98
99 let header = NodeHeader::decode(&mut input)?;
100 let contains_hash = header.contains_hash_of_value();
101
102 let branch_has_value = if let NodeHeader::Branch(has_value, _) = &header {
103 *has_value
104 } else {
105 true
107 };
108
109 match header {
110 NodeHeader::Null => Ok(NodePlan::Empty),
111 NodeHeader::HashedValueBranch(nibble_count) | NodeHeader::Branch(_, nibble_count) => {
112 let padding = nibble_count % nibble_ops::NIBBLE_PER_BYTE != 0;
113 if padding && nibble_ops::pad_left(data[input.offset]) != 0 {
115 return Err(Error::BadFormat)
116 }
117 let partial = input.take(
118 (nibble_count + (nibble_ops::NIBBLE_PER_BYTE - 1)) /
119 nibble_ops::NIBBLE_PER_BYTE,
120 )?;
121 let partial_padding = nibble_ops::number_padding(nibble_count);
122 let bitmap_range = input.take(BITMAP_LENGTH)?;
123 let bitmap = Bitmap::decode(&data[bitmap_range])?;
124 let value = if branch_has_value {
125 Some(if contains_hash {
126 ValuePlan::Node(input.take(H::LENGTH)?)
127 } else {
128 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
129 ValuePlan::Inline(input.take(count)?)
130 })
131 } else {
132 None
133 };
134 let mut children = [
135 None, None, None, None, None, None, None, None, None, None, None, None, None,
136 None, None, None,
137 ];
138 for i in 0..nibble_ops::NIBBLE_LENGTH {
139 if bitmap.value_at(i) {
140 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
141 let range = input.take(count)?;
142 children[i] = Some(if count == H::LENGTH {
143 NodeHandlePlan::Hash(range)
144 } else {
145 NodeHandlePlan::Inline(range)
146 });
147 }
148 }
149 Ok(NodePlan::NibbledBranch {
150 partial: NibbleSlicePlan::new(partial, partial_padding),
151 value,
152 children,
153 })
154 },
155 NodeHeader::HashedValueLeaf(nibble_count) | NodeHeader::Leaf(nibble_count) => {
156 let padding = nibble_count % nibble_ops::NIBBLE_PER_BYTE != 0;
157 if padding && nibble_ops::pad_left(data[input.offset]) != 0 {
159 return Err(Error::BadFormat)
160 }
161 let partial = input.take(
162 (nibble_count + (nibble_ops::NIBBLE_PER_BYTE - 1)) /
163 nibble_ops::NIBBLE_PER_BYTE,
164 )?;
165 let partial_padding = nibble_ops::number_padding(nibble_count);
166 let value = if contains_hash {
167 ValuePlan::Node(input.take(H::LENGTH)?)
168 } else {
169 let count = <Compact<u32>>::decode(&mut input)?.0 as usize;
170 ValuePlan::Inline(input.take(count)?)
171 };
172
173 Ok(NodePlan::Leaf {
174 partial: NibbleSlicePlan::new(partial, partial_padding),
175 value,
176 })
177 },
178 }
179 }
180
181 fn is_empty_node(data: &[u8]) -> bool {
182 data == <Self as NodeCodecT>::empty_node()
183 }
184
185 fn empty_node() -> &'static [u8] {
186 &[trie_constants::EMPTY_TRIE]
187 }
188
189 fn leaf_node(partial: impl Iterator<Item = u8>, number_nibble: usize, value: Value) -> Vec<u8> {
190 let contains_hash = matches!(&value, Value::Node(..));
191 let mut output = if contains_hash {
192 partial_from_iterator_encode(partial, number_nibble, NodeKind::HashedValueLeaf)
193 } else {
194 partial_from_iterator_encode(partial, number_nibble, NodeKind::Leaf)
195 };
196 match value {
197 Value::Inline(value) => {
198 Compact(value.len() as u32).encode_to(&mut output);
199 output.extend_from_slice(value);
200 },
201 Value::Node(hash) => {
202 debug_assert!(hash.len() == H::LENGTH);
203 output.extend_from_slice(hash);
204 },
205 }
206 output
207 }
208
209 fn extension_node(
210 _partial: impl Iterator<Item = u8>,
211 _nbnibble: usize,
212 _child: ChildReference<<H as Hasher>::Out>,
213 ) -> Vec<u8> {
214 unreachable!("No extension codec.")
215 }
216
217 fn branch_node(
218 _children: impl Iterator<Item = impl Borrow<Option<ChildReference<<H as Hasher>::Out>>>>,
219 _maybe_value: Option<Value>,
220 ) -> Vec<u8> {
221 unreachable!("No extension codec.")
222 }
223
224 fn branch_node_nibbled(
225 partial: impl Iterator<Item = u8>,
226 number_nibble: usize,
227 children: impl Iterator<Item = impl Borrow<Option<ChildReference<<H as Hasher>::Out>>>>,
228 value: Option<Value>,
229 ) -> Vec<u8> {
230 let contains_hash = matches!(&value, Some(Value::Node(..)));
231 let mut output = match (&value, contains_hash) {
232 (&None, _) =>
233 partial_from_iterator_encode(partial, number_nibble, NodeKind::BranchNoValue),
234 (_, false) =>
235 partial_from_iterator_encode(partial, number_nibble, NodeKind::BranchWithValue),
236 (_, true) =>
237 partial_from_iterator_encode(partial, number_nibble, NodeKind::HashedValueBranch),
238 };
239
240 let bitmap_index = output.len();
241 let mut bitmap: [u8; BITMAP_LENGTH] = [0; BITMAP_LENGTH];
242 (0..BITMAP_LENGTH).for_each(|_| output.push(0));
243 match value {
244 Some(Value::Inline(value)) => {
245 Compact(value.len() as u32).encode_to(&mut output);
246 output.extend_from_slice(value);
247 },
248 Some(Value::Node(hash)) => {
249 debug_assert!(hash.len() == H::LENGTH);
250 output.extend_from_slice(hash);
251 },
252 None => (),
253 }
254 Bitmap::encode(
255 children.map(|maybe_child| match maybe_child.borrow() {
256 Some(ChildReference::Hash(h)) => {
257 h.as_ref().encode_to(&mut output);
258 true
259 },
260 &Some(ChildReference::Inline(inline_data, len)) => {
261 inline_data.as_ref()[..len].encode_to(&mut output);
262 true
263 },
264 None => false,
265 }),
266 bitmap.as_mut(),
267 );
268 output[bitmap_index..bitmap_index + BITMAP_LENGTH]
269 .copy_from_slice(&bitmap[..BITMAP_LENGTH]);
270 output
271 }
272}
273
274fn partial_from_iterator_encode<I: Iterator<Item = u8>>(
279 partial: I,
280 nibble_count: usize,
281 node_kind: NodeKind,
282) -> Vec<u8> {
283 let mut output = Vec::with_capacity(4 + (nibble_count / nibble_ops::NIBBLE_PER_BYTE));
284 match node_kind {
285 NodeKind::Leaf => NodeHeader::Leaf(nibble_count).encode_to(&mut output),
286 NodeKind::BranchWithValue => NodeHeader::Branch(true, nibble_count).encode_to(&mut output),
287 NodeKind::BranchNoValue => NodeHeader::Branch(false, nibble_count).encode_to(&mut output),
288 NodeKind::HashedValueLeaf =>
289 NodeHeader::HashedValueLeaf(nibble_count).encode_to(&mut output),
290 NodeKind::HashedValueBranch =>
291 NodeHeader::HashedValueBranch(nibble_count).encode_to(&mut output),
292 };
293 output.extend(partial);
294 output
295}
296
297const BITMAP_LENGTH: usize = 2;
298
299pub(crate) struct Bitmap(u16);
304
305impl Bitmap {
306 pub fn decode(data: &[u8]) -> Result<Self, codec::Error> {
307 let value = u16::decode(&mut &data[..])?;
308 if value == 0 {
309 Err("Bitmap without a child.".into())
310 } else {
311 Ok(Bitmap(value))
312 }
313 }
314
315 pub fn value_at(&self, i: usize) -> bool {
316 self.0 & (1u16 << i) != 0
317 }
318
319 pub fn encode<I: Iterator<Item = bool>>(has_children: I, dest: &mut [u8]) {
320 let mut bitmap: u16 = 0;
321 let mut cursor: u16 = 1;
322 for v in has_children {
323 if v {
324 bitmap |= cursor
325 }
326 cursor <<= 1;
327 }
328 dest[0] = (bitmap % 256) as u8;
329 dest[1] = (bitmap / 256) as u8;
330 }
331}