jmt/types/
nibble.rs

1// Copyright (c) The Diem Core Contributors
2// SPDX-License-Identifier: Apache-2.0
3
4#![forbid(unsafe_code)]
5
6//! `Nibble` represents a four-bit unsigned integer.
7
8pub mod nibble_path;
9
10use core::fmt;
11
12#[cfg(any(test))]
13use proptest::prelude::*;
14use serde::{Deserialize, Serialize};
15
16use crate::{Bytes32Ext, KeyHash, ValueHash};
17
18/// The hardcoded maximum height of a state merkle tree in nibbles.
19pub const ROOT_NIBBLE_HEIGHT: usize = 32 * 2;
20
21#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq, Ord, PartialOrd, Serialize, Deserialize)]
22pub struct Nibble(u8);
23
24impl From<u8> for Nibble {
25    fn from(nibble: u8) -> Self {
26        assert!(nibble < 16, "Nibble out of range: {}", nibble);
27        Self(nibble)
28    }
29}
30
31impl From<Nibble> for u8 {
32    fn from(nibble: Nibble) -> Self {
33        nibble.0
34    }
35}
36
37impl fmt::LowerHex for Nibble {
38    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
39        write!(f, "{:x}", self.0)
40    }
41}
42
43impl Nibble {
44    pub fn as_usize(self) -> usize {
45        self.0 as usize
46    }
47}
48
49/// An iterator that iterates the index range (inclusive) of each different nibble at given
50/// `nibble_idx` of all the keys in a sorted key-value pairs.
51pub(crate) struct NibbleRangeIterator<'a> {
52    sorted_kvs: &'a [(KeyHash, ValueHash)],
53    nibble_idx: usize,
54    pos: usize,
55}
56
57impl<'a> NibbleRangeIterator<'a> {
58    pub fn new(sorted_kvs: &'a [(KeyHash, ValueHash)], nibble_idx: usize) -> Self {
59        assert!(nibble_idx < ROOT_NIBBLE_HEIGHT);
60        NibbleRangeIterator {
61            sorted_kvs,
62            nibble_idx,
63            pos: 0,
64        }
65    }
66}
67
68impl<'a> core::iter::Iterator for NibbleRangeIterator<'a> {
69    type Item = (usize, usize);
70
71    fn next(&mut self) -> Option<Self::Item> {
72        let left = self.pos;
73        if self.pos < self.sorted_kvs.len() {
74            let cur_nibble: u8 = self.sorted_kvs[left].0 .0.nibble(self.nibble_idx);
75            let (mut i, mut j) = (left, self.sorted_kvs.len() - 1);
76            // Find the last index of the cur_nibble.
77            while i < j {
78                let mid = j - (j - i) / 2;
79                if self.sorted_kvs[mid].0 .0.nibble(self.nibble_idx) > cur_nibble {
80                    j = mid - 1;
81                } else {
82                    i = mid;
83                }
84            }
85            self.pos = i + 1;
86            Some((left, i))
87        } else {
88            None
89        }
90    }
91}
92
93#[cfg(any(test))]
94impl Arbitrary for Nibble {
95    type Parameters = ();
96    type Strategy = BoxedStrategy<Self>;
97
98    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
99        (0..16u8).prop_map(Self::from).boxed()
100    }
101}