1use core::iter;
16
17use crate::{
18 nibble::{self, nibble_ops, NibbleSlice, NibbleVec},
19 node_codec::NodeCodec,
20 Bytes, CError, ChildReference, Result, TrieError, TrieHash, TrieLayout,
21};
22#[cfg(not(feature = "std"))]
23use alloc::{boxed::Box, vec::Vec};
24use hash_db::Hasher;
25
26use crate::rstd::{borrow::Borrow, mem, ops::Range};
27
28pub type NodeKey = (usize, nibble::BackingByteVec);
31
32#[derive(Debug, Clone, Copy, PartialEq, Eq)]
34pub enum NodeHandle<'a> {
35 Hash(&'a [u8]),
36 Inline(&'a [u8]),
37}
38
39impl NodeHandle<'_> {
40 pub fn to_owned_handle<L: TrieLayout>(
42 &self,
43 ) -> Result<NodeHandleOwned<TrieHash<L>>, TrieHash<L>, CError<L>> {
44 match self {
45 Self::Hash(h) => decode_hash::<L::Hash>(h)
46 .ok_or_else(|| Box::new(TrieError::InvalidHash(Default::default(), h.to_vec())))
47 .map(NodeHandleOwned::Hash),
48 Self::Inline(i) => match L::Codec::decode(i) {
49 Ok(node) => Ok(NodeHandleOwned::Inline(Box::new(node.to_owned_node::<L>()?))),
50 Err(e) => Err(Box::new(TrieError::DecoderError(Default::default(), e))),
51 },
52 }
53 }
54}
55
56#[derive(Clone, PartialEq, Eq)]
58#[cfg_attr(feature = "std", derive(Debug))]
59pub enum NodeHandleOwned<H> {
60 Hash(H),
61 Inline(Box<NodeOwned<H>>),
62}
63
64impl<H> NodeHandleOwned<H>
65where
66 H: Default + AsRef<[u8]> + AsMut<[u8]> + Copy,
67{
68 fn as_child_reference<C: NodeCodec<HashOut = H>>(&self) -> ChildReference<H> {
75 match self {
76 NodeHandleOwned::Hash(h) => ChildReference::Hash(*h),
77 NodeHandleOwned::Inline(n) => {
78 let encoded = n.to_encoded::<C>();
79 let mut store = H::default();
80 assert!(store.as_ref().len() > encoded.len(), "Invalid inline node handle");
81
82 store.as_mut()[..encoded.len()].copy_from_slice(&encoded);
83 ChildReference::Inline(store, encoded.len())
84 },
85 }
86 }
87}
88
89impl<H> NodeHandleOwned<H> {
90 pub fn as_inline(&self) -> Option<&NodeOwned<H>> {
92 match self {
93 Self::Hash(_) => None,
94 Self::Inline(node) => Some(&*node),
95 }
96 }
97}
98
99pub fn decode_hash<H: Hasher>(data: &[u8]) -> Option<H::Out> {
101 if data.len() != H::LENGTH {
102 return None
103 }
104 let mut hash = H::Out::default();
105 hash.as_mut().copy_from_slice(data);
106 Some(hash)
107}
108
109#[derive(Eq, PartialEq, Clone)]
111#[cfg_attr(feature = "std", derive(Debug))]
112pub enum Value<'a> {
113 Inline(&'a [u8]),
115 Node(&'a [u8]),
117}
118
119impl<'a> Value<'a> {
120 pub(crate) fn new_inline(value: &'a [u8], threshold: Option<u32>) -> Option<Self> {
121 if let Some(threshold) = threshold {
122 if value.len() >= threshold as usize {
123 return None
124 } else {
125 Some(Value::Inline(value))
126 }
127 } else {
128 Some(Value::Inline(value))
129 }
130 }
131
132 pub fn to_owned_value<L: TrieLayout>(&self) -> ValueOwned<TrieHash<L>> {
133 match self {
134 Self::Inline(data) => ValueOwned::Inline(Bytes::from(*data), L::Hash::hash(data)),
135 Self::Node(hash) => {
136 let mut res = TrieHash::<L>::default();
137 res.as_mut().copy_from_slice(hash);
138
139 ValueOwned::Node(res)
140 },
141 }
142 }
143}
144
145#[derive(Eq, PartialEq, Clone)]
147#[cfg_attr(feature = "std", derive(Debug))]
148pub enum ValueOwned<H> {
149 Inline(Bytes, H),
151 Node(H),
153}
154
155impl<H: AsRef<[u8]> + Copy> ValueOwned<H> {
156 pub fn as_value(&self) -> Value {
158 match self {
159 Self::Inline(data, _) => Value::Inline(&data),
160 Self::Node(hash) => Value::Node(hash.as_ref()),
161 }
162 }
163
164 pub fn data_hash(&self) -> Option<H> {
166 match self {
167 Self::Inline(_, hash) => Some(*hash),
168 Self::Node(hash) => Some(*hash),
169 }
170 }
171}
172
173impl<H> ValueOwned<H> {
174 pub fn data(&self) -> Option<&Bytes> {
176 match self {
177 Self::Inline(data, _) => Some(data),
178 Self::Node(_) => None,
179 }
180 }
181}
182
183#[derive(Eq, PartialEq, Clone)]
185#[cfg_attr(feature = "std", derive(Debug))]
186pub enum Node<'a> {
187 Empty,
189 Leaf(NibbleSlice<'a>, Value<'a>),
191 Extension(NibbleSlice<'a>, NodeHandle<'a>),
193 Branch([Option<NodeHandle<'a>>; nibble_ops::NIBBLE_LENGTH], Option<Value<'a>>),
196 NibbledBranch(
198 NibbleSlice<'a>,
199 [Option<NodeHandle<'a>>; nibble_ops::NIBBLE_LENGTH],
200 Option<Value<'a>>,
201 ),
202}
203
204impl Node<'_> {
205 pub fn to_owned_node<L: TrieLayout>(
207 &self,
208 ) -> Result<NodeOwned<TrieHash<L>>, TrieHash<L>, CError<L>> {
209 match self {
210 Self::Empty => Ok(NodeOwned::Empty),
211 Self::Leaf(n, d) => Ok(NodeOwned::Leaf((*n).into(), d.to_owned_value::<L>())),
212 Self::Extension(n, h) =>
213 Ok(NodeOwned::Extension((*n).into(), h.to_owned_handle::<L>()?)),
214 Self::Branch(childs, data) => {
215 let childs_owned = if let Some(last_existing_child_index) =
216 childs.iter().rposition(Option::is_some)
217 {
218 let mut childs_owned = Vec::with_capacity(last_existing_child_index + 1);
219
220 childs
221 .iter()
222 .enumerate()
223 .map(|(i, c)| {
224 if i <= last_existing_child_index {
225 childs_owned.push(
226 c.as_ref().map(|c| c.to_owned_handle::<L>()).transpose()?,
227 );
228 }
229 Ok(())
230 })
231 .collect::<Result<_, _, _>>()?;
232 childs_owned
233 } else {
234 Vec::with_capacity(0)
235 };
236
237 Ok(NodeOwned::Branch(
238 ChildrenNodesOwned(childs_owned),
239 data.as_ref().map(|d| d.to_owned_value::<L>()),
240 ))
241 },
242 Self::NibbledBranch(n, childs, data) => {
243 let childs_owned = if let Some(last_existing_child_index) =
244 childs.iter().rposition(Option::is_some)
245 {
246 let mut childs_owned = Vec::with_capacity(last_existing_child_index + 1);
247
248 childs
249 .iter()
250 .enumerate()
251 .map(|(i, c)| {
252 if i <= last_existing_child_index {
253 childs_owned.push(
254 c.as_ref().map(|c| c.to_owned_handle::<L>()).transpose()?,
255 );
256 }
257 Ok(())
258 })
259 .collect::<Result<_, _, _>>()?;
260 childs_owned
261 } else {
262 Vec::with_capacity(0)
263 };
264
265 Ok(NodeOwned::NibbledBranch(
266 (*n).into(),
267 ChildrenNodesOwned(childs_owned),
268 data.as_ref().map(|d| d.to_owned_value::<L>()),
269 ))
270 },
271 }
272 }
273}
274
275#[derive(Eq, PartialEq, Clone)]
277#[cfg_attr(feature = "std", derive(Debug))]
278pub enum NodeOwned<H> {
279 Empty,
281 Leaf(NibbleVec, ValueOwned<H>),
283 Extension(NibbleVec, NodeHandleOwned<H>),
285 Branch(ChildrenNodesOwned<H>, Option<ValueOwned<H>>),
288 NibbledBranch(NibbleVec, ChildrenNodesOwned<H>, Option<ValueOwned<H>>),
290 Value(Bytes, H),
295}
296
297#[derive(Eq, PartialEq, Clone)]
310#[cfg_attr(feature = "std", derive(Debug))]
311pub struct ChildrenNodesOwned<H>(Vec<Option<NodeHandleOwned<H>>>);
312
313impl<H> ChildrenNodesOwned<H> {
314 pub fn get(&self, index: usize) -> Option<&NodeHandleOwned<H>> {
316 self.0.get(index).and_then(|node| node.as_ref())
317 }
318
319 pub fn iter(&self) -> impl Iterator<Item = &Option<NodeHandleOwned<H>>> {
324 self.0.iter().chain(iter::repeat(&None)).take(nibble_ops::NIBBLE_LENGTH)
325 }
326
327 pub fn allocated_len(&self) -> usize {
329 self.0.len()
330 }
331}
332
333impl<H> NodeOwned<H>
334where
335 H: Default + AsRef<[u8]> + AsMut<[u8]> + Copy,
336{
337 pub fn to_encoded<C>(&self) -> Vec<u8>
339 where
340 C: NodeCodec<HashOut = H>,
341 {
342 match self {
343 Self::Empty => C::empty_node().to_vec(),
344 Self::Leaf(partial, value) =>
345 C::leaf_node(partial.right_iter(), partial.len(), value.as_value()),
346 Self::Extension(partial, child) => C::extension_node(
347 partial.right_iter(),
348 partial.len(),
349 child.as_child_reference::<C>(),
350 ),
351 Self::Branch(children, value) => C::branch_node(
352 children.iter().map(|child| child.as_ref().map(|c| c.as_child_reference::<C>())),
353 value.as_ref().map(|v| v.as_value()),
354 ),
355 Self::NibbledBranch(partial, children, value) => C::branch_node_nibbled(
356 partial.right_iter(),
357 partial.len(),
358 children.iter().map(|child| child.as_ref().map(|c| c.as_child_reference::<C>())),
359 value.as_ref().map(|v| v.as_value()),
360 ),
361 Self::Value(data, _) => data.to_vec(),
362 }
363 }
364
365 pub fn child_iter(&self) -> impl Iterator<Item = (Option<u8>, &NodeHandleOwned<H>)> {
367 enum ChildIter<'a, H> {
368 Empty,
369 Single(&'a NodeHandleOwned<H>, bool),
370 Array(&'a ChildrenNodesOwned<H>, usize),
371 }
372
373 impl<'a, H> Iterator for ChildIter<'a, H> {
374 type Item = (Option<u8>, &'a NodeHandleOwned<H>);
375
376 fn next(&mut self) -> Option<Self::Item> {
377 loop {
378 match self {
379 Self::Empty => break None,
380 Self::Single(child, returned) =>
381 break if *returned {
382 None
383 } else {
384 *returned = true;
385 Some((None, child))
386 },
387 Self::Array(childs, index) =>
388 if *index >= childs.allocated_len() {
389 break None
390 } else {
391 *index += 1;
392
393 if let Some(ref child) = childs.get(*index - 1) {
395 break Some((Some(*index as u8 - 1), child))
396 }
397 },
398 }
399 }
400 }
401 }
402
403 match self {
404 Self::Leaf(_, _) | Self::Empty | Self::Value(_, _) => ChildIter::Empty,
405 Self::Extension(_, child) => ChildIter::Single(child, false),
406 Self::Branch(children, _) | Self::NibbledBranch(_, children, _) =>
407 ChildIter::Array(children, 0),
408 }
409 }
410
411 pub fn data_hash(&self) -> Option<H> {
413 match &self {
414 Self::Empty => None,
415 Self::Leaf(_, value) => value.data_hash(),
416 Self::Extension(_, _) => None,
417 Self::Branch(_, value) => value.as_ref().and_then(|v| v.data_hash()),
418 Self::NibbledBranch(_, _, value) => value.as_ref().and_then(|v| v.data_hash()),
419 Self::Value(_, hash) => Some(*hash),
420 }
421 }
422}
423
424impl<H> NodeOwned<H> {
425 pub fn data(&self) -> Option<&Bytes> {
427 match &self {
428 Self::Empty => None,
429 Self::Leaf(_, value) => value.data(),
430 Self::Extension(_, _) => None,
431 Self::Branch(_, value) => value.as_ref().and_then(|v| v.data()),
432 Self::NibbledBranch(_, _, value) => value.as_ref().and_then(|v| v.data()),
433 Self::Value(data, _) => Some(data),
434 }
435 }
436
437 pub fn partial_key(&self) -> Option<&NibbleVec> {
439 match self {
440 Self::Branch(_, _) | Self::Value(_, _) | Self::Empty => None,
441 Self::Extension(partial, _) |
442 Self::Leaf(partial, _) |
443 Self::NibbledBranch(partial, _, _) => Some(partial),
444 }
445 }
446
447 pub fn size_in_bytes(&self) -> usize {
451 let self_size = mem::size_of::<Self>();
452
453 fn childs_size<'a, H: 'a>(
454 childs: impl Iterator<Item = &'a Option<NodeHandleOwned<H>>>,
455 ) -> usize {
456 childs
459 .filter_map(|c| c.as_ref())
460 .map(|c| c.as_inline().map_or(0, |n| n.size_in_bytes()))
461 .sum()
462 }
463
464 let dynamic_size = match self {
467 Self::Empty => 0,
468 Self::Leaf(nibbles, value) =>
469 nibbles.inner().len() + value.data().map_or(0, |b| b.len()),
470 Self::Value(bytes, _) => bytes.len(),
471 Self::Extension(nibbles, child) => {
472 nibbles.inner().len() + child.as_inline().map_or(0, |n| n.size_in_bytes())
475 },
476 Self::Branch(childs, value) =>
477 childs_size(childs.iter()) +
478 value.as_ref().and_then(|v| v.data()).map_or(0, |b| b.len()),
479 Self::NibbledBranch(nibbles, childs, value) =>
480 nibbles.inner().len() +
481 childs_size(childs.iter()) +
482 value.as_ref().and_then(|v| v.data()).map_or(0, |b| b.len()),
483 };
484
485 self_size + dynamic_size
486 }
487}
488
489#[derive(Debug, Clone, PartialEq, Eq)]
492pub enum NodeHandlePlan {
493 Hash(Range<usize>),
494 Inline(Range<usize>),
495}
496
497impl NodeHandlePlan {
498 pub fn build<'a, 'b>(&'a self, data: &'b [u8]) -> NodeHandle<'b> {
502 match self {
503 NodeHandlePlan::Hash(range) => NodeHandle::Hash(&data[range.clone()]),
504 NodeHandlePlan::Inline(range) => NodeHandle::Inline(&data[range.clone()]),
505 }
506 }
507}
508
509#[derive(Eq, PartialEq, Clone)]
512#[cfg_attr(feature = "std", derive(Debug))]
513pub struct NibbleSlicePlan {
514 bytes: Range<usize>,
515 offset: usize,
516}
517
518impl NibbleSlicePlan {
519 pub fn new(bytes: Range<usize>, offset: usize) -> Self {
521 NibbleSlicePlan { bytes, offset }
522 }
523
524 pub fn len(&self) -> usize {
526 (self.bytes.end - self.bytes.start) * nibble_ops::NIBBLE_PER_BYTE - self.offset
527 }
528
529 pub fn build<'a, 'b>(&'a self, data: &'b [u8]) -> NibbleSlice<'b> {
533 NibbleSlice::new_offset(&data[self.bytes.clone()], self.offset)
534 }
535}
536
537#[derive(Eq, PartialEq, Clone)]
539#[cfg_attr(feature = "std", derive(Debug))]
540pub enum ValuePlan {
541 Inline(Range<usize>),
543 Node(Range<usize>),
546}
547
548impl ValuePlan {
549 pub fn build<'a, 'b>(&'a self, data: &'b [u8]) -> Value<'b> {
551 match self {
552 ValuePlan::Inline(range) => Value::Inline(&data[range.clone()]),
553 ValuePlan::Node(range) => Value::Node(&data[range.clone()]),
554 }
555 }
556}
557
558#[derive(Eq, PartialEq, Clone)]
565#[cfg_attr(feature = "std", derive(Debug))]
566pub enum NodePlan {
567 Empty,
569 Leaf { partial: NibbleSlicePlan, value: ValuePlan },
571 Extension { partial: NibbleSlicePlan, child: NodeHandlePlan },
573 Branch {
576 value: Option<ValuePlan>,
577 children: [Option<NodeHandlePlan>; nibble_ops::NIBBLE_LENGTH],
578 },
579 NibbledBranch {
581 partial: NibbleSlicePlan,
582 value: Option<ValuePlan>,
583 children: [Option<NodeHandlePlan>; nibble_ops::NIBBLE_LENGTH],
584 },
585}
586
587impl NodePlan {
588 pub fn build<'a, 'b>(&'a self, data: &'b [u8]) -> Node<'b> {
592 match self {
593 NodePlan::Empty => Node::Empty,
594 NodePlan::Leaf { partial, value } => Node::Leaf(partial.build(data), value.build(data)),
595 NodePlan::Extension { partial, child } =>
596 Node::Extension(partial.build(data), child.build(data)),
597 NodePlan::Branch { value, children } => {
598 let mut child_slices = [None; nibble_ops::NIBBLE_LENGTH];
599 for i in 0..nibble_ops::NIBBLE_LENGTH {
600 child_slices[i] = children[i].as_ref().map(|child| child.build(data));
601 }
602 Node::Branch(child_slices, value.as_ref().map(|v| v.build(data)))
603 },
604 NodePlan::NibbledBranch { partial, value, children } => {
605 let mut child_slices = [None; nibble_ops::NIBBLE_LENGTH];
606 for i in 0..nibble_ops::NIBBLE_LENGTH {
607 child_slices[i] = children[i].as_ref().map(|child| child.build(data));
608 }
609 Node::NibbledBranch(
610 partial.build(data),
611 child_slices,
612 value.as_ref().map(|v| v.build(data)),
613 )
614 },
615 }
616 }
617
618 pub fn value_plan(&self) -> Option<&ValuePlan> {
621 match self {
622 NodePlan::Extension { .. } | NodePlan::Empty => None,
623 NodePlan::Leaf { value, .. } => Some(value),
624 NodePlan::Branch { value, .. } | NodePlan::NibbledBranch { value, .. } =>
625 value.as_ref(),
626 }
627 }
628
629 pub fn value_plan_mut(&mut self) -> Option<&mut ValuePlan> {
632 match self {
633 NodePlan::Extension { .. } | NodePlan::Empty => None,
634 NodePlan::Leaf { value, .. } => Some(value),
635 NodePlan::Branch { value, .. } | NodePlan::NibbledBranch { value, .. } =>
636 value.as_mut(),
637 }
638 }
639}
640
641#[cfg_attr(feature = "std", derive(Debug))]
644#[derive(PartialEq, Eq)]
645pub struct OwnedNode<D: Borrow<[u8]>> {
646 data: D,
647 plan: NodePlan,
648}
649
650impl<D: Borrow<[u8]>> OwnedNode<D> {
651 pub fn new<C: NodeCodec>(data: D) -> core::result::Result<Self, C::Error> {
653 let plan = C::decode_plan(data.borrow())?;
654 Ok(OwnedNode { data, plan })
655 }
656
657 pub fn data(&self) -> &[u8] {
659 self.data.borrow()
660 }
661
662 pub fn node_plan(&self) -> &NodePlan {
664 &self.plan
665 }
666
667 pub fn node_plan_mut(&mut self) -> &mut NodePlan {
669 &mut self.plan
670 }
671
672 pub fn node(&self) -> Node {
674 self.plan.build(self.data.borrow())
675 }
676}