tp_trie/
trie_stream.rs

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