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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125

use crate::{
    common::{Bytes32, StorageMap},
    sparse::{self, Primitive},
    storage::Mappable,
};

/// The table of the Sparse Merkle tree's nodes. [`MerkleTree`] works with it as a sparse merkle
/// tree, where the storage key is `Bytes32` and the value is the [`Buffer`](crate::sparse::Buffer)
/// (raw presentation of the [`Node`](crate::sparse::Node)).
pub struct NodesTable;

impl Mappable for NodesTable {
    type Key = Self::OwnedKey;
    type OwnedKey = Bytes32;
    type Value = Self::OwnedValue;
    type OwnedValue = Primitive;
}

type Storage = StorageMap<NodesTable>;
type SparseMerkleTree = sparse::MerkleTree<NodesTable, Storage>;

pub struct MerkleTree {
    tree: SparseMerkleTree,
}

impl MerkleTree {
    pub fn new() -> Self {
        Self {
            tree: SparseMerkleTree::new(Storage::new()),
        }
    }

    pub fn update(&mut self, key: &Bytes32, data: &[u8]) {
        let _ = self.tree.update(key, data);
    }

    pub fn delete(&mut self, key: &Bytes32) {
        let _ = self.tree.delete(key);
    }

    pub fn root(&self) -> Bytes32 {
        self.tree.root()
    }
}

impl Default for MerkleTree {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use sparse::hash::sum;

    #[test]
    fn test_empty_root() {
        let tree = MerkleTree::new();
        let root = tree.root();
        let expected_root = "0000000000000000000000000000000000000000000000000000000000000000";
        assert_eq!(hex::encode(root), expected_root);
    }

    #[test]
    fn test_update_1() {
        let mut tree = MerkleTree::new();

        tree.update(&sum(b"\x00\x00\x00\x00"), b"DATA");

        let root = tree.root();
        let expected_root = "39f36a7cb4dfb1b46f03d044265df6a491dffc1034121bc1071a34ddce9bb14b";
        assert_eq!(hex::encode(root), expected_root);
    }

    #[test]
    fn test_update_2() {
        let mut tree = MerkleTree::new();

        tree.update(&sum(b"\x00\x00\x00\x00"), b"DATA");
        tree.update(&sum(b"\x00\x00\x00\x01"), b"DATA");

        let root = tree.root();
        let expected_root = "8d0ae412ca9ca0afcb3217af8bcd5a673e798bd6fd1dfacad17711e883f494cb";
        assert_eq!(hex::encode(root), expected_root);
    }

    #[test]
    fn test_update_3() {
        let mut tree = MerkleTree::new();

        tree.update(&sum(b"\x00\x00\x00\x00"), b"DATA");
        tree.update(&sum(b"\x00\x00\x00\x01"), b"DATA");
        tree.update(&sum(b"\x00\x00\x00\x02"), b"DATA");

        let root = tree.root();
        let expected_root = "52295e42d8de2505fdc0cc825ff9fead419cbcf540d8b30c7c4b9c9b94c268b7";
        assert_eq!(hex::encode(root), expected_root);
    }

    #[test]
    fn test_update_1_delete_1() {
        let mut tree = MerkleTree::new();

        tree.update(&sum(b"\x00\x00\x00\x00"), b"DATA");
        tree.delete(&sum(b"\x00\x00\x00\x00"));

        let root = tree.root();
        let expected_root = "0000000000000000000000000000000000000000000000000000000000000000";
        assert_eq!(hex::encode(root), expected_root);
    }

    #[test]
    fn test_update_2_delete_1() {
        let mut tree = MerkleTree::new();

        tree.update(&sum(b"\x00\x00\x00\x00"), b"DATA");
        tree.update(&sum(b"\x00\x00\x00\x01"), b"DATA");
        tree.delete(&sum(b"\x00\x00\x00\x01"));

        let root = tree.root();
        let expected_root = "39f36a7cb4dfb1b46f03d044265df6a491dffc1034121bc1071a34ddce9bb14b";
        assert_eq!(hex::encode(root), expected_root);
    }
}