tallytree/proof/
util.rs

1use crate::hash::{hash_leaf, hash_node_ref, NodeHash, NULL_HASH};
2use crate::node::{is_leaf_node, is_wrapper_node};
3use crate::proof::{Proof, ProofBranch, ProofHash, ProofNode};
4use crate::tally::{combine_tally, has_one_vote, TallyList};
5use crate::utilstring::to_hex;
6use crate::{NodeRef, Validation, VoteReference};
7use sha2::{Digest, Sha256};
8use std::collections::HashSet;
9
10/// The direction traversed when parsing a merkle proof branch
11#[derive(Clone, Debug, PartialEq, Copy)]
12pub enum TraverseDirection {
13    /// Merkle proof recurses to the left from current branch
14    Left = 0,
15    /// Merkle proof recurses to the right from current branch
16    Right = 1,
17}
18
19/// Set constraints on how branch is allowed to collapse the branches.
20pub enum CollapseBranchConstrains {
21    /// No constraints
22    None = 0,
23    /// Fail if the proof is not for the right-most branch of the tree.
24    /// The leaf can still hang on the left on its parent, but any parent
25    /// cannot be on the left of its parent.
26    /// (Don't allow collapsing to the left)
27    RejectLeftPath = 1,
28    /// Fail if the proof is not for the left-most branch of the tree.
29    /// The leaf can still hang on the right on its parent, but any parent
30    /// cannot be on the right of its parent.
31    /// (Don't allow collapsing to the right)
32    RejectRightPath = 2,
33}
34
35/**
36 * Check if vote reference is in proof
37*/
38pub fn is_vote_in_proof(vote: &VoteReference, proof: &Proof) -> bool {
39    if let Some(leaf_branch) = proof.branches.get(0) {
40        match leaf_branch {
41            (None, None) => false,
42            (Some(l), None) => &l.0 == vote,
43            (Some(l), Some(r)) => &l.0 == vote || &r.0 == vote,
44            (None, Some(r)) => &r.0 == vote,
45        }
46    } else {
47        false
48    }
49}
50
51/**
52 * Given two vote references and merkle tally proof of their inclusion, check
53 * if they are next to each other in the tree (neighbors). They do not need
54 * to share same parent.
55 *
56 * Returns error if merkle proof is invalid for vote reference.
57 *
58 * Example:
59 * ```
60 * use tallytree::generate::generate_tree;
61 * use tallytree::proof::util::{are_neighboors, create_merkle_proof};
62 * use tallytree::navigate::find_node_by_vote_reference;
63 * use tallytree::{Validation, VoteReference};
64 * let a_vote: VoteReference = [0xaa; 32];
65 * let b_vote: VoteReference = [0xbb; 32];
66 * let c_vote: VoteReference = [0xcc; 32];
67 *
68 * let tree = generate_tree(vec![
69 *      (a_vote, vec![1, 0]),
70 *      (b_vote, vec![1, 0]),
71 *      (c_vote, vec![0, 1]),
72 * ], true).unwrap();
73 *
74 * let v = &Validation::Strict;
75 * let a_proof = create_merkle_proof(&find_node_by_vote_reference(&tree, &a_vote).unwrap(), v).unwrap();
76 * let b_proof = create_merkle_proof(&find_node_by_vote_reference(&tree, &b_vote).unwrap(), v).unwrap();
77 * let c_proof = create_merkle_proof(&find_node_by_vote_reference(&tree, &c_vote).unwrap(), v).unwrap();
78 *
79 * assert!(are_neighboors((&a_vote, &a_proof), (&b_vote, &b_proof), v).unwrap());
80 * assert!(!are_neighboors((&a_vote, &a_proof), (&c_vote, &c_proof), v).unwrap());
81 * assert!(are_neighboors((&b_vote, &b_proof), (&c_vote, &c_proof), v).unwrap());
82 * ```
83 */
84pub fn are_neighboors(
85    left: (&VoteReference, &Proof),
86    right: (&VoteReference, &Proof),
87    v: &Validation,
88) -> Result<bool, String> {
89    let (left_vote, left_proof) = left;
90    let (right_vote, right_proof) = right;
91
92    if left_vote == right_vote {
93        return Ok(false);
94    }
95
96    if !is_vote_in_proof(left_vote, left_proof) {
97        return Err("Left vote reference does not belong to proof".to_string());
98    }
99    if !is_vote_in_proof(right_vote, right_proof) {
100        return Err("Right vote reference does not belong to proof".to_string());
101    }
102
103    // If left vote is also in right proof, then both votes must have same parent.
104    // In a binary tree they must be neighboors.
105    if is_vote_in_proof(left_vote, right_proof) {
106        if !is_vote_in_proof(right_vote, left_proof) {
107            return Err("If left vote reference is in right proof, then right vote refrence should be in left proof".to_string());
108        }
109        return Ok(true);
110    }
111
112    // For the two nodes to be next to each other, they should have a common parent
113    // where the left node is the right-most node in the parents subtree, and the
114    // right node is the left-mode node in the parents subtree.
115    if left_proof.branches.len() != right_proof.branches.len() {
116        return Err(
117            "Expected the branch length of left and right proof to be the same".to_string(),
118        );
119    }
120
121    let ((left_merkle_root, _), left_directions) = collapse_branches(
122        left_proof.branches.clone(),
123        &CollapseBranchConstrains::None,
124        v,
125    )?;
126    let ((right_merkle_root, _), right_directions) = collapse_branches(
127        right_proof.branches.clone(),
128        &CollapseBranchConstrains::None,
129        v,
130    )?;
131
132    if left_merkle_root != right_merkle_root {
133        return Err(format!(
134            "Merkle root hash for proofs don't match ({} != {})",
135            to_hex(&left_merkle_root)?,
136            to_hex(&right_merkle_root)?
137        ));
138    }
139
140    let left_index = find_node_index(
141        left_vote,
142        left_proof.branches.get(0).ok_or("Empty proof")?,
143        &left_directions,
144    )?;
145    let right_index = find_node_index(
146        right_vote,
147        right_proof.branches.get(0).ok_or("Empty proof")?,
148        &right_directions,
149    )?;
150
151    Ok(left_index + 1 == right_index || left_index != 0 && left_index - 1 == right_index)
152}
153
154/// Collapse merkle proof branches into a single hash and tally.
155///
156/// If `allow_left` is set, we restrict the function so that it fails if the
157/// proof is not for the right-most node of the tree. This is used in the
158/// 'verify proof of vote count' function.
159///
160/// ```
161/// use tallytree::proof::util::{collapse_branches, create_merkle_proof, CollapseBranchConstrains};
162/// use tallytree::navigate::find_node_by_vote_reference;
163/// use tallytree::{Validation, VoteReference};
164/// use tallytree::hash::hash_node_ref;
165/// use tallytree::generate::generate_tree;
166///
167/// let tree = generate_tree(vec![
168///      ([0xaa; 32], vec![1, 0]),
169///      ([0xbb; 32], vec![1, 0]),
170///      ([0xcc; 32], vec![0, 1]),
171/// ], true).unwrap();
172///
173/// let v = &Validation::Strict;
174/// let proof = create_merkle_proof(
175///     &find_node_by_vote_reference(&tree, &[0xaa; 32]).unwrap(),
176///     v
177/// ).unwrap();
178///
179/// let ((branches_hash, _), _) = collapse_branches(
180///     proof.branches.clone(),
181///     &CollapseBranchConstrains::None,
182///     v
183/// ).unwrap();
184/// let (root_hash, _) = hash_node_ref(&tree, v).unwrap().unwrap();
185/// assert_eq!(*root_hash.as_array(), branches_hash);
186/// ```
187pub fn collapse_branches(
188    mut branches: Vec<ProofBranch>,
189    constraints: &CollapseBranchConstrains,
190    v: &Validation,
191) -> Result<(ProofNode, Vec<TraverseDirection>), String> {
192    if branches.is_empty() {
193        return Err("Can't collapse empty branches".to_string());
194    }
195    let has_leaf = branches.len() == 1;
196
197    if has_leaf {
198        let branch = derive_wrapper_branch(&branches[0], v)?;
199        let dir = if branch.1.as_ref().unwrap().0 == *NULL_HASH.as_array() {
200            TraverseDirection::Left
201        } else {
202            // Both left and right have valid leafs. We don't know which one
203            // the proof was for.
204            TraverseDirection::Right
205        };
206        return Ok((
207            tally_and_hash(
208                &branch.0.ok_or("Invalid leaf branch")?,
209                &branch.1.ok_or("Invalid leaf branch")?,
210            )?,
211            vec![dir],
212        ));
213    };
214
215    let (proof_left, proof_right) = branches.pop().unwrap();
216
217    if proof_left.is_none() && proof_right.is_none() {
218        return Err("Both left and right nodes cannot be None".to_string());
219    }
220
221    if proof_left.is_some() && proof_right.is_some() {
222        return Err("Both left and right cannot exist for non-leaf branches".to_string());
223    }
224
225    let mut directions: Vec<TraverseDirection> = vec![];
226    let left = if let Some(l) = proof_left {
227        l
228    } else {
229        if matches!(constraints, CollapseBranchConstrains::RejectLeftPath) {
230            // ```text
231            //   N
232            //  / \
233            // N   Ø
234            // ```
235            //
236            // In case the right node is a Ø-node (NULL_HASH), we must still
237            // allow taking the left path. Ø-node are just fillers to balance
238            // leaf nodes at same height.
239
240            let (hash, _) = proof_right.as_ref().ok_or("Invalid right proof")?;
241            if hash != NULL_HASH.as_array() {
242                return Err("Proof is constrained from collapsing to the left".to_string());
243            }
244        }
245        let (n, d) = collapse_branches(branches.clone(), constraints, v)?;
246        directions = [vec![TraverseDirection::Left], d].concat();
247        n
248    };
249    let right = if let Some(r) = proof_right {
250        r
251    } else {
252        if matches!(constraints, CollapseBranchConstrains::RejectRightPath) {
253            return Err("Proof is constrained from collapsing to the right".to_string());
254        }
255        let (n, d) = collapse_branches(branches, constraints, v)?;
256        directions = [vec![TraverseDirection::Right], d].concat();
257        n
258    };
259    Ok((tally_and_hash(&left, &right)?, directions))
260}
261
262/// Given the deepest proof branch, derive the above wrapper nodes.
263///
264/// Example: Given the branch with nodes (V3 and Ø), we return (F and Ø).
265///          Given the branch with nodes (V1 and V2), we return (D and E)
266/// ```text
267///            A
268///          /  \
269///         B    C
270///        / \   | \
271///       D   E  F  Ø
272///       |   |  |
273///       V1  V2 V3
274/// ```
275pub fn derive_wrapper_branch(branch: &ProofBranch, v: &Validation) -> Result<ProofBranch, String> {
276    let (left_proof_hash, left_tally) = branch.0.as_ref().ok_or("Invalid left leaf")?;
277    let (right_proof_hash, right_tally) = branch.1.as_ref().ok_or("Invalid right leaf")?;
278
279    let left_tally = match left_tally {
280        None => Err("Invalid proof, no tally on left node.".to_string()),
281        Some(t) => {
282            if !matches!(v, Validation::Relaxed) && !has_one_vote(t) {
283                Err("Invalid proof. Left node should have exactly one vote.".to_string())
284            } else {
285                Ok(t)
286            }
287        }
288    }?;
289    if right_proof_hash != NULL_HASH.as_array()
290        && !has_one_vote(right_tally.as_ref().unwrap_or(&vec![]))
291    {
292        return Err(
293            "Invalid proof. Right node should be null or have exactly one vote.".to_string(),
294        );
295    }
296
297    // The deepest nodes in the proof point to leafs or Ø. We need to find the
298    // hash for the wrapper nodes.
299    let left_hash = derive_wrapper_hash(&hash_leaf(left_proof_hash, left_tally), left_tally);
300    let right_hash = if right_proof_hash == NULL_HASH.as_array() {
301        *NULL_HASH.as_array()
302    } else {
303        derive_wrapper_hash(
304            &hash_leaf(right_proof_hash, right_tally.as_ref().unwrap()),
305            right_tally.as_ref().unwrap(),
306        )
307    };
308
309    Ok((
310        Some((left_hash, Some(left_tally.clone()))),
311        Some((right_hash, right_tally.clone())),
312    ))
313}
314
315/// Given a leaf hash and tally, calculate the hash of its wrapper node in a
316/// merkle tally tree.
317///
318/// Example:
319/// ```
320/// use tallytree::generate::generate_tree;
321/// use tallytree::proof::util::derive_wrapper_hash;
322/// use tallytree::hash::{hash_node_ref, hash_leaf};
323/// use tallytree::Validation;
324///
325/// let v = &Validation::Strict;
326/// let (vote_reference, vote) = ([0xaa; 32], vec![1, 0]);
327/// let tree = generate_tree(vec![
328///     (vote_reference, vote.clone()),
329/// ], false).unwrap();
330/// let wrapper_hash = derive_wrapper_hash(&hash_leaf(&vote_reference, &vote), &vote);
331/// assert_eq!(&wrapper_hash,
332///     hash_node_ref(&tree.unwrap().as_ref().left, v).unwrap().unwrap().0.as_array());
333/// ```
334pub fn derive_wrapper_hash(hash: &NodeHash, tally: &[u32]) -> ProofHash {
335    let mut hasher = Sha256::new();
336    tally
337        .iter()
338        .map(|t| t.to_be_bytes())
339        .for_each(|b| hasher.update(b));
340    hasher.update(hash.as_slice());
341    hasher
342        .finalize()
343        .as_slice()
344        .try_into()
345        .expect("Expected hash to be 32 bytes")
346}
347
348/// Calculate a index to a node given traverse directions to it in the merkle
349/// tree.
350///
351/// Note: Directions given by `collapse_branches` may cause the index to
352/// be off-by-one. Use `find_node_index` for always accurate result.
353///
354/// The index is zero based numbered.
355///
356/// Example:
357/// ```
358/// use tallytree::proof::util::{collapse_branches, create_merkle_proof, CollapseBranchConstrains, directions_to_node_index};
359/// use tallytree::navigate::find_node_by_vote_reference;
360/// use tallytree::{Validation, VoteReference};
361/// use tallytree::generate::generate_tree;
362///
363/// let tree = generate_tree(vec![
364///      ([0xaa; 32], vec![1, 0]),
365///      ([0xbb; 32], vec![1, 0]),
366///      ([0xcc; 32], vec![0, 1]),
367/// ], true).unwrap();
368///
369/// let v = &Validation::Strict;
370/// let proof = create_merkle_proof(
371///     &find_node_by_vote_reference(&tree, &[0xcc; 32]).unwrap(),
372///     v
373/// ).unwrap();
374///
375/// let ((_, _), directions) = collapse_branches(
376///     proof.branches.clone(),
377///     &CollapseBranchConstrains::None,
378///     v
379/// ).unwrap();
380/// assert_eq!(2, directions_to_node_index(&directions));
381/// ```
382pub fn directions_to_node_index(directions: &[TraverseDirection]) -> usize {
383    let mut index: usize = 0;
384    for dir in directions.iter() {
385        index <<= 1;
386        index |= *dir as usize;
387    }
388    index
389}
390
391/// Given a path to a node, create a merkle proof.
392///
393/// Example:
394/// ```
395/// use tallytree::proof::util::create_merkle_proof;
396/// use tallytree::navigate::find_node_by_vote_reference;
397/// use tallytree::generate::generate_tree;
398/// use tallytree::Validation;
399///
400/// let tree = generate_tree(vec![
401///      ([0xaa; 32], vec![1, 0]),
402///      ([0xbb; 32], vec![1, 0]),
403///      ([0xcc; 32], vec![0, 1]),
404/// ], true).unwrap();
405///
406/// let proof = create_merkle_proof(
407///     &find_node_by_vote_reference(&tree, &[0xcc; 32]).unwrap(),
408///     &Validation::Strict,
409/// );
410/// ```
411pub fn create_merkle_proof(path: &[NodeRef], v: &Validation) -> Result<Proof, String> {
412    let mut proof = Proof::default();
413
414    let mut added_nodes: HashSet<ProofHash> = HashSet::new();
415
416    for b in path.iter().rev() {
417        if is_leaf_node(b) || is_wrapper_node(b) {
418            continue;
419        }
420        if let Some(h) = hash_node_ref(b, v)? {
421            // This node can be derived from it's children in the merkle proof
422            // and shall not be explicitly added in a later branch.
423            added_nodes.insert(*h.0.as_array());
424        }
425
426        let none_if_added = |node: Option<(NodeHash, Option<TallyList>)>| {
427            if let Some((hash, tally)) = node {
428                if added_nodes.contains(hash.as_array()) {
429                    None
430                } else {
431                    Some((*hash.as_array(), tally))
432                }
433            } else {
434                None
435            }
436        };
437
438        let leaf_to_proof_node = |leaf: &NodeRef| -> Result<Option<ProofNode>, String> {
439            let leaf = leaf.as_ref().ok_or("Missing leaf node")?;
440            let vote = leaf
441                .vote
442                .as_ref()
443                .ok_or("Invalid leaf node, vote missing")?;
444
445            // Leaf nodes are represented with their vote reference instead of hash.
446            // The hash can be derived from the vote reference and vote cast (tally).
447            Ok(Some((vote.0, Some(vote.1.clone()))))
448        };
449
450        let b = b.as_ref().ok_or("None node in path")?;
451        let left = if is_wrapper_node(&b.left) {
452            leaf_to_proof_node(&b.left.as_ref().unwrap().left)?
453        } else {
454            none_if_added(hash_node_ref(&b.left, v)?)
455        };
456        let right = if is_wrapper_node(&b.right) {
457            leaf_to_proof_node(&b.right.as_ref().unwrap().left)?
458        } else {
459            none_if_added(hash_node_ref(&b.right, v)?)
460        };
461
462        proof.branches.push((left, right))
463    }
464    Ok(proof)
465}
466
467/// Find the index of a vote in a merkle tree using the merkle proof for the vote.
468/// Use `collapse_branches` on the proof to get the directions parameter.
469/// `leaf_branch` is the zero-indexed branch in the merkle proof.
470///
471/// Example:
472/// ```
473/// use tallytree::proof::util::{collapse_branches, create_merkle_proof, CollapseBranchConstrains, find_node_index};
474/// use tallytree::navigate::find_node_by_vote_reference;
475/// use tallytree::{Validation, VoteReference};
476/// use tallytree::generate::generate_tree;
477///
478/// let tree = generate_tree(vec![
479///      ([0xaa; 32], vec![1, 0]),
480///      ([0xbb; 32], vec![1, 0]),
481///      ([0xcc; 32], vec![0, 1]),
482/// ], true).unwrap();
483///
484/// let v = &Validation::Strict;
485/// let proof = create_merkle_proof(
486///     &find_node_by_vote_reference(&tree, &[0xbb; 32]).unwrap(),
487///     v
488/// ).unwrap();
489///
490/// let ((_, _), directions) = collapse_branches(
491///     proof.branches.clone(),
492///     &CollapseBranchConstrains::None,
493///     v
494/// ).unwrap();
495/// assert_eq!(1, find_node_index(&[0xbb; 32], &proof.branches[0], &directions).unwrap());
496/// ```
497pub fn find_node_index(
498    reference: &VoteReference,
499    leaf_branch: &ProofBranch,
500    directions: &[TraverseDirection],
501) -> Result<usize, String> {
502    let directions_index = directions_to_node_index(directions);
503
504    // If leaf branch has two nodes, directions index assumes it's the right one.
505    // We need to check if it's the left one.
506    let (left, right) = leaf_branch;
507    let left = left.as_ref().ok_or("Left proof node missing")?;
508    let right = right.as_ref().ok_or("Right proof node missing")?;
509
510    if &right.0 != NULL_HASH.as_array() && &left.0 == reference {
511        Ok(directions_index - 1)
512    } else {
513        Ok(directions_index)
514    }
515}
516
517/// Get the tally and hash for parent node of two sibling nodes.
518pub fn tally_and_hash(left: &ProofNode, right: &ProofNode) -> Result<ProofNode, String> {
519    let mut hasher = Sha256::new();
520    let tally = combine_tally(&left.1, &right.1)?;
521    tally
522        .iter()
523        .map(|t| t.to_be_bytes())
524        .for_each(|b| hasher.update(b));
525
526    hasher.update(left.0.as_slice());
527    hasher.update(right.0.as_slice());
528    let hash: [u8; 32] = hasher
529        .finalize()
530        .as_slice()
531        .try_into()
532        .expect("Expected hash to be 32 bytes");
533
534    Ok((hash, Some(tally)))
535}
536
537#[cfg(test)]
538mod tests {
539    use super::*;
540    use crate::generate::generate_tree;
541    use crate::navigate::find_node_by_vote_reference;
542    use crate::proof::votecount::create_proof_of_vote_count;
543
544    #[test]
545    fn test_are_neighboors() {
546        // ```text
547        //            A
548        //          /  \
549        //         B    C
550        //        / \   | \
551        //       D   E  F  Ø
552        //       |   |  |
553        //       V1  V2 V3
554        // ```
555        let a_vote: VoteReference = [0xaa; 32];
556        let b_vote: VoteReference = [0xbb; 32];
557        let c_vote: VoteReference = [0xcc; 32];
558
559        let root = generate_tree(
560            vec![
561                (a_vote, vec![1, 0]),
562                (b_vote, vec![1, 0]),
563                (c_vote, vec![0, 1]),
564            ],
565            true,
566        )
567        .unwrap();
568
569        let v = &Validation::Strict;
570        let a_proof =
571            create_merkle_proof(&find_node_by_vote_reference(&root, &a_vote).unwrap(), v).unwrap();
572        let b_proof =
573            create_merkle_proof(&find_node_by_vote_reference(&root, &b_vote).unwrap(), v).unwrap();
574        let c_proof =
575            create_merkle_proof(&find_node_by_vote_reference(&root, &c_vote).unwrap(), v).unwrap();
576
577        assert!(are_neighboors((&a_vote, &a_proof), (&b_vote, &b_proof), v).unwrap());
578        assert!(!are_neighboors((&a_vote, &a_proof), (&c_vote, &c_proof), v).unwrap());
579        assert!(are_neighboors((&b_vote, &b_proof), (&c_vote, &c_proof), v).unwrap());
580
581        let err = are_neighboors((&a_vote, &a_proof), (&c_vote, &a_proof), v);
582        assert_eq!(
583            err,
584            Err("Right vote reference does not belong to proof".to_string())
585        );
586    }
587
588    #[test]
589    fn test_create_merkle_proof() {
590        // ```text
591        //            A
592        //          /  \
593        //         B    C
594        //        / \   | \
595        //       D   E  F  Ø
596        //       |   |  |
597        //       V1  V2 V3
598        // ```
599        let root = generate_tree(
600            vec![
601                ([0xaa; 32], vec![1, 0]),
602                ([0xbb; 32], vec![1, 0]),
603                ([0xcc; 32], vec![0, 1]),
604            ],
605            true,
606        )
607        .unwrap();
608        let a = root.clone();
609        let b = a.as_ref().unwrap().left.clone();
610        let c = a.as_ref().unwrap().right.clone();
611        let d = b.as_ref().unwrap().left.clone();
612        let v1 = d.as_ref().unwrap().left.clone();
613        let path = [a, b, d, v1.clone()];
614        let proof = create_merkle_proof(&path, &Validation::Strict).unwrap();
615
616        // Proof should be
617        // * V1, V2
618        // * (None), C
619        assert_eq!(proof.branches.len(), 2);
620        let (proof_v1, proof_v2) = &proof.branches[0];
621        assert_eq!(proof_v1.as_ref().unwrap(), &([0xaa; 32], Some(vec![1, 0])));
622        assert_eq!(proof_v2.as_ref().unwrap(), &([0xbb; 32], Some(vec![1, 0])));
623
624        let (left, proof_c) = &proof.branches[1];
625        assert!(left.is_none());
626        let (c_hash, c_tally) = hash_node_ref(&c, &Validation::Strict).unwrap().unwrap();
627        assert_eq!(proof_c, &Some((*c_hash.as_array(), c_tally)));
628    }
629
630    #[test]
631    fn test_collapse_branches_directions() {
632        // ```text
633        //           N
634        //          / \
635        //         /   \
636        //        /     \
637        //       N       N
638        //      / \     / \
639        //     N   N   N   Ø
640        //    / \  |\  | \
641        //   V1 V  V V V5 Ø
642        // ```
643        //
644        // Collapsing the merkle proof for V5, Ø should give directions "Right, Left"
645        let root = generate_tree(
646            vec![
647                ([0xaa; 32], vec![1, 0]),
648                ([0xbb; 32], vec![0, 1]),
649                ([0xcc; 32], vec![0, 1]),
650                ([0xdd; 32], vec![1, 0]),
651                ([0xee; 32], vec![1, 0]),
652            ],
653            true,
654        )
655        .unwrap();
656        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
657
658        let (_, directions) = collapse_branches(
659            proof.branches,
660            &CollapseBranchConstrains::RejectLeftPath,
661            &Validation::Strict,
662        )
663        .unwrap();
664
665        assert_eq!(
666            directions,
667            vec![
668                TraverseDirection::Right,
669                TraverseDirection::Left,
670                TraverseDirection::Left
671            ]
672        );
673    }
674
675    #[test]
676    fn test_directions_to_node_index() {
677        let root = generate_tree(
678            vec![
679                ([0xaa; 32], vec![1, 0]),
680                ([0xbb; 32], vec![0, 1]),
681                ([0xcc; 32], vec![0, 1]),
682                ([0xdd; 32], vec![1, 0]),
683                ([0xee; 32], vec![1, 0]),
684            ],
685            true,
686        )
687        .unwrap();
688        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
689
690        let (_, directions) = collapse_branches(
691            proof.branches,
692            &CollapseBranchConstrains::RejectLeftPath,
693            &Validation::Strict,
694        )
695        .unwrap();
696
697        assert_eq!(directions_to_node_index(&directions), 4);
698    }
699
700    #[test]
701    fn test_collapse_branches_constraints() {
702        // ```text
703        //            A
704        //          /  \
705        //         B    C
706        //        / \   | \
707        //       D   E  F  Ø
708        //       |   |  |
709        //       V1  V2 V3
710        // ```
711
712        let tree = generate_tree(
713            vec![
714                ([0xaa; 32], vec![1, 0]),
715                ([0xbb; 32], vec![1, 0]),
716                ([0xcc; 32], vec![0, 1]),
717            ],
718            true,
719        )
720        .unwrap();
721        let v = &Validation::Strict;
722        let root_hash_and_tally = hash_node_ref(&tree, v).unwrap().unwrap();
723        let expected = (*root_hash_and_tally.0.as_array(), root_hash_and_tally.1);
724
725        let left_most_node = find_node_by_vote_reference(&tree, &[0xaa; 32]).unwrap();
726        let middle_node = find_node_by_vote_reference(&tree, &[0xbb; 32]).unwrap();
727        let right_most_node = find_node_by_vote_reference(&tree, &[0xcc; 32]).unwrap();
728
729        let left_proof = create_merkle_proof(&left_most_node, v).unwrap();
730        let middle_proof = create_merkle_proof(&middle_node, v).unwrap();
731        let right_proof = create_merkle_proof(&right_most_node, &v).unwrap();
732
733        // Left
734        assert_eq!(
735            expected,
736            collapse_branches(
737                left_proof.branches.clone(),
738                &CollapseBranchConstrains::None,
739                &v
740            )
741            .unwrap()
742            .0
743        );
744
745        assert_eq!(
746            Err("Proof is constrained from collapsing to the left".to_string()),
747            collapse_branches(
748                left_proof.branches.clone(),
749                &CollapseBranchConstrains::RejectLeftPath,
750                &v
751            )
752        );
753
754        assert_eq!(
755            expected,
756            collapse_branches(
757                left_proof.branches,
758                &CollapseBranchConstrains::RejectRightPath,
759                &v
760            )
761            .unwrap()
762            .0
763        );
764
765        // Middle
766        assert_eq!(
767            expected,
768            collapse_branches(
769                middle_proof.branches.clone(),
770                &CollapseBranchConstrains::None,
771                &v
772            )
773            .unwrap()
774            .0
775        );
776
777        assert_eq!(
778            Err("Proof is constrained from collapsing to the left".to_string()),
779            collapse_branches(
780                middle_proof.branches.clone(),
781                &CollapseBranchConstrains::RejectLeftPath,
782                &v
783            )
784        );
785
786        // Although the node is on the right, it's not hang on a parent that is on the right.
787        assert_eq!(
788            expected,
789            collapse_branches(
790                middle_proof.branches,
791                &CollapseBranchConstrains::RejectRightPath,
792                &v
793            )
794            .unwrap()
795            .0
796        );
797
798        // Right
799        assert_eq!(
800            expected,
801            collapse_branches(
802                right_proof.branches.clone(),
803                &CollapseBranchConstrains::None,
804                &v
805            )
806            .unwrap()
807            .0
808        );
809
810        assert_eq!(
811            expected,
812            collapse_branches(
813                right_proof.branches.clone(),
814                &CollapseBranchConstrains::RejectLeftPath,
815                &v
816            )
817            .unwrap()
818            .0
819        );
820
821        assert_eq!(
822            Err("Proof is constrained from collapsing to the right".to_string()),
823            collapse_branches(
824                right_proof.branches,
825                &CollapseBranchConstrains::RejectRightPath,
826                &v
827            )
828        );
829    }
830
831    #[test]
832    fn test_find_node_index() {
833        // ```text
834        //            A
835        //          /  \
836        //         B    C
837        //        / \   | \
838        //       D   E  F  Ø
839        //       |   |  |
840        //       V1  V2 V3
841        // ```
842        let root = generate_tree(
843            vec![
844                ([0xaa; 32], vec![1, 0]),
845                ([0xbb; 32], vec![1, 0]),
846                ([0xcc; 32], vec![0, 1]),
847            ],
848            true,
849        )
850        .unwrap();
851        let v = &Validation::Strict;
852        let a_proof =
853            create_merkle_proof(&find_node_by_vote_reference(&root, &[0xaa; 32]).unwrap(), v)
854                .unwrap();
855        let b_proof =
856            create_merkle_proof(&find_node_by_vote_reference(&root, &[0xbb; 32]).unwrap(), v)
857                .unwrap();
858        let c_proof =
859            create_merkle_proof(&find_node_by_vote_reference(&root, &[0xcc; 32]).unwrap(), v)
860                .unwrap();
861
862        let (_, a_directions) =
863            collapse_branches(a_proof.branches.clone(), &CollapseBranchConstrains::None, v)
864                .unwrap();
865        let (_, b_directions) =
866            collapse_branches(b_proof.branches.clone(), &CollapseBranchConstrains::None, v)
867                .unwrap();
868        let (_, c_directions) =
869            collapse_branches(c_proof.branches.clone(), &CollapseBranchConstrains::None, v)
870                .unwrap();
871
872        assert_eq!(
873            0,
874            find_node_index(&[0xaa; 32], &a_proof.branches[0], &a_directions).unwrap()
875        );
876        assert_eq!(
877            1,
878            find_node_index(&[0xbb; 32], &b_proof.branches[0], &b_directions).unwrap()
879        );
880        assert_eq!(
881            2,
882            find_node_index(&[0xcc; 32], &c_proof.branches[0], &c_directions).unwrap()
883        );
884    }
885}
tallytree/proof/util.rs

tallytree/proof/
util.rs
