push_decoder/
merkle.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
use alloc::vec::Vec;

use bitcoin::consensus::Encodable;
use bitcoin::hashes::Hash;
use bitcoin::hash_types::TxMerkleNode;

/// An incremental merkle tree hasher
pub struct IncrementalHasher {
    stack: Vec<Option<TxMerkleNode>>,
}

fn hash(left: TxMerkleNode, right: TxMerkleNode) -> TxMerkleNode {
    let mut encoder = TxMerkleNode::engine();
    left.consensus_encode(&mut encoder)
        .expect("in-memory writers don't error");
    right
        .consensus_encode(&mut encoder)
        .expect("in-memory writers don't error");
    TxMerkleNode::from_engine(encoder)
}

impl IncrementalHasher {
    /// Create a hasher
    pub fn new() -> Self {
        Self { stack: Vec::new() }
    }

    /// Add a node to the tree
    pub fn add(&mut self, mut node: TxMerkleNode) {
        for height in 0.. {
            if self.stack.len() <= height {
                self.stack.push(Some(node));
                break;
            }
            let left = self.stack[height].take();
            if let Some(left) = left {
                node = hash(left, node);
            } else {
                self.stack[height] = Some(node);
                break;
            }
        }
    }

    /// Finish hashing the tree and return the root.
    /// Will panic if no nodes have been added.
    pub fn finish(self) -> TxMerkleNode {
        let height = self.stack.len();

        let mut node = None;
        for (h, left) in self.stack.into_iter().enumerate() {
            let is_last = h == height - 1;
            if let Some(left) = left {
                if let Some(right) = node {
                    node = Some(hash(left, right));
                } else {
                    node = Some(if is_last { left } else { hash(left, left) });
                }
            } else {
                if let Some(single) = node {
                    node = Some(if is_last {
                        single
                    } else {
                        hash(single, single)
                    });
                }
            }
        }
        node.expect("empty merkle tree")
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bitcoin::hashes::Hash;
    use bitcoin::hash_types::TxMerkleNode;

    #[test]
    fn test_merkle() {
        for len in 1..63 {
            run_one(len);
        }
    }

    fn run_one(len: usize) {
        let nodes = (0..len)
            .map(|i| TxMerkleNode::from_slice(&[i as u8; 32]).unwrap())
            .collect::<Vec<_>>();
        let root = bitcoin::merkle_tree::calculate_root(nodes.clone().into_iter());
        let mut incremental = IncrementalHasher::new();
        for node in nodes.iter() {
            incremental.add(*node);
        }
        let incremental_root = incremental.finish();
        assert_eq!(root, Some(incremental_root), "mismatch for len={}", len);
    }
}