sp_trie/
trie_stream.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//! `TrieStream` implementation for Substrate's trie format.
19
20use crate::{
21	node_header::{size_and_prefix_iterator, NodeKind},
22	trie_constants,
23};
24use alloc::vec::Vec;
25use codec::{Compact, Encode};
26use hash_db::Hasher;
27use trie_root;
28
29/// Codec-flavored TrieStream.
30#[derive(Default, Clone)]
31pub struct TrieStream {
32	/// Current node buffer.
33	buffer: Vec<u8>,
34}
35
36impl TrieStream {
37	// useful for debugging but not used otherwise
38	pub fn as_raw(&self) -> &[u8] {
39		&self.buffer
40	}
41}
42
43fn branch_node_bit_mask(has_children: impl Iterator<Item = bool>) -> (u8, u8) {
44	let mut bitmap: u16 = 0;
45	let mut cursor: u16 = 1;
46	for v in has_children {
47		if v {
48			bitmap |= cursor
49		}
50		cursor <<= 1;
51	}
52	((bitmap % 256) as u8, (bitmap / 256) as u8)
53}
54
55/// Create a leaf/branch node, encoding a number of nibbles.
56fn fuse_nibbles_node(nibbles: &[u8], kind: NodeKind) -> impl Iterator<Item = u8> + '_ {
57	let size = nibbles.len();
58	let iter_start = match kind {
59		NodeKind::Leaf => size_and_prefix_iterator(size, trie_constants::LEAF_PREFIX_MASK, 2),
60		NodeKind::BranchNoValue => {
61			size_and_prefix_iterator(size, trie_constants::BRANCH_WITHOUT_MASK, 2)
62		},
63		NodeKind::BranchWithValue => {
64			size_and_prefix_iterator(size, trie_constants::BRANCH_WITH_MASK, 2)
65		},
66		NodeKind::HashedValueLeaf => {
67			size_and_prefix_iterator(size, trie_constants::ALT_HASHING_LEAF_PREFIX_MASK, 3)
68		},
69		NodeKind::HashedValueBranch => {
70			size_and_prefix_iterator(size, trie_constants::ALT_HASHING_BRANCH_WITH_MASK, 4)
71		},
72	};
73	iter_start
74		.chain(if nibbles.len() % 2 == 1 { Some(nibbles[0]) } else { None })
75		.chain(nibbles[nibbles.len() % 2..].chunks(2).map(|ch| (ch[0] << 4) | ch[1]))
76}
77
78use trie_root::Value as TrieStreamValue;
79impl trie_root::TrieStream for TrieStream {
80	fn new() -> Self {
81		Self { buffer: Vec::new() }
82	}
83
84	fn append_empty_data(&mut self) {
85		self.buffer.push(trie_constants::EMPTY_TRIE);
86	}
87
88	fn append_leaf(&mut self, key: &[u8], value: TrieStreamValue) {
89		let kind = match &value {
90			TrieStreamValue::Inline(..) => NodeKind::Leaf,
91			TrieStreamValue::Node(..) => NodeKind::HashedValueLeaf,
92		};
93		self.buffer.extend(fuse_nibbles_node(key, kind));
94		match &value {
95			TrieStreamValue::Inline(value) => {
96				Compact(value.len() as u32).encode_to(&mut self.buffer);
97				self.buffer.extend_from_slice(value);
98			},
99			TrieStreamValue::Node(hash) => {
100				self.buffer.extend_from_slice(hash.as_slice());
101			},
102		};
103	}
104
105	fn begin_branch(
106		&mut self,
107		maybe_partial: Option<&[u8]>,
108		maybe_value: Option<TrieStreamValue>,
109		has_children: impl Iterator<Item = bool>,
110	) {
111		if let Some(partial) = maybe_partial {
112			let kind = match &maybe_value {
113				None => NodeKind::BranchNoValue,
114				Some(TrieStreamValue::Inline(..)) => NodeKind::BranchWithValue,
115				Some(TrieStreamValue::Node(..)) => NodeKind::HashedValueBranch,
116			};
117
118			self.buffer.extend(fuse_nibbles_node(partial, kind));
119			let bm = branch_node_bit_mask(has_children);
120			self.buffer.extend([bm.0, bm.1].iter());
121		} else {
122			unreachable!("trie stream codec only for no extension trie");
123		}
124		match maybe_value {
125			None => (),
126			Some(TrieStreamValue::Inline(value)) => {
127				Compact(value.len() as u32).encode_to(&mut self.buffer);
128				self.buffer.extend_from_slice(value);
129			},
130			Some(TrieStreamValue::Node(hash)) => {
131				self.buffer.extend_from_slice(hash.as_slice());
132			},
133		}
134	}
135
136	fn append_extension(&mut self, _key: &[u8]) {
137		debug_assert!(false, "trie stream codec only for no extension trie");
138	}
139
140	fn append_substream<H: Hasher>(&mut self, other: Self) {
141		let data = other.out();
142		match data.len() {
143			0..=31 => data.encode_to(&mut self.buffer),
144			_ => H::hash(&data).as_ref().encode_to(&mut self.buffer),
145		}
146	}
147
148	fn out(self) -> Vec<u8> {
149		self.buffer
150	}
151}
sp_trie/trie_stream.rs

sp_trie/
trie_stream.rs
