sp_trie/
node_header.rs

1// This file is part of Substrate.
2
3// Copyright (C) Parity Technologies (UK) Ltd.
4// SPDX-License-Identifier: Apache-2.0
5
6// Licensed under the Apache License, Version 2.0 (the "License");
7// you may not use this file except in compliance with the License.
8// You may obtain a copy of the License at
9//
10// 	http://www.apache.org/licenses/LICENSE-2.0
11//
12// Unless required by applicable law or agreed to in writing, software
13// distributed under the License is distributed on an "AS IS" BASIS,
14// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15// See the License for the specific language governing permissions and
16// limitations under the License.
17
18//! The node header.
19
20use crate::trie_constants;
21use codec::{Decode, Encode, Input, Output};
22use core::iter::once;
23
24/// A node header
25#[derive(Copy, Clone, PartialEq, Eq, sp_core::RuntimeDebug)]
26pub(crate) enum NodeHeader {
27	Null,
28	// contains wether there is a value and nibble count
29	Branch(bool, usize),
30	// contains nibble count
31	Leaf(usize),
32	// contains nibble count.
33	HashedValueBranch(usize),
34	// contains nibble count.
35	HashedValueLeaf(usize),
36}
37
38impl NodeHeader {
39	pub(crate) fn contains_hash_of_value(&self) -> bool {
40		matches!(self, NodeHeader::HashedValueBranch(_) | NodeHeader::HashedValueLeaf(_))
41	}
42}
43
44/// NodeHeader without content
45pub(crate) enum NodeKind {
46	Leaf,
47	BranchNoValue,
48	BranchWithValue,
49	HashedValueLeaf,
50	HashedValueBranch,
51}
52
53impl Encode for NodeHeader {
54	fn encode_to<T: Output + ?Sized>(&self, output: &mut T) {
55		match self {
56			NodeHeader::Null => output.push_byte(trie_constants::EMPTY_TRIE),
57			NodeHeader::Branch(true, nibble_count) =>
58				encode_size_and_prefix(*nibble_count, trie_constants::BRANCH_WITH_MASK, 2, output),
59			NodeHeader::Branch(false, nibble_count) => encode_size_and_prefix(
60				*nibble_count,
61				trie_constants::BRANCH_WITHOUT_MASK,
62				2,
63				output,
64			),
65			NodeHeader::Leaf(nibble_count) =>
66				encode_size_and_prefix(*nibble_count, trie_constants::LEAF_PREFIX_MASK, 2, output),
67			NodeHeader::HashedValueBranch(nibble_count) => encode_size_and_prefix(
68				*nibble_count,
69				trie_constants::ALT_HASHING_BRANCH_WITH_MASK,
70				4,
71				output,
72			),
73			NodeHeader::HashedValueLeaf(nibble_count) => encode_size_and_prefix(
74				*nibble_count,
75				trie_constants::ALT_HASHING_LEAF_PREFIX_MASK,
76				3,
77				output,
78			),
79		}
80	}
81}
82
83impl codec::EncodeLike for NodeHeader {}
84
85impl Decode for NodeHeader {
86	fn decode<I: Input>(input: &mut I) -> Result<Self, codec::Error> {
87		let i = input.read_byte()?;
88		if i == trie_constants::EMPTY_TRIE {
89			return Ok(NodeHeader::Null)
90		}
91		match i & (0b11 << 6) {
92			trie_constants::LEAF_PREFIX_MASK => Ok(NodeHeader::Leaf(decode_size(i, input, 2)?)),
93			trie_constants::BRANCH_WITH_MASK =>
94				Ok(NodeHeader::Branch(true, decode_size(i, input, 2)?)),
95			trie_constants::BRANCH_WITHOUT_MASK =>
96				Ok(NodeHeader::Branch(false, decode_size(i, input, 2)?)),
97			trie_constants::EMPTY_TRIE => {
98				if i & (0b111 << 5) == trie_constants::ALT_HASHING_LEAF_PREFIX_MASK {
99					Ok(NodeHeader::HashedValueLeaf(decode_size(i, input, 3)?))
100				} else if i & (0b1111 << 4) == trie_constants::ALT_HASHING_BRANCH_WITH_MASK {
101					Ok(NodeHeader::HashedValueBranch(decode_size(i, input, 4)?))
102				} else {
103					// do not allow any special encoding
104					Err("Unallowed encoding".into())
105				}
106			},
107			_ => unreachable!(),
108		}
109	}
110}
111
112/// Returns an iterator over encoded bytes for node header and size.
113/// Size encoding allows unlimited, length inefficient, representation, but
114/// is bounded to 16 bit maximum value to avoid possible DOS.
115pub(crate) fn size_and_prefix_iterator(
116	size: usize,
117	prefix: u8,
118	prefix_mask: usize,
119) -> impl Iterator<Item = u8> {
120	let max_value = 255u8 >> prefix_mask;
121	let l1 = core::cmp::min((max_value as usize).saturating_sub(1), size);
122	let (first_byte, mut rem) = if size == l1 {
123		(once(prefix + l1 as u8), 0)
124	} else {
125		(once(prefix + max_value as u8), size - l1)
126	};
127	let next_bytes = move || {
128		if rem > 0 {
129			if rem < 256 {
130				let result = rem - 1;
131				rem = 0;
132				Some(result as u8)
133			} else {
134				rem = rem.saturating_sub(255);
135				Some(255)
136			}
137		} else {
138			None
139		}
140	};
141	first_byte.chain(core::iter::from_fn(next_bytes))
142}
143
144/// Encodes size and prefix to a stream output.
145fn encode_size_and_prefix<W>(size: usize, prefix: u8, prefix_mask: usize, out: &mut W)
146where
147	W: Output + ?Sized,
148{
149	for b in size_and_prefix_iterator(size, prefix, prefix_mask) {
150		out.push_byte(b)
151	}
152}
153
154/// Decode size only from stream input and header byte.
155fn decode_size(
156	first: u8,
157	input: &mut impl Input,
158	prefix_mask: usize,
159) -> Result<usize, codec::Error> {
160	let max_value = 255u8 >> prefix_mask;
161	let mut result = (first & max_value) as usize;
162	if result < max_value as usize {
163		return Ok(result)
164	}
165	result -= 1;
166	loop {
167		let n = input.read_byte()? as usize;
168		if n < 255 {
169			return Ok(result + n + 1)
170		}
171		result += 255;
172	}
173}
sp_trie/node_header.rs

sp_trie/
node_header.rs
