tallytree 0.3.4

use crate::hash::{hash_leaf, hash_node_ref, NodeHash, NULL_HASH};
use crate::node::{is_leaf_node, is_wrapper_node};
use crate::proof::{Proof, ProofBranch, ProofHash, ProofNode};
use crate::tally::{combine_tally, has_one_vote, TallyList};
use crate::utilstring::to_hex;
use crate::{NodeRef, Validation, VoteReference};
use sha2::{Digest, Sha256};
use std::collections::HashSet;

/// The direction traversed when parsing a merkle proof branch
#[derive(Clone, Debug, PartialEq, Copy)]
pub enum TraverseDirection {
    /// Merkle proof recurses to the left from current branch
    Left = 0,
    /// Merkle proof recurses to the right from current branch
    Right = 1,
}

/// Set constraints on how branch is allowed to collapse the branches.
pub enum CollapseBranchConstrains {
    /// No constraints
    None = 0,
    /// Fail if the proof is not for the right-most branch of the tree.
    /// The leaf can still hang on the left on its parent, but any parent
    /// cannot be on the left of its parent.
    /// (Don't allow collapsing to the left)
    RejectLeftPath = 1,
    /// Fail if the proof is not for the left-most branch of the tree.
    /// The leaf can still hang on the right on its parent, but any parent
    /// cannot be on the right of its parent.
    /// (Don't allow collapsing to the right)
    RejectRightPath = 2,
}

/**
 * Check if vote reference is in proof
*/
pub fn is_vote_in_proof(vote: &VoteReference, proof: &Proof) -> bool {
    if let Some(leaf_branch) = proof.branches.get(0) {
        match leaf_branch {
            (None, None) => false,
            (Some(l), None) => &l.0 == vote,
            (Some(l), Some(r)) => &l.0 == vote || &r.0 == vote,
            (None, Some(r)) => &r.0 == vote,
        }
    } else {
        false
    }
}

/**
 * Given two vote references and merkle tally proof of their inclusion, check
 * if they are next to each other in the tree (neighbors). They do not need
 * to share same parent.
 *
 * Returns error if merkle proof is invalid for vote reference.
 *
 * Example:
 * ```
 * use tallytree::generate::generate_tree;
 * use tallytree::proof::util::{are_neighboors, create_merkle_proof};
 * use tallytree::navigate::find_node_by_vote_reference;
 * use tallytree::{Validation, VoteReference};
 * let a_vote: VoteReference = [0xaa; 32];
 * let b_vote: VoteReference = [0xbb; 32];
 * let c_vote: VoteReference = [0xcc; 32];
 *
 * let tree = generate_tree(vec![
 *      (a_vote, vec![1, 0]),
 *      (b_vote, vec![1, 0]),
 *      (c_vote, vec![0, 1]),
 * ], true).unwrap();
 *
 * let v = &Validation::Strict;
 * let a_proof = create_merkle_proof(&find_node_by_vote_reference(&tree, &a_vote).unwrap(), v).unwrap();
 * let b_proof = create_merkle_proof(&find_node_by_vote_reference(&tree, &b_vote).unwrap(), v).unwrap();
 * let c_proof = create_merkle_proof(&find_node_by_vote_reference(&tree, &c_vote).unwrap(), v).unwrap();
 *
 * assert!(are_neighboors((&a_vote, &a_proof), (&b_vote, &b_proof), v).unwrap());
 * assert!(!are_neighboors((&a_vote, &a_proof), (&c_vote, &c_proof), v).unwrap());
 * assert!(are_neighboors((&b_vote, &b_proof), (&c_vote, &c_proof), v).unwrap());
 * ```
 */
pub fn are_neighboors(
    left: (&VoteReference, &Proof),
    right: (&VoteReference, &Proof),
    v: &Validation,
) -> Result<bool, String> {
    let (left_vote, left_proof) = left;
    let (right_vote, right_proof) = right;

    if left_vote == right_vote {
        return Ok(false);
    }

    if !is_vote_in_proof(left_vote, left_proof) {
        return Err("Left vote reference does not belong to proof".to_string());
    }
    if !is_vote_in_proof(right_vote, right_proof) {
        return Err("Right vote reference does not belong to proof".to_string());
    }

    // If left vote is also in right proof, then both votes must have same parent.
    // In a binary tree they must be neighboors.
    if is_vote_in_proof(left_vote, right_proof) {
        if !is_vote_in_proof(right_vote, left_proof) {
            return Err("If left vote reference is in right proof, then right vote refrence should be in left proof".to_string());
        }
        return Ok(true);
    }

    // For the two nodes to be next to each other, they should have a common parent
    // where the left node is the right-most node in the parents subtree, and the
    // right node is the left-mode node in the parents subtree.
    if left_proof.branches.len() != right_proof.branches.len() {
        return Err(
            "Expected the branch length of left and right proof to be the same".to_string(),
        );
    }

    let ((left_merkle_root, _), left_directions) = collapse_branches(
        left_proof.branches.clone(),
        &CollapseBranchConstrains::None,
        v,
    )?;
    let ((right_merkle_root, _), right_directions) = collapse_branches(
        right_proof.branches.clone(),
        &CollapseBranchConstrains::None,
        v,
    )?;

    if left_merkle_root != right_merkle_root {
        return Err(format!(
            "Merkle root hash for proofs don't match ({} != {})",
            to_hex(&left_merkle_root)?,
            to_hex(&right_merkle_root)?
        ));
    }

    let left_index = find_node_index(
        left_vote,
        left_proof.branches.get(0).ok_or("Empty proof")?,
        &left_directions,
    )?;
    let right_index = find_node_index(
        right_vote,
        right_proof.branches.get(0).ok_or("Empty proof")?,
        &right_directions,
    )?;

    Ok(left_index + 1 == right_index || left_index != 0 && left_index - 1 == right_index)
}

/// Collapse merkle proof branches into a single hash and tally.
///
/// If `allow_left` is set, we restrict the function so that it fails if the
/// proof is not for the right-most node of the tree. This is used in the
/// 'verify proof of vote count' function.
///
/// ```
/// use tallytree::proof::util::{collapse_branches, create_merkle_proof, CollapseBranchConstrains};
/// use tallytree::navigate::find_node_by_vote_reference;
/// use tallytree::{Validation, VoteReference};
/// use tallytree::hash::hash_node_ref;
/// use tallytree::generate::generate_tree;
///
/// let tree = generate_tree(vec![
///      ([0xaa; 32], vec![1, 0]),
///      ([0xbb; 32], vec![1, 0]),
///      ([0xcc; 32], vec![0, 1]),
/// ], true).unwrap();
///
/// let v = &Validation::Strict;
/// let proof = create_merkle_proof(
///     &find_node_by_vote_reference(&tree, &[0xaa; 32]).unwrap(),
///     v
/// ).unwrap();
///
/// let ((branches_hash, _), _) = collapse_branches(
///     proof.branches.clone(),
///     &CollapseBranchConstrains::None,
///     v
/// ).unwrap();
/// let (root_hash, _) = hash_node_ref(&tree, v).unwrap().unwrap();
/// assert_eq!(*root_hash.as_array(), branches_hash);
/// ```
pub fn collapse_branches(
    mut branches: Vec<ProofBranch>,
    constraints: &CollapseBranchConstrains,
    v: &Validation,
) -> Result<(ProofNode, Vec<TraverseDirection>), String> {
    if branches.is_empty() {
        return Err("Can't collapse empty branches".to_string());
    }
    let has_leaf = branches.len() == 1;

    if has_leaf {
        let branch = derive_wrapper_branch(&branches[0], v)?;
        let dir = if branch.1.as_ref().unwrap().0 == *NULL_HASH.as_array() {
            TraverseDirection::Left
        } else {
            // Both left and right have valid leafs. We don't know which one
            // the proof was for.
            TraverseDirection::Right
        };
        return Ok((
            tally_and_hash(
                &branch.0.ok_or("Invalid leaf branch")?,
                &branch.1.ok_or("Invalid leaf branch")?,
            )?,
            vec![dir],
        ));
    };

    let (proof_left, proof_right) = branches.pop().unwrap();

    if proof_left.is_none() && proof_right.is_none() {
        return Err("Both left and right nodes cannot be None".to_string());
    }

    if proof_left.is_some() && proof_right.is_some() {
        return Err("Both left and right cannot exist for non-leaf branches".to_string());
    }

    let mut directions: Vec<TraverseDirection> = vec![];
    let left = if let Some(l) = proof_left {
        l
    } else {
        if matches!(constraints, CollapseBranchConstrains::RejectLeftPath) {
            // ```text
            //   N
            //  / \
            // N   Ø
            // ```
            //
            // In case the right node is a Ø-node (NULL_HASH), we must still
            // allow taking the left path. Ø-node are just fillers to balance
            // leaf nodes at same height.

            let (hash, _) = proof_right.as_ref().ok_or("Invalid right proof")?;
            if hash != NULL_HASH.as_array() {
                return Err("Proof is constrained from collapsing to the left".to_string());
            }
        }
        let (n, d) = collapse_branches(branches.clone(), constraints, v)?;
        directions = [vec![TraverseDirection::Left], d].concat();
        n
    };
    let right = if let Some(r) = proof_right {
        r
    } else {
        if matches!(constraints, CollapseBranchConstrains::RejectRightPath) {
            return Err("Proof is constrained from collapsing to the right".to_string());
        }
        let (n, d) = collapse_branches(branches, constraints, v)?;
        directions = [vec![TraverseDirection::Right], d].concat();
        n
    };
    Ok((tally_and_hash(&left, &right)?, directions))
}

/// Given the deepest proof branch, derive the above wrapper nodes.
///
/// Example: Given the branch with nodes (V3 and Ø), we return (F and Ø).
///          Given the branch with nodes (V1 and V2), we return (D and E)
/// ```text
///            A
///          /  \
///         B    C
///        / \   | \
///       D   E  F  Ø
///       |   |  |
///       V1  V2 V3
/// ```
pub fn derive_wrapper_branch(branch: &ProofBranch, v: &Validation) -> Result<ProofBranch, String> {
    let (left_proof_hash, left_tally) = branch.0.as_ref().ok_or("Invalid left leaf")?;
    let (right_proof_hash, right_tally) = branch.1.as_ref().ok_or("Invalid right leaf")?;

    let left_tally = match left_tally {
        None => Err("Invalid proof, no tally on left node.".to_string()),
        Some(t) => {
            if !matches!(v, Validation::Relaxed) && !has_one_vote(t) {
                Err("Invalid proof. Left node should have exactly one vote.".to_string())
            } else {
                Ok(t)
            }
        }
    }?;
    if right_proof_hash != NULL_HASH.as_array()
        && !has_one_vote(right_tally.as_ref().unwrap_or(&vec![]))
    {
        return Err(
            "Invalid proof. Right node should be null or have exactly one vote.".to_string(),
        );
    }

    // The deepest nodes in the proof point to leafs or Ø. We need to find the
    // hash for the wrapper nodes.
    let left_hash = derive_wrapper_hash(&hash_leaf(left_proof_hash, left_tally), left_tally);
    let right_hash = if right_proof_hash == NULL_HASH.as_array() {
        *NULL_HASH.as_array()
    } else {
        derive_wrapper_hash(
            &hash_leaf(right_proof_hash, right_tally.as_ref().unwrap()),
            right_tally.as_ref().unwrap(),
        )
    };

    Ok((
        Some((left_hash, Some(left_tally.clone()))),
        Some((right_hash, right_tally.clone())),
    ))
}

/// Given a leaf hash and tally, calculate the hash of its wrapper node in a
/// merkle tally tree.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::proof::util::derive_wrapper_hash;
/// use tallytree::hash::{hash_node_ref, hash_leaf};
/// use tallytree::Validation;
///
/// let v = &Validation::Strict;
/// let (vote_reference, vote) = ([0xaa; 32], vec![1, 0]);
/// let tree = generate_tree(vec![
///     (vote_reference, vote.clone()),
/// ], false).unwrap();
/// let wrapper_hash = derive_wrapper_hash(&hash_leaf(&vote_reference, &vote), &vote);
/// assert_eq!(&wrapper_hash,
///     hash_node_ref(&tree.unwrap().as_ref().left, v).unwrap().unwrap().0.as_array());
/// ```
pub fn derive_wrapper_hash(hash: &NodeHash, tally: &[u32]) -> ProofHash {
    let mut hasher = Sha256::new();
    tally
        .iter()
        .map(|t| t.to_be_bytes())
        .for_each(|b| hasher.update(b));
    hasher.update(hash.as_slice());
    hasher
        .finalize()
        .as_slice()
        .try_into()
        .expect("Expected hash to be 32 bytes")
}

/// Calculate a index to a node given traverse directions to it in the merkle
/// tree.
///
/// Note: Directions given by `collapse_branches` may cause the index to
/// be off-by-one. Use `find_node_index` for always accurate result.
///
/// The index is zero based numbered.
///
/// Example:
/// ```
/// use tallytree::proof::util::{collapse_branches, create_merkle_proof, CollapseBranchConstrains, directions_to_node_index};
/// use tallytree::navigate::find_node_by_vote_reference;
/// use tallytree::{Validation, VoteReference};
/// use tallytree::generate::generate_tree;
///
/// let tree = generate_tree(vec![
///      ([0xaa; 32], vec![1, 0]),
///      ([0xbb; 32], vec![1, 0]),
///      ([0xcc; 32], vec![0, 1]),
/// ], true).unwrap();
///
/// let v = &Validation::Strict;
/// let proof = create_merkle_proof(
///     &find_node_by_vote_reference(&tree, &[0xcc; 32]).unwrap(),
///     v
/// ).unwrap();
///
/// let ((_, _), directions) = collapse_branches(
///     proof.branches.clone(),
///     &CollapseBranchConstrains::None,
///     v
/// ).unwrap();
/// assert_eq!(2, directions_to_node_index(&directions));
/// ```
pub fn directions_to_node_index(directions: &[TraverseDirection]) -> usize {
    let mut index: usize = 0;
    for dir in directions.iter() {
        index <<= 1;
        index |= *dir as usize;
    }
    index
}

/// Given a path to a node, create a merkle proof.
///
/// Example:
/// ```
/// use tallytree::proof::util::create_merkle_proof;
/// use tallytree::navigate::find_node_by_vote_reference;
/// use tallytree::generate::generate_tree;
/// use tallytree::Validation;
///
/// let tree = generate_tree(vec![
///      ([0xaa; 32], vec![1, 0]),
///      ([0xbb; 32], vec![1, 0]),
///      ([0xcc; 32], vec![0, 1]),
/// ], true).unwrap();
///
/// let proof = create_merkle_proof(
///     &find_node_by_vote_reference(&tree, &[0xcc; 32]).unwrap(),
///     &Validation::Strict,
/// );
/// ```
pub fn create_merkle_proof(path: &[NodeRef], v: &Validation) -> Result<Proof, String> {
    let mut proof = Proof::default();

    let mut added_nodes: HashSet<ProofHash> = HashSet::new();

    for b in path.iter().rev() {
        if is_leaf_node(b) || is_wrapper_node(b) {
            continue;
        }
        if let Some(h) = hash_node_ref(b, v)? {
            // This node can be derived from it's children in the merkle proof
            // and shall not be explicitly added in a later branch.
            added_nodes.insert(*h.0.as_array());
        }

        let none_if_added = |node: Option<(NodeHash, Option<TallyList>)>| {
            if let Some((hash, tally)) = node {
                if added_nodes.contains(hash.as_array()) {
                    None
                } else {
                    Some((*hash.as_array(), tally))
                }
            } else {
                None
            }
        };

        let leaf_to_proof_node = |leaf: &NodeRef| -> Result<Option<ProofNode>, String> {
            let leaf = leaf.as_ref().ok_or("Missing leaf node")?;
            let vote = leaf
                .vote
                .as_ref()
                .ok_or("Invalid leaf node, vote missing")?;

            // Leaf nodes are represented with their vote reference instead of hash.
            // The hash can be derived from the vote reference and vote cast (tally).
            Ok(Some((vote.0, Some(vote.1.clone()))))
        };

        let b = b.as_ref().ok_or("None node in path")?;
        let left = if is_wrapper_node(&b.left) {
            leaf_to_proof_node(&b.left.as_ref().unwrap().left)?
        } else {
            none_if_added(hash_node_ref(&b.left, v)?)
        };
        let right = if is_wrapper_node(&b.right) {
            leaf_to_proof_node(&b.right.as_ref().unwrap().left)?
        } else {
            none_if_added(hash_node_ref(&b.right, v)?)
        };

        proof.branches.push((left, right))
    }
    Ok(proof)
}

/// Find the index of a vote in a merkle tree using the merkle proof for the vote.
/// Use `collapse_branches` on the proof to get the directions parameter.
/// `leaf_branch` is the zero-indexed branch in the merkle proof.
///
/// Example:
/// ```
/// use tallytree::proof::util::{collapse_branches, create_merkle_proof, CollapseBranchConstrains, find_node_index};
/// use tallytree::navigate::find_node_by_vote_reference;
/// use tallytree::{Validation, VoteReference};
/// use tallytree::generate::generate_tree;
///
/// let tree = generate_tree(vec![
///      ([0xaa; 32], vec![1, 0]),
///      ([0xbb; 32], vec![1, 0]),
///      ([0xcc; 32], vec![0, 1]),
/// ], true).unwrap();
///
/// let v = &Validation::Strict;
/// let proof = create_merkle_proof(
///     &find_node_by_vote_reference(&tree, &[0xbb; 32]).unwrap(),
///     v
/// ).unwrap();
///
/// let ((_, _), directions) = collapse_branches(
///     proof.branches.clone(),
///     &CollapseBranchConstrains::None,
///     v
/// ).unwrap();
/// assert_eq!(1, find_node_index(&[0xbb; 32], &proof.branches[0], &directions).unwrap());
/// ```
pub fn find_node_index(
    reference: &VoteReference,
    leaf_branch: &ProofBranch,
    directions: &[TraverseDirection],
) -> Result<usize, String> {
    let directions_index = directions_to_node_index(directions);

    // If leaf branch has two nodes, directions index assumes it's the right one.
    // We need to check if it's the left one.
    let (left, right) = leaf_branch;
    let left = left.as_ref().ok_or("Left proof node missing")?;
    let right = right.as_ref().ok_or("Right proof node missing")?;

    if &right.0 != NULL_HASH.as_array() && &left.0 == reference {
        Ok(directions_index - 1)
    } else {
        Ok(directions_index)
    }
}

/// Get the tally and hash for parent node of two sibling nodes.
pub fn tally_and_hash(left: &ProofNode, right: &ProofNode) -> Result<ProofNode, String> {
    let mut hasher = Sha256::new();
    let tally = combine_tally(&left.1, &right.1)?;
    tally
        .iter()
        .map(|t| t.to_be_bytes())
        .for_each(|b| hasher.update(b));

    hasher.update(left.0.as_slice());
    hasher.update(right.0.as_slice());
    let hash: [u8; 32] = hasher
        .finalize()
        .as_slice()
        .try_into()
        .expect("Expected hash to be 32 bytes");

    Ok((hash, Some(tally)))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::generate::generate_tree;
    use crate::navigate::find_node_by_vote_reference;
    use crate::proof::votecount::create_proof_of_vote_count;

    #[test]
    fn test_are_neighboors() {
        // ```text
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        // ```
        let a_vote: VoteReference = [0xaa; 32];
        let b_vote: VoteReference = [0xbb; 32];
        let c_vote: VoteReference = [0xcc; 32];

        let root = generate_tree(
            vec![
                (a_vote, vec![1, 0]),
                (b_vote, vec![1, 0]),
                (c_vote, vec![0, 1]),
            ],
            true,
        )
        .unwrap();

        let v = &Validation::Strict;
        let a_proof =
            create_merkle_proof(&find_node_by_vote_reference(&root, &a_vote).unwrap(), v).unwrap();
        let b_proof =
            create_merkle_proof(&find_node_by_vote_reference(&root, &b_vote).unwrap(), v).unwrap();
        let c_proof =
            create_merkle_proof(&find_node_by_vote_reference(&root, &c_vote).unwrap(), v).unwrap();

        assert!(are_neighboors((&a_vote, &a_proof), (&b_vote, &b_proof), v).unwrap());
        assert!(!are_neighboors((&a_vote, &a_proof), (&c_vote, &c_proof), v).unwrap());
        assert!(are_neighboors((&b_vote, &b_proof), (&c_vote, &c_proof), v).unwrap());

        let err = are_neighboors((&a_vote, &a_proof), (&c_vote, &a_proof), v);
        assert_eq!(
            err,
            Err("Right vote reference does not belong to proof".to_string())
        );
    }

    #[test]
    fn test_create_merkle_proof() {
        // ```text
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        // ```
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xcc; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        let a = root.clone();
        let b = a.as_ref().unwrap().left.clone();
        let c = a.as_ref().unwrap().right.clone();
        let d = b.as_ref().unwrap().left.clone();
        let v1 = d.as_ref().unwrap().left.clone();
        let path = [a, b, d, v1.clone()];
        let proof = create_merkle_proof(&path, &Validation::Strict).unwrap();

        // Proof should be
        // * V1, V2
        // * (None), C
        assert_eq!(proof.branches.len(), 2);
        let (proof_v1, proof_v2) = &proof.branches[0];
        assert_eq!(proof_v1.as_ref().unwrap(), &([0xaa; 32], Some(vec![1, 0])));
        assert_eq!(proof_v2.as_ref().unwrap(), &([0xbb; 32], Some(vec![1, 0])));

        let (left, proof_c) = &proof.branches[1];
        assert!(left.is_none());
        let (c_hash, c_tally) = hash_node_ref(&c, &Validation::Strict).unwrap().unwrap();
        assert_eq!(proof_c, &Some((*c_hash.as_array(), c_tally)));
    }

    #[test]
    fn test_collapse_branches_directions() {
        // ```text
        //           N
        //          / \
        //         /   \
        //        /     \
        //       N       N
        //      / \     / \
        //     N   N   N   Ø
        //    / \  |\  | \
        //   V1 V  V V V5 Ø
        // ```
        //
        // Collapsing the merkle proof for V5, Ø should give directions "Right, Left"
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![0, 1]),
                ([0xcc; 32], vec![0, 1]),
                ([0xdd; 32], vec![1, 0]),
                ([0xee; 32], vec![1, 0]),
            ],
            true,
        )
        .unwrap();
        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();

        let (_, directions) = collapse_branches(
            proof.branches,
            &CollapseBranchConstrains::RejectLeftPath,
            &Validation::Strict,
        )
        .unwrap();

        assert_eq!(
            directions,
            vec![
                TraverseDirection::Right,
                TraverseDirection::Left,
                TraverseDirection::Left
            ]
        );
    }

    #[test]
    fn test_directions_to_node_index() {
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![0, 1]),
                ([0xcc; 32], vec![0, 1]),
                ([0xdd; 32], vec![1, 0]),
                ([0xee; 32], vec![1, 0]),
            ],
            true,
        )
        .unwrap();
        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();

        let (_, directions) = collapse_branches(
            proof.branches,
            &CollapseBranchConstrains::RejectLeftPath,
            &Validation::Strict,
        )
        .unwrap();

        assert_eq!(directions_to_node_index(&directions), 4);
    }

    #[test]
    fn test_collapse_branches_constraints() {
        // ```text
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        // ```

        let tree = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xcc; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        let v = &Validation::Strict;
        let root_hash_and_tally = hash_node_ref(&tree, v).unwrap().unwrap();
        let expected = (*root_hash_and_tally.0.as_array(), root_hash_and_tally.1);

        let left_most_node = find_node_by_vote_reference(&tree, &[0xaa; 32]).unwrap();
        let middle_node = find_node_by_vote_reference(&tree, &[0xbb; 32]).unwrap();
        let right_most_node = find_node_by_vote_reference(&tree, &[0xcc; 32]).unwrap();

        let left_proof = create_merkle_proof(&left_most_node, v).unwrap();
        let middle_proof = create_merkle_proof(&middle_node, v).unwrap();
        let right_proof = create_merkle_proof(&right_most_node, &v).unwrap();

        // Left
        assert_eq!(
            expected,
            collapse_branches(
                left_proof.branches.clone(),
                &CollapseBranchConstrains::None,
                &v
            )
            .unwrap()
            .0
        );

        assert_eq!(
            Err("Proof is constrained from collapsing to the left".to_string()),
            collapse_branches(
                left_proof.branches.clone(),
                &CollapseBranchConstrains::RejectLeftPath,
                &v
            )
        );

        assert_eq!(
            expected,
            collapse_branches(
                left_proof.branches,
                &CollapseBranchConstrains::RejectRightPath,
                &v
            )
            .unwrap()
            .0
        );

        // Middle
        assert_eq!(
            expected,
            collapse_branches(
                middle_proof.branches.clone(),
                &CollapseBranchConstrains::None,
                &v
            )
            .unwrap()
            .0
        );

        assert_eq!(
            Err("Proof is constrained from collapsing to the left".to_string()),
            collapse_branches(
                middle_proof.branches.clone(),
                &CollapseBranchConstrains::RejectLeftPath,
                &v
            )
        );

        // Although the node is on the right, it's not hang on a parent that is on the right.
        assert_eq!(
            expected,
            collapse_branches(
                middle_proof.branches,
                &CollapseBranchConstrains::RejectRightPath,
                &v
            )
            .unwrap()
            .0
        );

        // Right
        assert_eq!(
            expected,
            collapse_branches(
                right_proof.branches.clone(),
                &CollapseBranchConstrains::None,
                &v
            )
            .unwrap()
            .0
        );

        assert_eq!(
            expected,
            collapse_branches(
                right_proof.branches.clone(),
                &CollapseBranchConstrains::RejectLeftPath,
                &v
            )
            .unwrap()
            .0
        );

        assert_eq!(
            Err("Proof is constrained from collapsing to the right".to_string()),
            collapse_branches(
                right_proof.branches,
                &CollapseBranchConstrains::RejectRightPath,
                &v
            )
        );
    }

    #[test]
    fn test_find_node_index() {
        // ```text
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        // ```
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xcc; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        let v = &Validation::Strict;
        let a_proof =
            create_merkle_proof(&find_node_by_vote_reference(&root, &[0xaa; 32]).unwrap(), v)
                .unwrap();
        let b_proof =
            create_merkle_proof(&find_node_by_vote_reference(&root, &[0xbb; 32]).unwrap(), v)
                .unwrap();
        let c_proof =
            create_merkle_proof(&find_node_by_vote_reference(&root, &[0xcc; 32]).unwrap(), v)
                .unwrap();

        let (_, a_directions) =
            collapse_branches(a_proof.branches.clone(), &CollapseBranchConstrains::None, v)
                .unwrap();
        let (_, b_directions) =
            collapse_branches(b_proof.branches.clone(), &CollapseBranchConstrains::None, v)
                .unwrap();
        let (_, c_directions) =
            collapse_branches(c_proof.branches.clone(), &CollapseBranchConstrains::None, v)
                .unwrap();

        assert_eq!(
            0,
            find_node_index(&[0xaa; 32], &a_proof.branches[0], &a_directions).unwrap()
        );
        assert_eq!(
            1,
            find_node_index(&[0xbb; 32], &b_proof.branches[0], &b_directions).unwrap()
        );
        assert_eq!(
            2,
            find_node_index(&[0xcc; 32], &c_proof.branches[0], &c_directions).unwrap()
        );
    }
}