diff --git a/nssa/src/merkle_tree.rs b/nssa/src/merkle_tree.rs index 917a836..132d537 100644 --- a/nssa/src/merkle_tree.rs +++ b/nssa/src/merkle_tree.rs @@ -34,14 +34,13 @@ impl MerkleTree { } fn root_index(&self) -> usize { - let total_levels = self.capacity.trailing_zeros() as usize; - let used_levels = self.size().trailing_zeros() as usize; - let diff = total_levels - used_levels; + let capacity_depth = self.capacity.trailing_zeros() as usize; + let diff = capacity_depth - self.depth(); if diff == 0 { 0 } else { (1 << diff) - 1 } } - fn size(&self) -> usize { - self.values.len().next_power_of_two() + fn depth(&self) -> usize { + self.values.len().next_power_of_two().trailing_zeros() as usize } fn get_node(&self, index: &usize) -> &Node { @@ -71,7 +70,7 @@ impl MerkleTree { } } - fn extend_capacity(&mut self) { + fn reallocate_to_double_capacity(&mut self) { let mut this = Self::with_capacity(self.capacity << 1); for value in self.values.iter() { this.insert(*value); @@ -85,7 +84,7 @@ impl MerkleTree { } if self.capacity == self.values.len() { - self.extend_capacity(); + self.reallocate_to_double_capacity(); } let new_index = self.values.len(); @@ -98,7 +97,7 @@ impl MerkleTree { self.node_map.insert(layer_index, layer_node); let mut layer = 0; - let mut top_layer = self.size().trailing_zeros(); + let mut top_layer = self.depth(); while layer < top_layer { let is_left_child = layer_index & 1 == 1; @@ -139,6 +138,28 @@ impl MerkleTree { } this } + + pub fn get_authentication_path_for(&self, value: &Value) -> Option<(usize, Vec)> { + let mut result = Vec::with_capacity(self.depth()); + let value_index = self.index_map.get(value)?; + let base_length = self.capacity; + let mut layer_index = base_length + value_index - 1; + let mut layer = 0; + let top_layer = self.depth(); + while layer < top_layer { + let is_left_child = layer_index & 1 == 1; + let (sibling, parent_index) = if is_left_child { + (self.get_node(&(layer_index + 1)), (layer_index - 1) >> 1) + } else { + (self.get_node(&(layer_index - 1)), (layer_index - 2) >> 1) + }; + result.push(*sibling); + + layer += 1; + layer_index = parent_index; + } + Some((*value_index, result)) + } } #[cfg(test)] @@ -245,6 +266,18 @@ mod tests { assert_eq!(tree.capacity, 8); } + #[test] + fn test_merkle_tree_6() { + let values = vec![[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; + let tree = MerkleTree::new(values); + let expected_root = [ + 6, 156, 184, 37, 154, 6, 254, 110, 219, 63, 167, 255, 121, 51, 166, 221, 125, 202, 111, + 202, 41, 147, 20, 55, 151, 148, 166, 136, 146, 108, 55, 146, + ]; + + assert_eq!(tree.root(), expected_root); + } + #[test] fn test_with_capacity_4() { let tree = MerkleTree::with_capacity(4); @@ -381,6 +414,76 @@ mod tests { assert_eq!(expected_tree, tree); } + + #[test] + fn test_authentication_path_1() { + let values = vec![[1; 32], [2; 32], [3; 32], [4; 32]]; + let tree = MerkleTree::new(values); + let expected_authentication_path = ( + 2, + vec![ + [ + 159, 79, 182, 143, 62, 29, 172, 130, 32, 47, 154, 165, 129, 206, 11, 191, 31, + 118, 93, 240, 233, 172, 60, 140, 87, 226, 15, 104, 90, 186, 184, 237, + ], + [ + 80, 162, 125, 71, 70, 243, 87, 203, 112, 12, 190, 157, 72, 131, 183, 127, 182, + 79, 1, 40, 130, 138, 52, 137, 220, 106, 111, 33, 221, 191, 36, 20, + ], + ], + ); + + let authentication_path = tree.get_authentication_path_for(&[3; 32]).unwrap(); + assert_eq!(authentication_path, expected_authentication_path); + } + + #[test] + fn test_authentication_path_2() { + let values = vec![[1; 32], [2; 32], [3; 32]]; + let tree = MerkleTree::new(values); + let expected_authentication_path = ( + 0, + vec![ + [ + 117, 135, 123, 180, 29, 57, 59, 95, 184, 69, 92, 230, 14, 205, 141, 218, 0, 29, + 6, 49, 100, 150, 177, 77, 250, 127, 137, 86, 86, 238, 202, 74, + ], + [ + 164, 27, 133, 93, 45, 180, 222, 144, 82, 205, 123, 229, 236, 103, 214, 88, 102, + 41, 203, 159, 110, 50, 70, 164, 175, 165, 186, 49, 63, 7, 169, 197, + ], + ], + ); + + let authentication_path = tree.get_authentication_path_for(&[1; 32]).unwrap(); + assert_eq!(authentication_path, expected_authentication_path); + } + + #[test] + fn test_authentication_path_3() { + let values = vec![[1; 32], [2; 32], [3; 32], [4; 32], [5; 32]]; + let tree = MerkleTree::new(values); + let expected_authentication_path = ( + 4, + vec![ + [ + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, + ], + [ + 245, 165, 253, 66, 209, 106, 32, 48, 39, 152, 239, 110, 211, 9, 151, 155, 67, + 0, 61, 35, 32, 217, 240, 232, 234, 152, 49, 169, 39, 89, 251, 75, + ], + [ + 72, 199, 63, 120, 33, 165, 138, 141, 42, 112, 62, 91, 57, 197, 113, 192, 170, + 32, 207, 20, 171, 205, 10, 248, 242, 185, 85, 188, 32, 41, 152, 222, + ], + ], + ); + + let authentication_path = tree.get_authentication_path_for(&[5; 32]).unwrap(); + assert_eq!(authentication_path, expected_authentication_path); + } } mod default_values {