luct_core/
tree.rs

1use crate::store::{AppendableStore, Hashable, Store};
2pub use crate::tree::{
3    consistency::ConsistencyProof,
4    inclusion::AuditProof,
5    node::{Node, NodeKey},
6};
7use serde::{Deserialize, Serialize};
8use thiserror::Error;
9
10mod consistency;
11mod inclusion;
12mod node;
13
14pub(crate) type HashOutput = [u8; 32];
15
16#[derive(Clone, Debug, PartialEq, Eq, Error)]
17pub enum ProofGenerationError {
18    #[error("Index {index} not found in tree of size {tree_size}")]
19    InvalidIndex { tree_size: u64, index: u64 },
20
21    #[error("Invalid tree size {received} smaller than {expected}")]
22    InvalidTreeSize { expected: u64, received: u64 },
23
24    #[error("Failed to fetch key {0:?} from the store")]
25    KeyNotFound(NodeKey),
26}
27
28#[derive(Clone, Debug, PartialEq, Eq, Error)]
29pub enum ProofValidationError {
30    #[error("Found an unxexpected hash length (expected: {expected}, received: {received})")]
31    InvalidHashLength { expected: usize, received: usize },
32
33    #[error("Index {index} not found in tree of size {tree_size}")]
34    InvalidIndex { tree_size: u64, index: u64 },
35
36    #[error("Invalid tree size {received} smaller than {expected}")]
37    InvalidTreeSize { expected: u64, received: u64 },
38
39    #[error("Hash mismatch")]
40    HashMismatch,
41
42    #[error("Merkle path was too short")]
43    PathTooShort,
44
45    #[error("Merkle path was too long")]
46    PathTooLong,
47}
48
49#[derive(Debug, Clone)]
50
51pub struct Tree<N, L> {
52    nodes: N,
53    leafs: L,
54}
55
56impl<N, L> Tree<N, L> {
57    pub fn new(node_store: N, leaf_store: L) -> Self {
58        Self {
59            nodes: node_store,
60            leafs: leaf_store,
61        }
62    }
63
64    pub fn nodes(&self) -> &N {
65        &self.nodes
66    }
67}
68
69impl<N, L> Tree<N, L>
70where
71    N: Store<Key = NodeKey, Value = HashOutput>,
72    L: AppendableStore<Key = u64, Value: Hashable>,
73{
74    pub fn insert_entry(&self, entry: L::Value) {
75        let entry_hash = entry.hash();
76        let idx = self.leafs.append(entry);
77        let entry_key = NodeKey::leaf(idx);
78        self.nodes.insert(entry_key, entry_hash);
79
80        // Already update intermediate nodes, if they are power of twos
81        let end = idx + 1;
82        let mut diff = 2;
83
84        while end.is_multiple_of(diff) {
85            let start = end - diff;
86
87            let key = NodeKey { start, end };
88            let (left, right) = key.split();
89
90            let node = Node {
91                left: self.nodes.get(&left).unwrap(),
92                right: self.nodes.get(&right).unwrap(),
93            };
94
95            self.nodes.insert(key, node.hash());
96
97            diff <<= 1;
98        }
99    }
100
101    pub fn recompute_tree_head(&self) -> TreeHead {
102        let tree_size = self.leafs.len() as u64;
103        let mut current_key = NodeKey::full_range(tree_size);
104        let mut balanced_nodes = vec![];
105
106        while !current_key.is_balanced() {
107            let (left, right) = current_key.split();
108            assert!(left.is_balanced());
109            balanced_nodes.push(left);
110            current_key = right;
111        }
112
113        let mut current_node_hash = self.nodes.get(&current_key).unwrap();
114        while let Some(left_key) = balanced_nodes.pop() {
115            let current_node = Node {
116                left: self.nodes.get(&left_key).unwrap(),
117                right: self.nodes.get(&current_key).unwrap(),
118            };
119
120            current_key = left_key.merge(&current_key).unwrap();
121            current_node_hash = current_node.hash();
122            self.nodes.insert(current_key.clone(), current_node_hash);
123        }
124
125        TreeHead {
126            tree_size,
127            head: current_node_hash,
128        }
129    }
130
131    pub fn get_latest_tree_head(&self) -> Option<TreeHead> {
132        let idx = self.leafs.len() as u64;
133        self.nodes
134            .get(&NodeKey::full_range(idx))
135            .map(|head| TreeHead {
136                tree_size: idx,
137                head,
138            })
139    }
140}
141
142#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
143pub struct TreeHead {
144    pub(crate) tree_size: u64,
145    pub(crate) head: HashOutput,
146}
147
148impl TreeHead {
149    pub fn tree_size(&self) -> u64 {
150        self.tree_size
151    }
152}
153
154#[cfg(test)]
155mod tests {
156    use super::*;
157    use sha2::{Digest, Sha256};
158
159    impl Hashable for String {
160        fn hash(&self) -> HashOutput {
161            Sha256::digest(self.as_bytes()).into()
162        }
163    }
164
165    impl Hashable for HashOutput {
166        fn hash(&self) -> HashOutput {
167            Sha256::digest(self).into()
168        }
169    }
170}
luct_core/tree.rs

luct_core/
tree.rs
