diff --git a/src/tree.rs b/src/tree.rs index 6b31896..d7e2d68 100644 --- a/src/tree.rs +++ b/src/tree.rs @@ -205,13 +205,14 @@ where let root_key = Key(0, 0); subtree.insert(root_key, self.root); - self.fill_nodes(root_key, self.next_index, end, &mut subtree)?; - - ////// TODO! Refactor later - for i in self.next_index..end { - subtree.insert(Key(self.depth, i), leaves[i - self.next_index]); - } - ////// + self.fill_nodes( + root_key, + self.next_index, + end, + &mut subtree, + leaves, + self.next_index, + )?; let subtree = Arc::new(RwLock::new(subtree)); @@ -248,8 +249,11 @@ where start: usize, end: usize, subtree: &mut HashMap, + leaves: &[H::Fr], + from: usize, ) -> Result<()> { if key.0 == self.depth { + subtree.insert(key, leaves[key.1 - from]); return Ok(()); } @@ -265,11 +269,11 @@ where let half = 1 << (self.depth - key.0 - 1); if start < half { - self.fill_nodes(left, start, min(end, half), subtree)?; + self.fill_nodes(left, start, min(end, half), subtree, leaves, from)?; } if end > half { - self.fill_nodes(right, 0, end - half, subtree)?; + self.fill_nodes(right, 0, end - half, subtree, leaves, from)?; } Ok(()) diff --git a/tests/sled_keccak.rs b/tests/sled_keccak.rs index 8c2b253..9eae448 100644 --- a/tests/sled_keccak.rs +++ b/tests/sled_keccak.rs @@ -129,7 +129,7 @@ fn insert_delete() -> Result<()> { #[test] fn batch_insertions() -> Result<()> { - let mut mt = MerkleTree::::new(2, "abacabas")?; + let mut mt = MerkleTree::::new(2, "abacabasa")?; let leaves = [ hex!("0000000000000000000000000000000000000000000000000000000000000001"), @@ -145,7 +145,7 @@ fn batch_insertions() -> Result<()> { hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36") ); - fs::remove_dir_all("abacabas").expect("Error removing db"); + fs::remove_dir_all("abacabasa").expect("Error removing db"); Ok(()) }