trie_db/nibble/
nibblevec.rs

1// Copyright 2017, 2018 Parity Technologies
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! An owning, nibble-oriented byte vector.
16
17use super::NibbleVec;
18use crate::{
19	nibble::{nibble_ops, BackingByteVec, NibbleSlice},
20	node::NodeKey,
21	node_codec::Partial,
22};
23use hash_db::Prefix;
24
25impl Default for NibbleVec {
26	fn default() -> Self {
27		NibbleVec::new()
28	}
29}
30
31impl NibbleVec {
32	/// Make a new `NibbleVec`.
33	pub fn new() -> Self {
34		NibbleVec { inner: BackingByteVec::new(), len: 0 }
35	}
36
37	/// Length of the `NibbleVec`.
38	#[inline(always)]
39	pub fn len(&self) -> usize {
40		self.len
41	}
42
43	/// Returns true if `NibbleVec` has zero length.
44	pub fn is_empty(&self) -> bool {
45		self.len == 0
46	}
47
48	/// Try to get the nibble at the given offset.
49	#[inline]
50	pub fn at(&self, idx: usize) -> u8 {
51		let ix = idx / nibble_ops::NIBBLE_PER_BYTE;
52		let pad = idx % nibble_ops::NIBBLE_PER_BYTE;
53		nibble_ops::at_left(pad as u8, self.inner[ix])
54	}
55
56	/// Push a nibble onto the `NibbleVec`. Ignores the high 4 bits.
57	pub fn push(&mut self, nibble: u8) {
58		let i = self.len % nibble_ops::NIBBLE_PER_BYTE;
59
60		if i == 0 {
61			self.inner.push(nibble_ops::push_at_left(0, nibble, 0));
62		} else {
63			let output = self
64				.inner
65				.last_mut()
66				.expect("len != 0 since len % 2 != 0; inner has a last element; qed");
67			*output = nibble_ops::push_at_left(i as u8, nibble, *output);
68		}
69		self.len += 1;
70	}
71
72	/// Try to pop a nibble off the `NibbleVec`. Fails if len == 0.
73	pub fn pop(&mut self) -> Option<u8> {
74		if self.is_empty() {
75			return None
76		}
77		let byte = self.inner.pop().expect("len != 0; inner has last elem; qed");
78		self.len -= 1;
79		let i_new = self.len % nibble_ops::NIBBLE_PER_BYTE;
80		if i_new != 0 {
81			self.inner.push(nibble_ops::pad_left(byte));
82		}
83		Some(nibble_ops::at_left(i_new as u8, byte))
84	}
85
86	/// Remove then n last nibbles in a faster way than popping n times.
87	pub fn drop_lasts(&mut self, n: usize) {
88		if n == 0 {
89			return
90		}
91		if n >= self.len {
92			self.clear();
93			return
94		}
95		let end = self.len - n;
96		let end_index = end / nibble_ops::NIBBLE_PER_BYTE +
97			if end % nibble_ops::NIBBLE_PER_BYTE == 0 { 0 } else { 1 };
98		(end_index..self.inner.len()).for_each(|_| {
99			self.inner.pop();
100		});
101		self.len = end;
102		let pos = self.len % nibble_ops::NIBBLE_PER_BYTE;
103		if pos != 0 {
104			let kl = self.inner.len() - 1;
105			self.inner[kl] = nibble_ops::pad_left(self.inner[kl]);
106		}
107	}
108
109	/// Get `Prefix` representation of this `NibbleVec`.
110	pub fn as_prefix(&self) -> Prefix {
111		let split = self.len / nibble_ops::NIBBLE_PER_BYTE;
112		let pos = (self.len % nibble_ops::NIBBLE_PER_BYTE) as u8;
113		if pos == 0 {
114			(&self.inner[..split], None)
115		} else {
116			(&self.inner[..split], Some(nibble_ops::pad_left(self.inner[split])))
117		}
118	}
119
120	/// Append another `NibbleVec`. Can be slow (alignement of second vec).
121	pub fn append(&mut self, v: &NibbleVec) {
122		if v.len == 0 {
123			return
124		}
125
126		let final_len = self.len + v.len;
127		let offset = self.len % nibble_ops::NIBBLE_PER_BYTE;
128		let final_offset = final_len % nibble_ops::NIBBLE_PER_BYTE;
129		let last_index = self.len / nibble_ops::NIBBLE_PER_BYTE;
130		if offset > 0 {
131			let (s1, s2) = nibble_ops::SPLIT_SHIFTS;
132			self.inner[last_index] =
133				nibble_ops::pad_left(self.inner[last_index]) | (v.inner[0] >> s2);
134			(0..v.inner.len() - 1)
135				.for_each(|i| self.inner.push(v.inner[i] << s1 | v.inner[i + 1] >> s2));
136			if final_offset > 0 {
137				self.inner.push(v.inner[v.inner.len() - 1] << s1);
138			}
139		} else {
140			(0..v.inner.len()).for_each(|i| self.inner.push(v.inner[i]));
141		}
142		self.len += v.len;
143	}
144
145	/// Append a `Partial`. Can be slow (alignement of partial).
146	pub fn append_partial(&mut self, (start_byte, sl): Partial) {
147		if start_byte.0 == 1 {
148			self.push(nibble_ops::at_left(1, start_byte.1));
149		}
150		let pad = self.inner.len() * nibble_ops::NIBBLE_PER_BYTE - self.len;
151		if pad == 0 {
152			self.inner.extend_from_slice(&sl[..]);
153		} else {
154			let kend = self.inner.len() - 1;
155			if sl.len() > 0 {
156				self.inner[kend] = nibble_ops::pad_left(self.inner[kend]);
157				let (s1, s2) = nibble_ops::SPLIT_SHIFTS;
158				self.inner[kend] |= sl[0] >> s1;
159				(0..sl.len() - 1).for_each(|i| self.inner.push(sl[i] << s2 | sl[i + 1] >> s1));
160				self.inner.push(sl[sl.len() - 1] << s2);
161			}
162		}
163		self.len += sl.len() * nibble_ops::NIBBLE_PER_BYTE;
164	}
165
166	/// Utility function for chaining two optional appending
167	/// of `NibbleSlice` and/or a byte.
168	/// Can be slow.
169	pub(crate) fn append_optional_slice_and_nibble(
170		&mut self,
171		o_slice: Option<&NibbleSlice>,
172		o_index: Option<u8>,
173	) -> usize {
174		let mut res = 0;
175		if let Some(slice) = o_slice {
176			self.append_partial(slice.right());
177			res += slice.len();
178		}
179		if let Some(ix) = o_index {
180			self.push(ix);
181			res += 1;
182		}
183		res
184	}
185
186	/// Utility function for `append_optional_slice_and_nibble` after a clone.
187	/// Can be slow.
188	#[cfg(feature = "std")]
189	pub(crate) fn clone_append_optional_slice_and_nibble(
190		&self,
191		o_slice: Option<&NibbleSlice>,
192		o_index: Option<u8>,
193	) -> Self {
194		let mut p = self.clone();
195		p.append_optional_slice_and_nibble(o_slice, o_index);
196		p
197	}
198
199	/// Get the underlying byte slice.
200	pub fn inner(&self) -> &[u8] {
201		&self.inner[..]
202	}
203
204	/// clear
205	pub fn clear(&mut self) {
206		self.inner.clear();
207		self.len = 0;
208	}
209
210	/// Try to treat this `NibbleVec` as a `NibbleSlice`. Works only if there is no padding.
211	pub fn as_nibbleslice(&self) -> Option<NibbleSlice> {
212		if self.len % nibble_ops::NIBBLE_PER_BYTE == 0 {
213			Some(NibbleSlice::new(self.inner()))
214		} else {
215			None
216		}
217	}
218
219	/// Do we start with the same nibbles as the whole of `them`?
220	pub fn starts_with(&self, other: &Self) -> bool {
221		if self.len() < other.len() {
222			return false
223		}
224		let byte_len = other.len() / nibble_ops::NIBBLE_PER_BYTE;
225		if &self.inner[..byte_len] != &other.inner[..byte_len] {
226			return false
227		}
228		for pad in 0..(other.len() - byte_len * nibble_ops::NIBBLE_PER_BYTE) {
229			let self_nibble = nibble_ops::at_left(pad as u8, self.inner[byte_len]);
230			let other_nibble = nibble_ops::at_left(pad as u8, other.inner[byte_len]);
231			if self_nibble != other_nibble {
232				return false
233			}
234		}
235		true
236	}
237
238	/// Same as [`Self::starts_with`] but using [`NibbleSlice`].
239	pub fn starts_with_slice(&self, other: &NibbleSlice) -> bool {
240		if self.len() < other.len() {
241			return false
242		}
243
244		match self.as_nibbleslice() {
245			Some(slice) => slice.starts_with(&other),
246			None => {
247				for i in 0..other.len() {
248					if self.at(i) != other.at(i) {
249						return false
250					}
251				}
252				true
253			},
254		}
255	}
256
257	/// Return an iterator over `Partial` bytes representation.
258	pub fn right_iter<'a>(&'a self) -> impl Iterator<Item = u8> + 'a {
259		let require_padding = self.len % nibble_ops::NIBBLE_PER_BYTE != 0;
260		let mut ix = 0;
261		let inner = &self.inner;
262
263		let (left_s, right_s) = nibble_ops::SPLIT_SHIFTS;
264
265		crate::rstd::iter::from_fn(move || {
266			if require_padding && ix < inner.len() {
267				if ix == 0 {
268					ix += 1;
269					Some(inner[ix - 1] >> nibble_ops::BIT_PER_NIBBLE)
270				} else {
271					ix += 1;
272
273					Some(inner[ix - 2] << left_s | inner[ix - 1] >> right_s)
274				}
275			} else if ix < inner.len() {
276				ix += 1;
277
278				Some(inner[ix - 1])
279			} else {
280				None
281			}
282		})
283	}
284}
285
286impl<'a> From<NibbleSlice<'a>> for NibbleVec {
287	fn from(s: NibbleSlice<'a>) -> Self {
288		let mut v = NibbleVec::new();
289		for i in 0..s.len() {
290			v.push(s.at(i));
291		}
292		v
293	}
294}
295
296impl From<&NibbleVec> for NodeKey {
297	fn from(nibble: &NibbleVec) -> NodeKey {
298		if let Some(slice) = nibble.as_nibbleslice() {
299			slice.into()
300		} else {
301			(1, nibble.right_iter().collect())
302		}
303	}
304}
305
306#[cfg(test)]
307mod tests {
308	use super::*;
309	use crate::{nibble::nibble_ops, NibbleSlice};
310
311	#[test]
312	fn push_pop() {
313		let mut v = NibbleVec::new();
314
315		for i in 0..(nibble_ops::NIBBLE_PER_BYTE * 3) {
316			let iu8 = (i % nibble_ops::NIBBLE_PER_BYTE) as u8;
317			v.push(iu8);
318			assert_eq!(v.len() - 1, i);
319			assert_eq!(v.at(i), iu8);
320		}
321
322		for i in (0..(nibble_ops::NIBBLE_PER_BYTE * 3)).rev() {
323			let iu8 = (i % nibble_ops::NIBBLE_PER_BYTE) as u8;
324			let a = v.pop();
325			assert_eq!(a, Some(iu8));
326			assert_eq!(v.len(), i);
327		}
328	}
329
330	#[test]
331	fn append_partial() {
332		append_partial_inner(&[1, 2, 3], &[], ((1, 1), &[0x23]));
333		append_partial_inner(&[1, 2, 3], &[1], ((0, 0), &[0x23]));
334		append_partial_inner(&[0, 1, 2, 3], &[0], ((1, 1), &[0x23]));
335	}
336
337	fn append_partial_inner(res: &[u8], init: &[u8], partial: ((u8, u8), &[u8])) {
338		let mut resv = NibbleVec::new();
339		res.iter().for_each(|r| resv.push(*r));
340		let mut initv = NibbleVec::new();
341		init.iter().for_each(|r| initv.push(*r));
342		initv.append_partial(partial);
343		assert_eq!(resv, initv);
344	}
345
346	#[test]
347	fn drop_lasts_test() {
348		let test_trun = |a: &[u8], b: usize, c: (&[u8], usize)| {
349			let mut k = NibbleVec::new();
350			for v in a {
351				k.push(*v);
352			}
353			k.drop_lasts(b);
354			assert_eq!((&k.inner[..], k.len), c);
355		};
356		test_trun(&[1, 2, 3, 4], 0, (&[0x12, 0x34], 4));
357		test_trun(&[1, 2, 3, 4], 1, (&[0x12, 0x30], 3));
358		test_trun(&[1, 2, 3, 4], 2, (&[0x12], 2));
359		test_trun(&[1, 2, 3, 4], 3, (&[0x10], 1));
360		test_trun(&[1, 2, 3, 4], 4, (&[], 0));
361		test_trun(&[1, 2, 3, 4], 5, (&[], 0));
362		test_trun(&[1, 2, 3], 0, (&[0x12, 0x30], 3));
363		test_trun(&[1, 2, 3], 1, (&[0x12], 2));
364		test_trun(&[1, 2, 3], 2, (&[0x10], 1));
365		test_trun(&[1, 2, 3], 3, (&[], 0));
366		test_trun(&[1, 2, 3], 4, (&[], 0));
367	}
368
369	#[test]
370	fn right_iter_works() {
371		let data = vec![1, 2, 3, 4, 5, 234, 78, 99];
372
373		let nibble = NibbleSlice::new(&data);
374		let vec = NibbleVec::from(nibble);
375
376		nibble
377			.right_iter()
378			.zip(vec.right_iter())
379			.enumerate()
380			.for_each(|(i, (l, r))| assert_eq!(l, r, "Don't match at {}", i));
381
382		// Also try with using an offset.
383		let nibble = NibbleSlice::new_offset(&data, 3);
384		let vec = NibbleVec::from(nibble);
385
386		nibble
387			.right_iter()
388			.zip(vec.right_iter())
389			.enumerate()
390			.for_each(|(i, (l, r))| assert_eq!(l, r, "Don't match at {}", i));
391	}
392}