merkletree-rs/src/lib.rs
src/lib.rsRaw
1use std::collections::HashMap;
2
3#[macro_use]
4extern crate arrayref;
5
6extern crate tiny_keccak;
7
8extern crate rustc_hex;
9
10mod utils;
11mod node;
12
13const TYPENODEEMPTY: u8 = 0;
14const TYPENODENORMAL: u8 = 1;
15const TYPENODEFINAL: u8 = 2;
16const TYPENODEVALUE: u8 = 3;
17// const TYPENODEROOT: u8 = 4;
18const EMPTYNODEVALUE: [u8;32] = [0;32];
19
20pub struct TestValue {
21 bytes: Vec<u8>,
22 index_length: u32,
23}
24pub trait Value {
25 fn bytes(&self) -> &Vec<u8>;
26 fn index_length(&self) -> u32;
27}
28impl Value for TestValue {
29 fn bytes(&self) -> &Vec<u8> {
30 &self.bytes
31 }
32 fn index_length(&self) -> u32 {
33 self.index_length
34 }
35}
36
37#[allow(dead_code)]
38pub struct Db {
39 storage: HashMap<[u8;32], Vec<u8>>,
40}
41impl Db {
42 pub fn insert(&mut self, k: [u8; 32], t: u8, il: u32, b: &mut Vec<u8>) {
43 let mut v: Vec<u8>;
44 v = [t].to_vec();
45 let il_bytes = il.to_le_bytes();
46 v.append(&mut il_bytes.to_vec()); // il_bytes are [u8;4] (4 bytes)
47 v.append(b);
48 self.storage.insert(k, v);
49 }
50 pub fn get(&self, k: &[u8;32]) -> (u8, u32, Vec<u8>) {
51 if k.to_vec() == EMPTYNODEVALUE.to_vec() {
52 return (0, 0, EMPTYNODEVALUE.to_vec());
53 }
54 match self.storage.get(k) {
55 Some(x) => {
56 let t = x[0];
57 let il_bytes: [u8; 4] = [x[1], x[2], x[3], x[4]];
58 let il = u32::from_le_bytes(il_bytes);
59 let b = &x[5..];
60 return (t, il, b.to_vec());
61 },
62 None => return (TYPENODEEMPTY, 0, EMPTYNODEVALUE.to_vec()),
63 }
64 }
65}
66
67pub fn new_db()-> Db {
68 let db = Db {
69 storage: HashMap::new(),
70 };
71 db
72}
73
74pub struct MerkleTree {
75 #[allow(dead_code)]
76 root: [u8; 32],
77 #[allow(dead_code)]
78 num_levels: u32,
79 #[allow(dead_code)]
80 sto: Db,
81}
82
83
84pub fn new(num_levels: u32) -> MerkleTree {
85 let mt = MerkleTree {
86 root: EMPTYNODEVALUE,
87 num_levels: num_levels,
88 sto: new_db(),
89 };
90 mt
91}
92
93impl MerkleTree {
94 pub fn add(&mut self, v: &TestValue) {
95 #![allow(unused_variables)]
96 #[allow(dead_code)]
97
98 // println!("adding value: {:?}", v.bytes());
99 // add the leaf that we are adding
100 self.sto.insert(utils::hash_vec(v.bytes().to_vec()), TYPENODEVALUE, v.index_length(), &mut v.bytes().to_vec());
101
102 let index = v.index_length() as usize;
103 let hi = utils::hash_vec(v.bytes()[..index].to_vec());
104 let ht = utils::hash_vec(v.bytes().to_vec());
105 let path = utils::get_path(self.num_levels, hi);
106 let mut siblings: Vec<[u8;32]> = Vec::new();
107
108
109 let mut node_hash = self.root;
110
111 for i in (0..self.num_levels-1).rev() {
112 // get node
113 // let (t, il, node_bytes) = self.sto.get(&utils::hash_vec(node_hash.to_vec()));
114 let (t, il, node_bytes) = self.sto.get(&node_hash);
115 if t == TYPENODEFINAL {
116 let hi_child = utils::hash_vec(v.bytes().to_vec().split_off(il as usize));
117 let path_child = utils::get_path(self.num_levels, hi_child);
118 let pos_diff = utils::compare_paths(path_child.clone(), path.clone());
119 if pos_diff == 999 { // TODO use a match here, and instead of 999 return something better
120 println!("compare paths err");
121 return;
122 }
123 let final_node_1_hash = utils::calc_hash_from_leaf_and_level(pos_diff, path_child.clone(), utils::hash_vec(node_bytes.to_vec()));
124 self.sto.insert(final_node_1_hash, TYPENODEFINAL, il, &mut node_bytes.to_vec());
125 let final_node_2_hash = utils::calc_hash_from_leaf_and_level(pos_diff, path.clone(), utils::hash_vec(v.bytes().to_vec()));
126 self.sto.insert(final_node_2_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
127
128 // parent node
129 let parent_node: node::TreeNode;
130 if path[pos_diff as usize] {
131 parent_node = node::TreeNode {
132 child_l: final_node_1_hash,
133 child_r: final_node_2_hash,
134 };
135 } else {
136 parent_node = node::TreeNode {
137 child_l: final_node_2_hash,
138 child_r: final_node_1_hash,
139 };
140 }
141 let empties = utils::get_empties_between_i_and_pos(i, pos_diff+1);
142 for j in 0..empties.len() {
143 siblings.push(empties[j]);
144 }
145 let path_from_pos_diff = utils::cut_path(path.clone(), (pos_diff +1) as usize);
146 self.root = self.replace_leaf(path_from_pos_diff, siblings.clone(), parent_node.ht(), TYPENODENORMAL, 0, &mut parent_node.bytes().to_vec());
147 return;
148 }
149
150 let node = node::parse_node_bytes(node_bytes);
151
152 let sibling: [u8;32];
153 if path.clone().into_iter().nth(i as usize).unwrap() {
154 node_hash = node.child_l;
155 sibling = node.child_r;
156 } else {
157 sibling = node.child_l;
158 node_hash = node.child_r;
159 }
160 siblings.push(*array_ref!(sibling, 0, 32));
161 if node_hash == EMPTYNODEVALUE {
162 if i==self.num_levels-2 && siblings[siblings.len()-1]==EMPTYNODEVALUE {
163
164 let final_node_hash = utils::calc_hash_from_leaf_and_level(i+1, path.clone(), utils::hash_vec(v.bytes().to_vec()));
165 self.sto.insert(final_node_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
166 self.root = final_node_hash;
167 return;
168 }
169 let final_node_hash = utils::calc_hash_from_leaf_and_level(i, path.clone(), utils::hash_vec(v.bytes().to_vec()));
170 let path_from_i = utils::cut_path(path.clone(), i as usize);
171 self.root = self.replace_leaf(path_from_i, siblings.clone(), final_node_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
172 return;
173 }
174 }
175 self.root = self.replace_leaf(path, siblings, utils::hash_vec(v.bytes().to_vec()), TYPENODEVALUE, v.index_length(), &mut v.bytes().to_vec());
176 }
177
178 #[allow(dead_code)]
179 pub fn replace_leaf(&mut self, path: Vec<bool>, siblings: Vec<[u8;32]>, leaf_hash: [u8;32], node_type: u8, index_length: u32, leaf_value: &mut Vec<u8>) -> [u8;32] {
180 self.sto.insert(leaf_hash, node_type, index_length, leaf_value);
181 let mut curr_node = leaf_hash;
182
183 for i in 0..siblings.len() {
184 if path.clone().into_iter().nth(i as usize).unwrap() {
185 let node = node::TreeNode {
186 child_l: curr_node,
187 child_r: siblings[siblings.len()-1-i],
188 };
189 self.sto.insert(node.ht(), TYPENODENORMAL, 0, &mut node.bytes());
190 curr_node = node.ht();
191 } else {
192 let node = node::TreeNode {
193 child_l: siblings[siblings.len()-1-i],
194 child_r: curr_node,
195 };
196 self.sto.insert(node.ht(), TYPENODENORMAL, 0, &mut node.bytes());
197 curr_node = node.ht();
198 }
199 }
200 curr_node
201 }
202 }
203
204
205#[cfg(test)]
206 mod tests {
207 use super::*;
208 use rustc_hex::ToHex;
209
210 #[test]
211 fn test_hash_vec() {
212 let a: Vec<u8> = From::from("test");
213 let h = utils::hash_vec(a);
214 assert_eq!("9c22ff5f21f0b81b113e63f7db6da94fedef11b2119b4088b89664fb9a3cb658", h.to_hex());
215 }
216
217 #[test]
218 fn test_new_mt() {
219 let mt: MerkleTree = new(140);
220 assert_eq!(140, mt.num_levels);
221 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
222 let (_t, _il, b) = mt.sto.get(&[0;32]);
223 assert_eq!(mt.root.to_vec(), b);
224 }
225
226
227 #[test]
228 fn test_tree_node() {
229 let n = node::TreeNode {
230 child_l: [1;32],
231 child_r: [2;32],
232 };
233 assert_eq!("01010101010101010101010101010101010101010101010101010101010101010202020202020202020202020202020202020202020202020202020202020202",
234 n.bytes().to_hex());
235 assert_eq!("346d8c96a2454213fcc0daff3c96ad0398148181b9fa6488f7ae2c0af5b20aa0", n.ht().to_hex());
236 }
237
238 #[test]
239 fn test_add() {
240 let mut mt: MerkleTree = new(140);
241 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
242 let val = TestValue {
243 bytes: vec![1,2,3,4,5],
244 index_length: 3,
245 };
246 mt.add(&val);
247 let (_t, _il, b) = mt.sto.get(&utils::hash_vec(val.bytes().to_vec()));
248 assert_eq!(*val.bytes(), b);
249 assert_eq!("a0e72cc948119fcb71b413cf5ada12b2b825d5133299b20a6d9325ffc3e2fbf1", mt.root.to_hex());
250 }
251 #[test]
252 fn test_add_2() {
253 let mut mt: MerkleTree = new(140);
254 let val = TestValue {
255 bytes: "this is a test leaf".as_bytes().to_vec(),
256 index_length: 15,
257 };
258 assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
259 mt.add(&val);
260 let (_t, _il, b) = mt.sto.get(&utils::hash_vec(val.bytes().to_vec()));
261 assert_eq!(*val.bytes(), b);
262 assert_eq!("b4fdf8a653198f0e179ccb3af7e4fc09d76247f479d6cfc95cd92d6fda589f27", mt.root.to_hex());
263 }
264 }