basecoin_store/avl/
node.rs

1use core::borrow::Borrow;
2use core::mem;
3
4use sha2::{Digest, Sha256};
5use tendermint::hash::Hash;
6
7use super::proof::EMPTY_CHILD;
8use crate::avl::as_bytes::AsBytes;
9use crate::avl::{proof, HASH_ALGO};
10
11pub type NodeRef<T, V> = Option<Box<AvlNode<T, V>>>;
12
13/// A node in the AVL Tree.
14#[derive(Eq, PartialEq, Debug, Clone)]
15pub struct AvlNode<K: Ord, V> {
16    pub key: K,
17    pub value: V,
18    pub hash: Hash,
19    pub merkle_hash: Hash,
20    pub height: u32,
21    pub left: NodeRef<K, V>,
22    pub right: NodeRef<K, V>,
23}
24
25/// Wrap a key + value couple into a `NodeRef`.
26#[allow(clippy::unnecessary_wraps)]
27pub fn as_node_ref<K: Ord + AsBytes, V>(key: K, value: V) -> NodeRef<K, V>
28where
29    V: Borrow<[u8]>,
30{
31    Some(Box::new(AvlNode::new(key, value)))
32}
33
34impl<K: Ord + AsBytes, V> AvlNode<K, V>
35where
36    V: Borrow<[u8]>,
37{
38    fn new(key: K, value: V) -> Self {
39        let mut sha = Sha256::new();
40        sha.update(proof::LEAF_PREFIX);
41        sha.update(key.as_bytes().as_ref());
42        sha.update(value.borrow());
43        let hash = sha.finalize();
44
45        let mut sha = Sha256::new();
46        sha.update(EMPTY_CHILD);
47        sha.update(hash);
48        sha.update(EMPTY_CHILD);
49        let merkle_hash = sha.finalize();
50
51        let merkle_hash = Hash::from_bytes(HASH_ALGO, &merkle_hash).unwrap();
52        let hash = Hash::from_bytes(HASH_ALGO, &hash).unwrap();
53
54        AvlNode {
55            key,
56            value,
57            hash,
58            merkle_hash,
59            height: 0,
60            left: None,
61            right: None,
62        }
63    }
64
65    /// Set the value of the current node.
66    pub(crate) fn set_value(&mut self, value: V) -> V {
67        let hash = Self::local_hash(&self.key, &value);
68        self.hash = hash;
69        mem::replace(&mut self.value, value)
70    }
71
72    /// The left height, or `None` if there is no left child.
73    fn left_height(&self) -> Option<u32> {
74        self.left.as_ref().map(|left| left.height)
75    }
76
77    /// The right height, or `None` if there is no right child.
78    fn right_height(&self) -> Option<u32> {
79        self.right.as_ref().map(|right| right.height)
80    }
81
82    /// Compute the local hash for a given key and value.
83    fn local_hash(key: &K, value: &V) -> Hash {
84        let mut sha = Sha256::new();
85        sha.update(proof::LEAF_PREFIX);
86        sha.update(key.as_bytes());
87        sha.update(value.borrow());
88        let hash = sha.finalize();
89        Hash::from_bytes(HASH_ALGO, &hash).unwrap()
90    }
91
92    /// The left merkle hash, if any
93    pub fn left_hash(&self) -> Option<&[u8]> {
94        Some(self.left.as_ref()?.merkle_hash.as_bytes())
95    }
96
97    /// The right merkle hash, if any
98    pub fn right_hash(&self) -> Option<&[u8]> {
99        Some(self.right.as_ref()?.merkle_hash.as_bytes())
100    }
101
102    /// Update the height of this node by looking at the height of its two children.
103    /// The height of this node is computed as the maximum among the height of its two children, and
104    /// incremented by 1.
105    fn update_height(&mut self) {
106        match &self.right {
107            None => match &self.left {
108                None => self.height = 0,
109                Some(left) => self.height = left.height + 1,
110            },
111            Some(right) => match &self.left {
112                None => self.height = right.height + 1,
113                Some(left) => self.height = core::cmp::max(left.height, right.height) + 1,
114            },
115        }
116    }
117
118    /// Update the node's merkle hash by looking at the hashes of its two children.
119    fn update_hashes(&mut self) {
120        let mut sha = Sha256::new();
121        sha.update(self.left_hash().unwrap_or(&EMPTY_CHILD));
122        sha.update(self.hash.as_bytes());
123        sha.update(self.right_hash().unwrap_or(&EMPTY_CHILD));
124        self.merkle_hash = Hash::from_bytes(HASH_ALGO, sha.finalize().as_slice()).unwrap();
125    }
126
127    /// Update node meta data, such as its height and merkle hash, by looking at its two
128    /// children.
129    pub fn update(&mut self) {
130        self.update_hashes();
131        self.update_height();
132    }
133
134    /// Returns the node's balance factor (left_height - right_height).
135    pub fn balance_factor(&self) -> i32 {
136        match (self.left_height(), self.right_height()) {
137            (None, None) => 0,
138            (None, Some(h)) => -(h as i32),
139            (Some(h), None) => h as i32,
140            (Some(h_l), Some(h_r)) => (h_l as i32) - (h_r as i32),
141        }
142    }
143}