sp_trie/
node_codec.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//! `NodeCodec` implementation for Substrate's trie format.
19
20use 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
32/// Helper struct for trie node decoder. This implements `codec::Input` on a byte slice, while
33/// tracking the absolute position. This is similar to `std::io::Cursor` but does not implement
34/// `Read` and `io` are not in `core` or `alloc`.
35struct 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/// Concrete implementation of a [`NodeCodecT`] with SCALE encoding.
79///
80/// It is generic over `H` the [`Hasher`].
81#[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			// hashed_value_branch
106			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				// check that the padding is valid (if any)
114				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				// check that the padding is valid (if any)
158				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
274// utils
275
276/// Encode and allocate node type header (type and size), and partial value.
277/// It uses an iterator over encoded partial bytes as input.
278fn 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
299/// Radix 16 trie, bitmap encoding implementation,
300/// it contains children mapping information for a branch
301/// (children presence only), it encodes into
302/// a compact bitmap encoding representation.
303pub(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}