1use crate::hash::rpo::RpoDigest;
  2
  3use super::{
  4    super::{int_to_digest, int_to_node, NodeIndex},
  5    PartialMerkleTree, Rpo256,
  6};
  7
  8// TEST DATA
  9// ================================================================================================
 10
 11const NODE22: NodeIndex = NodeIndex::new_unchecked(2, 2);
 12
 13const NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2);
 14const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3);
 15
 16// TESTS
 17// ================================================================================================
 18
 19// with_paths CONSTRUCTOR TESTS
 20// ------------------------------------------------------------------------------------------------
 21
 22#[test]
 23fn get_root() {
 24    let leaf0 = int_to_digest(0);
 25    let leaf1 = int_to_digest(1);
 26    let leaf2 = int_to_digest(2);
 27    let leaf3 = int_to_digest(3);
 28
 29    let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
 30    let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
 31
 32    let root_exp = calculate_parent_hash(parent0, 0, parent1);
 33
 34    let set =
 35        super::PartialMerkleTree::with_paths([(0, leaf0, vec![*leaf1, *parent1].into())]).unwrap();
 36
 37    assert_eq!(set.root(), root_exp);
 38}
 39
 40#[test]
 41fn add_and_get_paths() {
 42    let value32 = int_to_digest(32);
 43    let value33 = int_to_digest(33);
 44    let value20 = int_to_digest(20);
 45    let value22 = int_to_digest(22);
 46    let value23 = int_to_digest(23);
 47
 48    let value21 = Rpo256::merge(&[value32, value33]);
 49    let value10 = Rpo256::merge(&[value20, value21]);
 50    let value11 = Rpo256::merge(&[value22, value23]);
 51
 52    let path_33 = vec![*value32, *value20, *value11];
 53
 54    let path_22 = vec![*value23, *value10];
 55
 56    let pmt = PartialMerkleTree::with_paths([
 57        (3, value33, path_33.clone().into()),
 58        (2, value22, path_22.clone().into()),
 59    ])
 60    .unwrap();
 61    let stored_path_33 = pmt.get_path(NODE33).unwrap();
 62    let stored_path_22 = pmt.get_path(NODE22).unwrap();
 63
 64    assert_eq!(path_33, *stored_path_33);
 65    assert_eq!(path_22, *stored_path_22);
 66}
 67
 68#[test]
 69fn get_node() {
 70    let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
 71    let hash_6 = int_to_digest(6);
 72    let index = NodeIndex::make(3, 6);
 73    let pmt = PartialMerkleTree::with_paths([(index.value(), hash_6, path_6.into())]).unwrap();
 74
 75    assert_eq!(int_to_digest(6u64), pmt.get_node(index).unwrap());
 76}
 77
 78#[test]
 79fn update_leaf() {
 80    let value32 = int_to_digest(32);
 81    let value33 = int_to_digest(33);
 82    let value20 = int_to_digest(20);
 83    let value22 = int_to_digest(22);
 84    let value23 = int_to_digest(23);
 85
 86    let value21 = Rpo256::merge(&[value32, value33]);
 87    let value10 = Rpo256::merge(&[value20, value21]);
 88    let value11 = Rpo256::merge(&[value22, value23]);
 89
 90    let path_33 = vec![*value32, *value20, *value11];
 91
 92    let path_22 = vec![*value23, *value10];
 93
 94    let mut pmt =
 95        PartialMerkleTree::with_paths([(3, value33, path_33.into()), (2, value22, path_22.into())])
 96            .unwrap();
 97
 98    let new_value32 = int_to_digest(132);
 99    let new_value21 = Rpo256::merge(&[new_value32, value33]);
100    let new_value10 = Rpo256::merge(&[value20, new_value21]);
101    let expected_root = Rpo256::merge(&[new_value10, value11]);
102
103    let old_leaf = pmt.update_leaf(NODE32, new_value32).unwrap();
104
105    assert_eq!(value32, old_leaf);
106
107    let new_root = pmt.root();
108
109    assert_eq!(new_root, expected_root);
110}
111
112#[test]
113fn check_leaf_depth() {
114    let value32 = int_to_digest(32);
115    let value33 = int_to_digest(33);
116    let value20 = int_to_digest(20);
117    let value22 = int_to_digest(22);
118    let value23 = int_to_digest(23);
119
120    let value11 = Rpo256::merge(&[value22, value23]);
121
122    let path_33 = vec![*value32, *value20, *value11];
123
124    let pmt = PartialMerkleTree::with_paths([(3, value33, path_33.into())]).unwrap();
125
126    assert_eq!(pmt.get_leaf_depth(0).unwrap(), 2);
127    assert_eq!(pmt.get_leaf_depth(1).unwrap(), 2);
128    assert_eq!(pmt.get_leaf_depth(2).unwrap(), 3);
129    assert_eq!(pmt.get_leaf_depth(3).unwrap(), 3);
130    assert_eq!(pmt.get_leaf_depth(4).unwrap(), 1);
131    assert_eq!(pmt.get_leaf_depth(5).unwrap(), 1);
132    assert_eq!(pmt.get_leaf_depth(6).unwrap(), 1);
133    assert_eq!(pmt.get_leaf_depth(7).unwrap(), 1);
134}
135
136// TODO: add test for add_path function and check correctness of leaf determination (requires
137// inner_nodes iter)
138
139// HELPER FUNCTIONS
140// --------------------------------------------------------------------------------------------
141
142/// Calculates the hash of the parent node by two sibling ones
143/// - node — current node
144/// - node_pos — position of the current node
145/// - sibling — neighboring vertex in the tree
146fn calculate_parent_hash(node: RpoDigest, node_pos: u64, sibling: RpoDigest) -> RpoDigest {
147    let parity = node_pos & 1;
148    if parity == 0 {
149        Rpo256::merge(&[node, sibling])
150    } else {
151        Rpo256::merge(&[sibling, node])
152    }
153}