implemented merkletree add() method, added few tests

src/lib.rs

@@ -89,3 +89,176 @@ pub fn new(num_levels: u32) -> MerkleTree {
    };
    mt
 }
+
+impl MerkleTree {
+   pub fn add(&mut self, v: &TestValue) {
+      #![allow(unused_variables)]
+      #[allow(dead_code)]
+
+      // println!("adding value: {:?}", v.bytes());
+      // add the leaf that we are adding
+      self.sto.insert(utils::hash_vec(v.bytes().to_vec()), TYPENODEVALUE, v.index_length(), &mut v.bytes().to_vec());
+
+      let index = v.index_length() as usize;
+      let hi = utils::hash_vec(v.bytes()[..index].to_vec());
+      let ht = utils::hash_vec(v.bytes().to_vec());
+      let path = utils::get_path(self.num_levels, hi);
+      let mut siblings: Vec<[u8;32]> = Vec::new();
+
+
+      let mut node_hash = self.root;
+
+      for i in (0..self.num_levels-1).rev() {
+         // get node
+         // let (t, il, node_bytes) = self.sto.get(&utils::hash_vec(node_hash.to_vec()));
+         let (t, il, node_bytes) = self.sto.get(&node_hash);
+         if t == TYPENODEFINAL {
+            let hi_child = utils::hash_vec(v.bytes().to_vec().split_off(il as usize));
+            let path_child = utils::get_path(self.num_levels, hi_child);
+            let pos_diff = utils::compare_paths(path_child.clone(), path.clone());
+            if pos_diff == 999 { // TODO use a match here, and instead of 999 return something better
+               println!("compare paths err");
+               return;
+            }
+            let final_node_1_hash = utils::calc_hash_from_leaf_and_level(pos_diff, path_child.clone(), utils::hash_vec(node_bytes.to_vec()));
+            self.sto.insert(final_node_1_hash, TYPENODEFINAL, il, &mut node_bytes.to_vec());
+            let final_node_2_hash = utils::calc_hash_from_leaf_and_level(pos_diff, path.clone(), utils::hash_vec(v.bytes().to_vec()));
+            self.sto.insert(final_node_2_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
+
+            // parent node
+            let parent_node: node::TreeNode;
+            if path[pos_diff as usize] {
+               parent_node = node::TreeNode {
+                  child_l: final_node_1_hash,
+                  child_r: final_node_2_hash,
+               };
+            } else {
+               parent_node = node::TreeNode {
+                  child_l: final_node_2_hash,
+                  child_r: final_node_1_hash,
+               };
+            }
+            let empties = utils::get_empties_between_i_and_pos(i, pos_diff+1);
+            for j in 0..empties.len() {
+               siblings.push(empties[j]);
+            }
+            let path_from_pos_diff = utils::cut_path(path.clone(), (pos_diff +1) as usize);
+            self.root = self.replace_leaf(path_from_pos_diff, siblings.clone(), parent_node.ht(), TYPENODENORMAL, 0, &mut parent_node.bytes().to_vec());
+            return;
+         }
+
+         let node = node::parse_node_bytes(node_bytes);
+
+         let sibling: [u8;32];
+         if path.clone().into_iter().nth(i as usize).unwrap() {
+            node_hash = node.child_l;
+            sibling = node.child_r;
+         } else {
+            sibling = node.child_l;
+            node_hash = node.child_r;
+         }
+         siblings.push(*array_ref!(sibling, 0, 32));
+         if node_hash == EMPTYNODEVALUE {
+            if i==self.num_levels-2 && siblings[siblings.len()-1]==EMPTYNODEVALUE {
+
+               let final_node_hash = utils::calc_hash_from_leaf_and_level(i+1, path.clone(), utils::hash_vec(v.bytes().to_vec()));
+               self.sto.insert(final_node_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
+               self.root = final_node_hash;
+               return;
+            }
+            let final_node_hash = utils::calc_hash_from_leaf_and_level(i, path.clone(), utils::hash_vec(v.bytes().to_vec()));
+            let path_from_i = utils::cut_path(path.clone(), i as usize);
+            self.root = self.replace_leaf(path_from_i, siblings.clone(), final_node_hash, TYPENODEFINAL, v.index_length(), &mut v.bytes().to_vec());
+            return;
+         }
+      }
+      self.root = self.replace_leaf(path, siblings, utils::hash_vec(v.bytes().to_vec()), TYPENODEVALUE, v.index_length(), &mut v.bytes().to_vec());
+      }
+
+      #[allow(dead_code)]
+      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] {
+         self.sto.insert(leaf_hash, node_type, index_length, leaf_value);
+         let mut curr_node = leaf_hash;
+
+         for i in 0..siblings.len() {
+            if path.clone().into_iter().nth(i as usize).unwrap() {
+               let node = node::TreeNode {
+                  child_l: curr_node,
+                  child_r: siblings[siblings.len()-1-i],
+               };
+               self.sto.insert(node.ht(), TYPENODENORMAL, 0, &mut node.bytes());
+               curr_node = node.ht();
+            } else {
+               let node = node::TreeNode {
+                  child_l: siblings[siblings.len()-1-i],
+                  child_r: curr_node,
+               };
+               self.sto.insert(node.ht(), TYPENODENORMAL, 0, &mut node.bytes());
+               curr_node = node.ht();
+            }
+         }
+         curr_node
+      }
+   }
+
+
+#[cfg(test)]
+   mod tests {
+      use super::*;
+      use rustc_hex::ToHex;
+
+      #[test]
+      fn test_hash_vec() {
+         let a: Vec<u8> = From::from("test");
+         let h = utils::hash_vec(a);
+         assert_eq!("9c22ff5f21f0b81b113e63f7db6da94fedef11b2119b4088b89664fb9a3cb658", h.to_hex());
+      }
+
+      #[test]
+      fn test_new_mt() {
+         let mt: MerkleTree = new(140);
+         assert_eq!(140, mt.num_levels);
+         assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
+         let (_t, _il, b) = mt.sto.get(&[0;32]);
+         assert_eq!(mt.root.to_vec(), b);
+      }
+
+
+      #[test]
+      fn test_tree_node() {
+         let n = node::TreeNode {
+            child_l: [1;32],
+            child_r: [2;32],
+         };
+         assert_eq!("01010101010101010101010101010101010101010101010101010101010101010202020202020202020202020202020202020202020202020202020202020202",
+                    n.bytes().to_hex());
+         assert_eq!("346d8c96a2454213fcc0daff3c96ad0398148181b9fa6488f7ae2c0af5b20aa0", n.ht().to_hex());
+      }
+
+      #[test]
+      fn test_add() {
+         let mut mt: MerkleTree = new(140);
+         assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
+         let val = TestValue {
+            bytes: vec![1,2,3,4,5],
+            index_length: 3,
+         };
+         mt.add(&val);
+         let (_t, _il, b) = mt.sto.get(&utils::hash_vec(val.bytes().to_vec()));
+         assert_eq!(*val.bytes(), b);
+         assert_eq!("a0e72cc948119fcb71b413cf5ada12b2b825d5133299b20a6d9325ffc3e2fbf1", mt.root.to_hex());
+      }
+      #[test]
+      fn test_add_2() {
+          let mut mt: MerkleTree = new(140);
+          let val = TestValue {
+              bytes: "this is a test leaf".as_bytes().to_vec(),
+              index_length: 15,
+          };
+          assert_eq!("0000000000000000000000000000000000000000000000000000000000000000", mt.root.to_hex());
+          mt.add(&val);
+         let (_t, _il, b) = mt.sto.get(&utils::hash_vec(val.bytes().to_vec()));
+          assert_eq!(*val.bytes(), b);
+          assert_eq!("b4fdf8a653198f0e179ccb3af7e4fc09d76247f479d6cfc95cd92d6fda589f27", mt.root.to_hex());
+      }
+   }