1use crate::hash::rpo::RpoDigest;
  2
  3use super::{
  4    super::{int_to_node, NodeIndex},
  5    InnerNodeInfo, MerkleError, PartialMerkleTree, Rpo256, Vec, Word,
  6};
  7
  8// TEST DATA
  9// ================================================================================================
 10
 11const ROOT_NODE: NodeIndex = NodeIndex::new_unchecked(0, 0);
 12
 13const NODE10: NodeIndex = NodeIndex::new_unchecked(1, 0);
 14const NODE11: NodeIndex = NodeIndex::new_unchecked(1, 1);
 15
 16const NODE20: NodeIndex = NodeIndex::new_unchecked(2, 0);
 17const NODE21: NodeIndex = NodeIndex::new_unchecked(2, 1);
 18const NODE22: NodeIndex = NodeIndex::new_unchecked(2, 2);
 19const NODE23: NodeIndex = NodeIndex::new_unchecked(2, 3);
 20
 21const NODE32: NodeIndex = NodeIndex::new_unchecked(3, 2);
 22const NODE33: NodeIndex = NodeIndex::new_unchecked(3, 3);
 23
 24// TESTS
 25// ================================================================================================
 26
 27// with_paths CONSTRUCTOR TESTS
 28// ------------------------------------------------------------------------------------------------
 29
 30#[test]
 31fn get_root() {
 32    let leaf0 = int_to_node(0);
 33    let leaf1 = int_to_node(1);
 34    let leaf2 = int_to_node(2);
 35    let leaf3 = int_to_node(3);
 36
 37    let parent0 = calculate_parent_hash(leaf0, 0, leaf1);
 38    let parent1 = calculate_parent_hash(leaf2, 2, leaf3);
 39
 40    let root_exp = calculate_parent_hash(parent0, 0, parent1);
 41
 42    let set = super::PartialMerkleTree::with_paths([(NODE20, leaf0, vec![leaf1, parent1].into())])
 43        .unwrap();
 44
 45    assert_eq!(set.root(), root_exp);
 46}
 47
 48#[test]
 49fn add_and_get_paths() {
 50    let value32 = int_to_node(32).into();
 51    let value33 = int_to_node(33).into();
 52    let value20 = int_to_node(20).into();
 53    let value22 = int_to_node(22).into();
 54    let value23 = int_to_node(23).into();
 55
 56    let value21 = Rpo256::merge(&[value32, value33]);
 57    let value10 = Rpo256::merge(&[value20, value21]);
 58    let value11 = Rpo256::merge(&[value22, value23]);
 59
 60    let path_33 = vec![*value32, *value20, *value11];
 61
 62    let path_22 = vec![*value23, *value10];
 63
 64    let pmt = PartialMerkleTree::with_paths([
 65        (NODE33, *value33, path_33.clone().into()),
 66        (NODE22, *value22, path_22.clone().into()),
 67    ])
 68    .unwrap();
 69    let stored_path_33 = pmt.get_path(NODE33).unwrap();
 70    let stored_path_22 = pmt.get_path(NODE22).unwrap();
 71
 72    assert_eq!(path_33, *stored_path_33);
 73    assert_eq!(path_22, *stored_path_22);
 74}
 75
 76#[test]
 77fn get_node() {
 78    let path_6 = vec![int_to_node(7), int_to_node(45), int_to_node(123)];
 79    let hash_6 = int_to_node(6);
 80    let index = NodeIndex::make(3, 6);
 81    let pmt = PartialMerkleTree::with_paths([(index, hash_6, path_6.into())]).unwrap();
 82
 83    assert_eq!(int_to_node(6u64), *pmt.get_node(index).unwrap());
 84}
 85
 86#[test]
 87fn update_leaf() {
 88    let value32 = int_to_node(32).into();
 89    let value33 = int_to_node(33).into();
 90    let value20 = int_to_node(20).into();
 91    let value22 = int_to_node(22).into();
 92    let value23 = int_to_node(23).into();
 93
 94    let value21 = Rpo256::merge(&[value32, value33]);
 95    let value10 = Rpo256::merge(&[value20, value21]);
 96    let value11 = Rpo256::merge(&[value22, value23]);
 97
 98    let path_33 = vec![*value32, *value20, *value11];
 99
100    let path_22 = vec![*value23, *value10];
101
102    let mut pmt = PartialMerkleTree::with_paths([
103        (NODE33, *value33, path_33.into()),
104        (NODE22, *value22, path_22.into()),
105    ])
106    .unwrap();
107
108    let new_value32 = int_to_node(132).into();
109    let new_value21 = Rpo256::merge(&[new_value32, value33]);
110    let new_value10 = Rpo256::merge(&[value20, new_value21]);
111    let expected_root = Rpo256::merge(&[new_value10, value11]);
112
113    let old_leaf = pmt.update_leaf(NODE32, new_value32).unwrap();
114
115    assert_eq!(value32, old_leaf);
116
117    let new_root = pmt.root();
118
119    assert_eq!(new_root, *expected_root);
120}
121
122#[test]
123fn test_inner_node_iterator() -> Result<(), MerkleError> {
124    let value32 = int_to_node(32).into();
125    let value33 = int_to_node(33).into();
126    let value20 = int_to_node(20).into();
127    let value22 = int_to_node(22).into();
128    let value23 = int_to_node(23).into();
129
130    let value21 = Rpo256::merge(&[value32, value33]);
131    let value10 = Rpo256::merge(&[value20, value21]);
132    let value11 = Rpo256::merge(&[value22, value23]);
133    let root = Rpo256::merge(&[value10, value11]);
134
135    let path_33 = vec![*value32, *value20, *value11];
136
137    let path_22 = vec![*value23, *value10];
138
139    let pmt = PartialMerkleTree::with_paths([
140        (NODE33, *value33, path_33.into()),
141        (NODE22, *value22, path_22.into()),
142    ])
143    .unwrap();
144
145    assert_eq!(root, pmt.get_node(ROOT_NODE).unwrap());
146    assert_eq!(value10, pmt.get_node(NODE10).unwrap());
147    assert_eq!(value11, pmt.get_node(NODE11).unwrap());
148    assert_eq!(value20, pmt.get_node(NODE20).unwrap());
149    assert_eq!(value21, pmt.get_node(NODE21).unwrap());
150    assert_eq!(value22, pmt.get_node(NODE22).unwrap());
151    assert_eq!(value23, pmt.get_node(NODE23).unwrap());
152    assert_eq!(value32, pmt.get_node(NODE32).unwrap());
153    assert_eq!(value33, pmt.get_node(NODE33).unwrap());
154
155    let nodes: Vec<InnerNodeInfo> = pmt.inner_nodes().collect();
156    let expected = vec![
157        InnerNodeInfo {
158            value: *root,
159            left: *value10,
160            right: *value11,
161        },
162        InnerNodeInfo {
163            value: *value10,
164            left: *value20,
165            right: *value21,
166        },
167        InnerNodeInfo {
168            value: *value11,
169            left: *value22,
170            right: *value23,
171        },
172        InnerNodeInfo {
173            value: *value21,
174            left: *value32,
175            right: *value33,
176        },
177    ];
178    assert_eq!(nodes, expected);
179
180    Ok(())
181}
182
183#[test]
184fn check_leaf_depth() {
185    let value32: RpoDigest = int_to_node(32).into();
186    let value33: RpoDigest = int_to_node(33).into();
187    let value20: RpoDigest = int_to_node(20).into();
188    let value22 = int_to_node(22).into();
189    let value23 = int_to_node(23).into();
190
191    let value11 = Rpo256::merge(&[value22, value23]);
192
193    let path_33 = vec![*value32, *value20, *value11];
194
195    let pmt = PartialMerkleTree::with_paths([(NODE33, *value33, path_33.into())]).unwrap();
196
197    assert_eq!(pmt.get_leaf_depth(0).unwrap(), 2);
198    assert_eq!(pmt.get_leaf_depth(1).unwrap(), 2);
199    assert_eq!(pmt.get_leaf_depth(2).unwrap(), 3);
200    assert_eq!(pmt.get_leaf_depth(3).unwrap(), 3);
201    assert_eq!(pmt.get_leaf_depth(4).unwrap(), 1);
202    assert_eq!(pmt.get_leaf_depth(5).unwrap(), 1);
203    assert_eq!(pmt.get_leaf_depth(6).unwrap(), 1);
204    assert_eq!(pmt.get_leaf_depth(7).unwrap(), 1);
205}
206
207// HELPER FUNCTIONS
208// --------------------------------------------------------------------------------------------
209
210/// Calculates the hash of the parent node by two sibling ones
211/// - node — current node
212/// - node_pos — position of the current node
213/// - sibling — neighboring vertex in the tree
214fn calculate_parent_hash(node: Word, node_pos: u64, sibling: Word) -> Word {
215    let parity = node_pos & 1;
216    if parity == 0 {
217        Rpo256::merge(&[node.into(), sibling.into()]).into()
218    } else {
219        Rpo256::merge(&[sibling.into(), node.into()]).into()
220    }
221}