clvm_utils/
tree_hash.rs

1use chia_sha2::Sha256;
2use clvmr::allocator::{Allocator, NodePtr, SExp};
3use clvmr::serde::node_from_bytes_backrefs_record;
4use std::collections::{HashMap, HashSet};
5use std::ops::Deref;
6use std::{fmt, io};
7
8#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub struct TreeHash([u8; 32]);
10
11impl TreeHash {
12    pub const fn new(hash: [u8; 32]) -> Self {
13        Self(hash)
14    }
15
16    pub const fn to_bytes(&self) -> [u8; 32] {
17        self.0
18    }
19
20    pub fn to_vec(&self) -> Vec<u8> {
21        self.0.to_vec()
22    }
23}
24
25impl fmt::Debug for TreeHash {
26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
27        write!(f, "TreeHash({self})")
28    }
29}
30
31impl fmt::Display for TreeHash {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        write!(f, "{}", hex::encode(self.0))
34    }
35}
36
37impl From<[u8; 32]> for TreeHash {
38    fn from(hash: [u8; 32]) -> Self {
39        Self::new(hash)
40    }
41}
42
43impl From<TreeHash> for [u8; 32] {
44    fn from(hash: TreeHash) -> [u8; 32] {
45        hash.0
46    }
47}
48
49impl AsRef<[u8]> for TreeHash {
50    fn as_ref(&self) -> &[u8] {
51        &self.0
52    }
53}
54
55impl Deref for TreeHash {
56    type Target = [u8];
57
58    fn deref(&self) -> &Self::Target {
59        &self.0
60    }
61}
62
63enum TreeOp {
64    SExp(NodePtr),
65    Cons,
66    ConsAddCache(NodePtr),
67}
68
69pub fn tree_hash_atom(bytes: &[u8]) -> TreeHash {
70    let mut sha256 = Sha256::new();
71    sha256.update([1]);
72    sha256.update(bytes);
73    TreeHash::new(sha256.finalize())
74}
75
76pub fn tree_hash_pair(first: TreeHash, rest: TreeHash) -> TreeHash {
77    let mut sha256 = Sha256::new();
78    sha256.update([2]);
79    sha256.update(first);
80    sha256.update(rest);
81    TreeHash::new(sha256.finalize())
82}
83
84pub fn tree_hash(a: &Allocator, node: NodePtr) -> TreeHash {
85    let mut hashes = Vec::new();
86    let mut ops = vec![TreeOp::SExp(node)];
87
88    while let Some(op) = ops.pop() {
89        match op {
90            TreeOp::SExp(node) => match a.sexp(node) {
91                SExp::Atom => {
92                    hashes.push(tree_hash_atom(a.atom(node).as_ref()));
93                }
94                SExp::Pair(left, right) => {
95                    ops.push(TreeOp::Cons);
96                    ops.push(TreeOp::SExp(left));
97                    ops.push(TreeOp::SExp(right));
98                }
99            },
100            TreeOp::Cons => {
101                let first = hashes.pop().unwrap();
102                let rest = hashes.pop().unwrap();
103                hashes.push(tree_hash_pair(first, rest));
104            }
105            TreeOp::ConsAddCache(_) => unreachable!(),
106        }
107    }
108
109    assert!(hashes.len() == 1);
110    hashes[0]
111}
112
113pub fn tree_hash_cached<S>(
114    a: &Allocator,
115    node: NodePtr,
116    backrefs: &HashSet<NodePtr, S>,
117    cache: &mut HashMap<NodePtr, TreeHash, S>,
118) -> TreeHash
119where
120    S: std::hash::BuildHasher,
121{
122    let mut hashes = Vec::new();
123    let mut ops = vec![TreeOp::SExp(node)];
124
125    while let Some(op) = ops.pop() {
126        match op {
127            TreeOp::SExp(node) => match a.sexp(node) {
128                SExp::Atom => {
129                    let hash = tree_hash_atom(a.atom(node).as_ref());
130                    if backrefs.contains(&node) {
131                        cache.insert(node, hash);
132                    }
133                    hashes.push(hash);
134                }
135                SExp::Pair(left, right) => {
136                    if let Some(hash) = cache.get(&node) {
137                        hashes.push(*hash);
138                    } else {
139                        if backrefs.contains(&node) {
140                            ops.push(TreeOp::ConsAddCache(node));
141                        } else {
142                            ops.push(TreeOp::Cons);
143                        }
144                        ops.push(TreeOp::SExp(left));
145                        ops.push(TreeOp::SExp(right));
146                    }
147                }
148            },
149            TreeOp::Cons => {
150                let first = hashes.pop().unwrap();
151                let rest = hashes.pop().unwrap();
152                hashes.push(tree_hash_pair(first, rest));
153            }
154            TreeOp::ConsAddCache(original_node) => {
155                let first = hashes.pop().unwrap();
156                let rest = hashes.pop().unwrap();
157                let hash = tree_hash_pair(first, rest);
158                hashes.push(hash);
159                cache.insert(original_node, hash);
160            }
161        }
162    }
163
164    assert!(hashes.len() == 1);
165    hashes[0]
166}
167
168pub fn tree_hash_from_bytes(buf: &[u8]) -> io::Result<TreeHash> {
169    let mut a = Allocator::new();
170    let (node, backrefs) = node_from_bytes_backrefs_record(&mut a, buf)?;
171    let mut cache = HashMap::<NodePtr, TreeHash>::new();
172    Ok(tree_hash_cached(&a, node, &backrefs, &mut cache))
173}
174
175#[test]
176fn test_tree_hash() {
177    let mut a = Allocator::new();
178    let atom1 = a.new_atom(&[1, 2, 3]).unwrap();
179    let atom2 = a.new_atom(&[4, 5, 6]).unwrap();
180    let root = a.new_pair(atom1, atom2).unwrap();
181
182    // test atom1 hash
183    let atom1_hash = {
184        let mut sha256 = Sha256::new();
185        sha256.update([1_u8]);
186        sha256.update([1, 2, 3]);
187        let atom1_hash = sha256.finalize();
188
189        assert_eq!(tree_hash(&a, atom1).as_ref(), atom1_hash.as_slice());
190        atom1_hash
191    };
192
193    // test atom2 hash
194    let atom2_hash = {
195        let mut sha256 = Sha256::new();
196        sha256.update([1_u8]);
197        sha256.update([4, 5, 6]);
198        let atom2_hash = sha256.finalize();
199
200        assert_eq!(tree_hash(&a, atom2).as_ref(), atom2_hash.as_slice());
201        atom2_hash
202    };
203
204    // test tree hash
205    let root_hash = {
206        let mut sha256 = Sha256::new();
207        sha256.update([2_u8]);
208        sha256.update(atom1_hash.as_slice());
209        sha256.update(atom2_hash.as_slice());
210        let root_hash = sha256.finalize();
211
212        assert_eq!(tree_hash(&a, root).as_ref(), root_hash.as_slice());
213        root_hash
214    };
215
216    let atom3 = a.new_atom(&[7, 8, 9]).unwrap();
217    let root2 = a.new_pair(root, atom3).unwrap();
218
219    let atom3_hash = {
220        let mut sha256 = Sha256::new();
221        sha256.update([1_u8]);
222        sha256.update([7, 8, 9]);
223        sha256.finalize()
224    };
225
226    // test deeper tree hash
227    {
228        let mut sha256 = Sha256::new();
229        sha256.update([2_u8]);
230        sha256.update(root_hash.as_slice());
231        sha256.update(atom3_hash.as_slice());
232
233        assert_eq!(tree_hash(&a, root2).as_ref(), sha256.finalize().as_slice());
234    }
235}
236
237#[test]
238fn test_tree_hash_from_bytes() {
239    use clvmr::serde::{node_to_bytes, node_to_bytes_backrefs};
240
241    let mut a = Allocator::new();
242    let atom1 = a.new_atom(&[1, 2, 3]).unwrap();
243    let atom2 = a.new_atom(&[4, 5, 6]).unwrap();
244    let node1 = a.new_pair(atom1, atom2).unwrap();
245    let node2 = a.new_pair(atom2, atom1).unwrap();
246
247    let node1 = a.new_pair(node1, node1).unwrap();
248    let node2 = a.new_pair(node2, node2).unwrap();
249
250    let root = a.new_pair(node1, node2).unwrap();
251
252    let serialized_clvm = node_to_bytes(&a, root).expect("node_to_bytes");
253    let serialized_clvm_backrefs =
254        node_to_bytes_backrefs(&a, root).expect("node_to_bytes_backrefs");
255
256    let hash1 = tree_hash_from_bytes(&serialized_clvm).expect("tree_hash_from_bytes");
257    let hash2 = tree_hash_from_bytes(&serialized_clvm_backrefs).expect("tree_hash_from_bytes");
258    let hash3 = tree_hash(&a, root);
259
260    assert!(serialized_clvm.len() > serialized_clvm_backrefs.len());
261    assert_eq!(hash1, hash2);
262    assert_eq!(hash1, hash3);
263}
264
265#[cfg(test)]
266use rstest::rstest;
267
268#[cfg(test)]
269#[rstest]
270#[case(
271    "block-1ee588dc",
272    "1cba0b22b84b597d265d77fbabb57fada01d963f75dc3956a6166a2385997ef2"
273)]
274#[case(
275    "block-6fe59b24",
276    "540c5afac7c26728ed6b7891d8ce2f5b26009c4b0090d7035403c2425dc54e1d"
277)]
278#[case(
279    "block-b45268ac",
280    "7cc321f5554126c9f430afbc7dd9c804f5d34a248e3192f275f5d585ecf8e873"
281)]
282#[case(
283    "block-c2a8df0d",
284    "2e25efa524e420111006fee77f50fb8fbd725920a5312d5480af239d81ab5e7e"
285)]
286#[case(
287    "block-e5002df2",
288    "c179ece232dceef984ba000f7e5b67ee3092582668bf6178969df10845eb8b18"
289)]
290#[case(
291    "block-4671894",
292    "3750f0e1bde9fcb407135f974aa276a4580e1e76a47e6d8d9bb2911d0fe91db1"
293)]
294#[case(
295    "block-225758",
296    "880df94c3c9e0f7c26c42ae99723e683a4cd37e73f74c6322d1dfabaa1d64d93"
297)]
298#[case(
299    "block-834752",
300    "be755b8ef03d917b8bd37ae152792a7daa7de81bbb0eaa21c530571c2105c130"
301)]
302#[case(
303    "block-834752-compressed",
304    "be755b8ef03d917b8bd37ae152792a7daa7de81bbb0eaa21c530571c2105c130"
305)]
306#[case(
307    "block-834760",
308    "77558768f74c5f863b36232a1390843a63a397fc22da1321fea3a05eab67be2c"
309)]
310#[case(
311    "block-834761",
312    "4bac8b299c6545a37a825883c863b79ce850e7f6c8f1d2abeec2865f5450f1c5"
313)]
314#[case(
315    "block-834765",
316    "b915ec5f9f8ea723e0a99b035df206673369b802766dd76b6c8f4c15ab7bca2c"
317)]
318#[case(
319    "block-834766",
320    "409559c3395fb18a6c3390ccccd55e82162b1e68b867490a90ccbddf78147c9d"
321)]
322#[case(
323    "block-834768",
324    "905441945a9a56558337c8b7a536a6b9606ad63e11a265a938f301747ccfb7af"
325)]
326fn test_tree_hash_cached(
327    #[case] name: &str,
328    #[case] expect: &str,
329    #[values(true, false)] compressed: bool,
330) {
331    use clvmr::serde::{
332        node_from_bytes_backrefs, node_from_bytes_backrefs_record, node_to_bytes_backrefs,
333    };
334    use std::fs::read_to_string;
335
336    let filename = format!("../../generator-tests/{name}.txt");
337    println!("file: {filename}",);
338    let test_file = read_to_string(filename).expect("test file not found");
339    let generator = test_file.lines().next().expect("invalid test file");
340    let generator = hex::decode(generator).expect("invalid hex encoded generator");
341
342    let generator = if compressed {
343        let mut a = Allocator::new();
344        let node = node_from_bytes_backrefs(&mut a, &generator).expect("node_from_bytes_backrefs");
345        node_to_bytes_backrefs(&a, node).expect("node_to_bytes_backrefs")
346    } else {
347        generator
348    };
349
350    let mut a = Allocator::new();
351    let mut cache = HashMap::<NodePtr, TreeHash>::new();
352    let (node, backrefs) = node_from_bytes_backrefs_record(&mut a, &generator)
353        .expect("node_from_bytes_backrefs_records");
354
355    let hash1 = tree_hash(&a, node);
356    let hash2 = tree_hash_cached(&a, node, &backrefs, &mut cache);
357    // for (key, value) in cache.iter() {
358    //     println!("  {key:?}: {}", hex::encode(value));
359    // }
360    assert_eq!(hash1, hash2);
361    assert_eq!(hash1.as_ref(), hex::decode(expect).unwrap().as_slice());
362    assert!(!compressed || !backrefs.is_empty());
363}
364
365#[cfg(test)]
366fn test_sha256_atom(buf: &[u8]) {
367    let hash = tree_hash_atom(buf);
368
369    let mut hasher = Sha256::new();
370    hasher.update([1_u8]);
371    if !buf.is_empty() {
372        hasher.update(buf);
373    }
374
375    assert_eq!(hash.as_ref(), hasher.finalize().as_slice());
376}
377
378#[test]
379fn test_tree_hash_atom() {
380    test_sha256_atom(&[]);
381    for val in 0..255 {
382        test_sha256_atom(&[val]);
383    }
384
385    for val in 0..255 {
386        test_sha256_atom(&[0, val]);
387    }
388
389    for val in 0..255 {
390        test_sha256_atom(&[0xff, val]);
391    }
392}
clvm_utils/tree_hash.rs

clvm_utils/
tree_hash.rs
