sparse_merkle_tree/
merge.rs

1use crate::h256::H256;
2use crate::traits::Hasher;
3
4const MERGE_NORMAL: u8 = 1;
5const MERGE_ZEROS: u8 = 2;
6
7#[derive(Debug, Eq, PartialEq, Clone)]
8pub enum MergeValue {
9    Value(H256),
10    MergeWithZero {
11        base_node: H256,
12        zero_bits: H256,
13        zero_count: u8,
14    },
15    #[cfg(feature = "trie")]
16    ShortCut {
17        key: H256,
18        value: H256,
19        height: u8,
20    },
21}
22
23impl MergeValue {
24    pub fn from_h256(v: H256) -> Self {
25        MergeValue::Value(v)
26    }
27
28    pub fn zero() -> Self {
29        MergeValue::Value(H256::zero())
30    }
31
32    pub fn is_zero(&self) -> bool {
33        match self {
34            MergeValue::Value(v) => v.is_zero(),
35            MergeValue::MergeWithZero { .. } => false,
36            #[cfg(feature = "trie")]
37            MergeValue::ShortCut { .. } => false,
38        }
39    }
40
41    #[cfg(feature = "trie")]
42    pub fn shortcut_or_value(key: H256, value: H256, height: u8) -> Self {
43        if height == 0 || value.is_zero() {
44            MergeValue::Value(value)
45        } else {
46            MergeValue::ShortCut { key, value, height }
47        }
48    }
49
50    #[cfg(feature = "trie")]
51    pub fn is_shortcut(&self) -> bool {
52        matches!(self, MergeValue::ShortCut { .. })
53    }
54
55    pub fn hash<H: Hasher + Default>(&self) -> H256 {
56        match self {
57            MergeValue::Value(v) => *v,
58            MergeValue::MergeWithZero {
59                base_node,
60                zero_bits,
61                zero_count,
62            } => {
63                let mut hasher = H::default();
64                hasher.write_byte(MERGE_ZEROS);
65                hasher.write_h256(base_node);
66                hasher.write_h256(zero_bits);
67                hasher.write_byte(*zero_count);
68                hasher.finish()
69            }
70            #[cfg(feature = "trie")]
71            MergeValue::ShortCut { key, value, height } => {
72                into_merge_value::<H>(*key, *value, *height).hash::<H>()
73            }
74        }
75    }
76}
77
78/// Helper function for Shortcut node
79/// Transform it into a MergeValue or MergeWithZero node
80#[cfg(feature = "trie")]
81pub fn into_merge_value<H: Hasher + Default>(key: H256, value: H256, height: u8) -> MergeValue {
82    // try keep hash same with MergeWithZero
83    if value.is_zero() || height == 0 {
84        MergeValue::from_h256(value)
85    } else {
86        let base_key = key.parent_path(0);
87        let base_node = hash_base_node::<H>(0, &base_key, &value);
88        let mut zero_bits = key;
89        for i in height..=core::u8::MAX {
90            if key.get_bit(i) {
91                zero_bits.clear_bit(i);
92            }
93        }
94        MergeValue::MergeWithZero {
95            base_node,
96            zero_bits,
97            zero_count: height,
98        }
99    }
100}
101
102/// Hash base node into a H256
103pub fn hash_base_node<H: Hasher + Default>(
104    base_height: u8,
105    base_key: &H256,
106    base_value: &H256,
107) -> H256 {
108    let mut hasher = H::default();
109    hasher.write_byte(base_height);
110    hasher.write_h256(base_key);
111    hasher.write_h256(base_value);
112    hasher.finish()
113}
114
115/// Merge two hash with node information
116/// this function optimized for ZERO_HASH
117/// if lhs and rhs both are ZERO_HASH return ZERO_HASH, otherwise hash all info.
118pub fn merge<H: Hasher + Default>(
119    height: u8,
120    node_key: &H256,
121    lhs: &MergeValue,
122    rhs: &MergeValue,
123) -> MergeValue {
124    if lhs.is_zero() && rhs.is_zero() {
125        return MergeValue::zero();
126    }
127    if lhs.is_zero() {
128        return merge_with_zero::<H>(height, node_key, rhs, true);
129    }
130    if rhs.is_zero() {
131        return merge_with_zero::<H>(height, node_key, lhs, false);
132    }
133    let mut hasher = H::default();
134    hasher.write_byte(MERGE_NORMAL);
135    hasher.write_byte(height);
136    hasher.write_h256(node_key);
137    hasher.write_h256(&lhs.hash::<H>());
138    hasher.write_h256(&rhs.hash::<H>());
139    MergeValue::Value(hasher.finish())
140}
141
142pub fn merge_with_zero<H: Hasher + Default>(
143    height: u8,
144    node_key: &H256,
145    value: &MergeValue,
146    set_bit: bool,
147) -> MergeValue {
148    match value {
149        MergeValue::Value(v) => {
150            let mut zero_bits = H256::zero();
151            if set_bit {
152                zero_bits.set_bit(height);
153            }
154            let base_node = hash_base_node::<H>(height, node_key, v);
155            MergeValue::MergeWithZero {
156                base_node,
157                zero_bits,
158                zero_count: 1,
159            }
160        }
161        MergeValue::MergeWithZero {
162            base_node,
163            zero_bits,
164            zero_count,
165        } => {
166            let mut zero_bits = *zero_bits;
167            if set_bit {
168                zero_bits.set_bit(height);
169            }
170            MergeValue::MergeWithZero {
171                base_node: *base_node,
172                zero_bits,
173                zero_count: zero_count.wrapping_add(1),
174            }
175        }
176        #[cfg(feature = "trie")]
177        MergeValue::ShortCut { key, value, .. } => {
178            if height == core::u8::MAX {
179                let base_key = key.parent_path(0);
180                let base_node = hash_base_node::<H>(0, &base_key, value);
181                MergeValue::MergeWithZero {
182                    base_node,
183                    zero_bits: *key,
184                    zero_count: 0,
185                }
186            } else {
187                MergeValue::ShortCut {
188                    key: *key,
189                    value: *value,
190                    height: height + 1,
191                }
192            }
193        }
194    }
195}
sparse_merkle_tree/merge.rs

sparse_merkle_tree/
merge.rs
