tallytree 0.3.4

use crate::hash::NodeHash;
use crate::navigate::{find_closest_siblings, find_node_by_vote_reference};
use crate::node::NodeRef;
use crate::proof::util::{
    are_neighboors, collapse_branches, create_merkle_proof, CollapseBranchConstrains,
};
use crate::proof::{Proof, ProofBranch};
use crate::tally::{has_one_vote, TallyList};
use crate::{Validation, VoteReference};

/**
 * A proof that a node exists in a merkle tally tree.
 */
pub type InclusionProof = Option<Proof>;

/**
 * A proof that a node does not exist in a merkle tally tree.
 */
pub type ExclusionProof = Option<(Option<Proof>, Option<Proof>)>;

/// Create proof of a vote references' existence or non-existance in a merkle
/// tree.
///
/// Return either inclusion proof or exclusion proof.
///
/// Inclusion proof includes the hash of the nodes needed to calculate the
/// merkle tree root hash from the votes position in the tree, including
/// the vote itself.
///
/// Exclusion proof includes the hash of the nodes needed to calculate the
/// merkle tree root hash from both neighboring nodes of where the vote
/// would have been positioned in the tree.
///
/// ```text
///            A
///          /  \
///         B    C
///        / \   | \
///       D   E  F  Ø
///       |   |  |
///       V1  V2 V4
/// ```
///
/// In the above tree, the inclusion proof for vote V4 would be:
/// - V4, Ø
/// - B, None
///
/// The other nodes in the path, F, C and A, can be derived from the nodes
/// provided.
///
/// In the above tree, the exclusion proof for non-existant vote V3 would be
/// - V1, V2
/// - None, C
///
/// ... proving the existence of V2 and
///
/// - V4, Ø
/// - B, None
///
/// ... proving the existence of V4. This is enough to prove that V3 does not
/// exist in the tree.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::proof::inclusion::create_inclusion_exclusion_proof;
/// use tallytree::Validation;
///
/// let tree = generate_tree(vec![
///   ([0xaa; 32], vec![1, 0]),
///   ([0xcc; 32], vec![1, 0]),
/// ], true).unwrap();
/// let (inclusion, _) = create_inclusion_exclusion_proof(&tree, &[0xaa; 32], &Validation::Strict).unwrap();
/// let (_, exclusion) = create_inclusion_exclusion_proof(&tree, &[0xbb; 32], &Validation::Strict).unwrap();
/// ```
pub fn create_inclusion_exclusion_proof(
    root: &NodeRef,
    vote_reference: &VoteReference,
    v: &Validation,
) -> Result<(InclusionProof, ExclusionProof), String> {
    let ((left, right), exists) = find_closest_siblings(root, vote_reference)?;
    if exists {
        // TODO: find_closest_siblings could provide the path, rather than
        // iterating the tree again to find it.
        let path = find_node_by_vote_reference(root, vote_reference).unwrap();
        Ok((Some(create_merkle_proof(&path, v)?), None))
    } else {
        let left_proof = if let Some(l) = left {
            let path = find_node_by_vote_reference(root, &l).unwrap();
            Some(create_merkle_proof(&path, v)?)
        } else {
            None
        };
        let right_proof = if let Some(r) = right {
            let path = find_node_by_vote_reference(root, &r).unwrap();
            Some(create_merkle_proof(&path, v)?)
        } else {
            None
        };
        Ok((None, Some((left_proof, right_proof))))
    }
}

/// Verify that a inclusion proof is valid for a vote reference.
///
/// Returns:
/// - Merkle root hash of merkle proof
/// - Tally of the entire tree
/// - Tally for given vote reference
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::proof::inclusion::{create_inclusion_exclusion_proof, verify_inclusion_proof};
/// use tallytree::Validation;
/// use tallytree::hash::hash_node_ref;
///
/// let tree = generate_tree(vec![
///   ([0xaa; 32], vec![1, 0]),
///   ([0xcc; 32], vec![1, 0]),
/// ], true).unwrap();
/// let v = &Validation::Strict;
/// let (inclusion, _) = create_inclusion_exclusion_proof(&tree, &[0xaa; 32], v).unwrap();
/// let (merkle_tree_hash, _, _) = verify_inclusion_proof(&inclusion, &[0xaa; 32], v).unwrap();
/// assert_eq!(merkle_tree_hash, hash_node_ref(&tree, v).unwrap().unwrap().0);
/// ```
pub fn verify_inclusion_proof(
    proof: &InclusionProof,
    vote_reference: &VoteReference,
    v: &Validation,
) -> Result<(NodeHash, TallyList, TallyList), String> {
    let null_hash = [0x00; 32];
    if vote_reference == &null_hash {
        return Err("Vote reference cannot be null".to_string());
    }
    let proof = proof.as_ref().ok_or("Proof is None")?;

    if proof.branches.is_empty() {
        return Err("Empty proof".to_string());
    }

    // Lowest branch should contain the included vote reference.
    let (left, right) = &proof.branches[0];

    let mut vote_cast = None;
    if let Some(l) = left {
        if l.0 == *vote_reference {
            vote_cast = l.1.clone()
        }
    }
    if let Some(r) = right {
        if r.0 == *vote_reference {
            vote_cast = r.1.clone()
        }
    }
    let vote_cast = vote_cast.ok_or("Proof does not contain vote")?;
    if !matches!(v, Validation::Relaxed) && !has_one_vote(&vote_cast) {
        return Err("Invalid proof. Vote reference should cast exactly one vote.".to_string());
    }
    let (root, _) = collapse_branches(proof.branches.clone(), &CollapseBranchConstrains::None, v)?;
    Ok((
        NodeHash(root.0),
        root.1.ok_or("Inclusion proof contains no tally")?,
        vote_cast,
    ))
}

/// Verify that an exclusion proof is valid for a vote reference. Returns error
/// if proof fails to validate.
///
/// Returns:
/// - Merkle root hash of merkle proof
/// - Tally of the entire tree
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::proof::inclusion::{create_inclusion_exclusion_proof, verify_exclusion_proof};
/// use tallytree::Validation;
/// use tallytree::hash::hash_node_ref;
///
/// let tree = generate_tree(vec![
///   ([0xaa; 32], vec![1, 0]),
///   ([0xcc; 32], vec![1, 0]),
/// ], true).unwrap();
/// let v = &Validation::Strict;
/// let (_, exclusion) = create_inclusion_exclusion_proof(&tree, &[0xbb; 32], v).unwrap();
/// let (merkle_tree_hash, _) = verify_exclusion_proof(&exclusion, &[0xbb; 32], v).unwrap();
/// assert_eq!(merkle_tree_hash, hash_node_ref(&tree, v).unwrap().unwrap().0);
/// ```
pub fn verify_exclusion_proof(
    proof: &ExclusionProof,
    vote: &VoteReference,
    v: &Validation,
) -> Result<(NodeHash, TallyList), String> {
    let proof = proof.as_ref().ok_or("Invalid proof (None)")?;

    match proof {
        (None, None) => Err("Invalid proof (Both trees are None)".to_string()),
        (Some(proof_left), None) => {
            // Must be the right-most node. Cannot walk left, and the
            // right-most leaf must be less than the vote.
            let run_check = |hash| {
                if hash == vote {
                    return Err("Invalid exclusion proof. Node is in proof.".to_string());
                }
                if hash > vote {
                    return Err(
                        "Invalid exclusion proof. Right-most node is larger than vote.".to_string(),
                    );
                }
                Ok(())
            };
            let leaf_branch: &ProofBranch = proof_left.branches.get(0).ok_or("Invalid proof")?;
            match leaf_branch {
                (Some((left_hash, _)), Some(_)) => {
                    run_check(left_hash)?;
                }
                (Some((left_hash, _)), None) => {
                    run_check(left_hash)?;
                }
                _ => return Err("Invalid branch in proof".to_string()),
            }

            let (merkle_root, _) = collapse_branches(
                proof_left.branches.clone(),
                &CollapseBranchConstrains::RejectLeftPath,
                v,
            )?;

            Ok((
                NodeHash(merkle_root.0),
                merkle_root.1.ok_or("Invalid proof, tally missing.")?,
            ))
        }
        (None, Some(proof_right)) => {
            // Must be the left-most node. Cannot walk right, and the
            // left-most leaf must be larger than the vote.

            let leaf_branch: &ProofBranch = proof_right.branches.get(0).ok_or("Invalid proof")?;
            let run_check = |hash| {
                if hash == vote {
                    return Err("Invalid exclusion proof. Node is in proof.".to_string());
                }
                if hash < vote {
                    return Err(
                        "Invalid exclusion proof. Left-most node is less than vote.".to_string()
                    );
                }
                Ok(())
            };
            match leaf_branch {
                (Some(_left), Some((right_hash, _))) => {
                    run_check(right_hash)?;
                }
                (Some((left_hash, _)), None) => {
                    run_check(left_hash)?;
                }
                _ => return Err("Invalid branch in proof".to_string()),
            }

            let (merkle_root, _) = collapse_branches(
                proof_right.branches.clone(),
                &CollapseBranchConstrains::RejectRightPath,
                v,
            )?;

            Ok((
                NodeHash(merkle_root.0),
                merkle_root.1.ok_or("Invalid proof, tally missing")?,
            ))
        }
        (Some(proof_left), Some(proof_right)) => {
            let leaf_branch_left = proof_left.branches.get(0).ok_or("Empty proof")?;
            let leaf_branch_right = proof_right.branches.get(0).ok_or("Empty proof")?;

            let (left_node, right_node) = if leaf_branch_left == leaf_branch_right {
                // The two would-be neigbours to excluded vote have the same parent.
                let (l, r) = leaf_branch_left;
                (l, r)
            } else {
                // The two would-be neigbours to excluded vote have different parents.
                let (_, left_node) = leaf_branch_left;
                let (right_node, _) = leaf_branch_right;
                (left_node, right_node)
            };
            let left_node = left_node
                .as_ref()
                .ok_or("Expected node if left proof missing")?;
            let right_node = right_node
                .as_ref()
                .ok_or("Expected node if right proof missing")?;

            if &left_node.0 >= vote {
                return Err("Invalid proof. Node to left is not less than vote.".to_string());
            }
            if vote >= &right_node.0 {
                return Err("Invalid proof. Node to right is not greater than vote".to_string());
            }

            if !are_neighboors((&left_node.0, proof_left), (&right_node.0, proof_right), v)? {
                return Err("Invalid proof. Left and right proof are not neighboors.".to_string());
            }
            let (merkle_root, _) = collapse_branches(
                proof_left.branches.clone(),
                &CollapseBranchConstrains::None,
                v,
            )?;
            Ok((
                NodeHash(merkle_root.0),
                merkle_root.1.ok_or("Invalid proof, tally missing")?,
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::generate::generate_tree;
    use crate::hash::hash_node_ref;
    #[cfg(feature = "with-benchmarks")]
    use crate::utiltest::*;
    #[cfg(feature = "with-benchmarks")]
    use test::Bencher;

    #[test]
    fn test_create_inclusion_proof() {
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xdd; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        let (inclusion, exclusion) =
            create_inclusion_exclusion_proof(&root, &[0xaa; 32], &Validation::Strict).unwrap();
        assert!(exclusion.is_none());
        let inclusion = inclusion.unwrap();
        assert_eq!(2, inclusion.branches.len());
        let (v1, v2) = &inclusion.branches[0];
        assert_eq!(v1, &Some(([0xaa; 32], Some(vec![1, 0]))));
        assert_eq!(v2, &Some(([0xbb; 32], Some(vec![1, 0]))));
        let (b, c) = &inclusion.branches[1];
        assert!(b.is_none());
        let (c_hash, c_tally) = hash_node_ref(&root.as_ref().unwrap().right, &Validation::Strict)
            .unwrap()
            .unwrap();
        assert_eq!(c, &Some((*c_hash.as_array(), c_tally)));
    }

    #[test]
    fn test_create_exclusion_proof() {
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xdd; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        // Voter 0xcc does not exist. Should return merkle tree to 0xaa and 0xdd to prove this.
        let (inclusion, exclusion) =
            create_inclusion_exclusion_proof(&root, &[0xcc; 32], &Validation::Strict).unwrap();
        assert!(inclusion.is_none());

        let (exclusion_left, exclusion_right) = exclusion.unwrap();
        assert_eq!(&2, &exclusion_left.as_ref().unwrap().branches.len());
        assert_eq!(&2, &exclusion_right.as_ref().unwrap().branches.len());
        let (_, v2) = &exclusion_left.unwrap().branches[0];
        let (v3, _) = &exclusion_right.unwrap().branches[0];
        assert_eq!(v2.as_ref().unwrap(), &([0xbb; 32], Some(vec![1, 0])));
        assert_eq!(v3.as_ref().unwrap(), &([0xdd; 32], Some(vec![0, 1])));
    }

    #[test]
    fn test_verify_valid_inclusion_proof() {
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xdd; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();

        let v = &Validation::Strict;
        let (inclusion, _) = create_inclusion_exclusion_proof(&root, &[0xdd; 32], v).unwrap();

        let (merkle_root_hash, final_tally, vote_cast) =
            verify_inclusion_proof(&inclusion, &[0xdd; 32], v).unwrap();

        assert_eq!(
            merkle_root_hash,
            hash_node_ref(&root, v).unwrap().unwrap().0
        );
        assert_eq!(final_tally, vec![2, 1]);
        assert_eq!(vote_cast, vec![0, 1]);
    }

    #[test]
    fn test_verify_invalid_inclusion_proof() {
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xdd; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();

        let v = &Validation::Strict;
        let (inclusion, _) = create_inclusion_exclusion_proof(&root, &[0xdd; 32], v).unwrap();

        // 0xaa is not in this proof
        let err = verify_inclusion_proof(&inclusion, &[0xaa; 32], v);
        assert_eq!(err, Err("Proof does not contain vote".to_string()));

        // Empty proofs
        let err = verify_inclusion_proof(&None, &[0xdd; 32], v);
        assert_eq!(err, Err("Proof is None".to_string()));

        let err = verify_inclusion_proof(&Some(Proof::default()), &[0xaa; 32], v);
        assert_eq!(err, Err("Empty proof".to_string()));
    }

    #[test]
    fn test_verify_valid_exclusion_proof() {
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        let root = generate_tree(
            vec![
                ([0xbb; 32], vec![1, 0]),
                ([0xcc; 32], vec![1, 0]),
                ([0xee; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        let v = &Validation::Strict;
        let merkle_hash_and_tally = hash_node_ref(&root, v).unwrap().unwrap();
        let expected = (merkle_hash_and_tally.0, merkle_hash_and_tally.1.unwrap());

        // Exclusion proof where the excluded vote is LEFT OF V1
        let (_, exclusion) = create_inclusion_exclusion_proof(&root, &[0xaa; 32], v).unwrap();
        let merkle_root_hash = verify_exclusion_proof(&exclusion, &[0xaa; 32], v).unwrap();
        assert_eq!(expected, merkle_root_hash);

        // Exclusion proof where the excluded vote is BETWEEN V2 and V3
        let (_, exclusion) = create_inclusion_exclusion_proof(&root, &[0xdd; 32], v).unwrap();
        let merkle_root_hash = verify_exclusion_proof(&exclusion, &[0xdd; 32], v).unwrap();
        assert_eq!(expected, merkle_root_hash);

        // Exclusion proof where the excluded vote is RIGHT OF V3
        let (_, exclusion) = create_inclusion_exclusion_proof(&root, &[0xff; 32], v).unwrap();
        let merkle_root_hash = verify_exclusion_proof(&exclusion, &[0xff; 32], v).unwrap();
        assert_eq!(expected, merkle_root_hash);
    }

    #[test]
    fn test_verify_valid_exclusion_proof_2() {
        // Test case with only two votes.
        let tree = generate_tree(
            vec![([0xaa; 32], vec![1, 0]), ([0xcc; 32], vec![1, 0])],
            true,
        )
        .unwrap();
        let v = &Validation::Strict;
        let (_, exclusion) = create_inclusion_exclusion_proof(&tree, &[0xbb; 32], v).unwrap();
        let (merkle_tree_hash, _) = verify_exclusion_proof(&exclusion, &[0xbb; 32], v).unwrap();
        assert_eq!(
            merkle_tree_hash,
            hash_node_ref(&tree, v).unwrap().unwrap().0
        );
    }

    #[test]
    fn test_verify_invalid_exclusion_proof() {
        let v = &Validation::Strict;
        let err = verify_exclusion_proof(&Some((None, None)), &[0xdd; 32], v);
        assert_eq!(err, Err("Invalid proof (Both trees are None)".to_string()))
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_inclusion_proof_1k_nocache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(1000);
        let tree = generate_tree(votes, false).unwrap();
        let vote: VoteReference = dummy_hash_from_number(42);
        assert!(
            create_inclusion_exclusion_proof(&tree, &vote, &Validation::Relaxed)
                .unwrap()
                .0
                .is_some()
        );
        b.iter(|| create_inclusion_exclusion_proof(&tree, &vote, &Validation::Relaxed).unwrap())
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_inclusion_proof_1k_cache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(1000);
        let tree = generate_tree(votes, true).unwrap();
        let vote: VoteReference = dummy_hash_from_number(42);
        assert!(
            create_inclusion_exclusion_proof(&tree, &vote, &Validation::Relaxed)
                .unwrap()
                .0
                .is_some()
        );
        b.iter(|| create_inclusion_exclusion_proof(&tree, &vote, &Validation::Relaxed).unwrap())
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_inclusion_proof_5k_nocache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(5000);
        let tree = generate_tree(votes, false).unwrap();
        let vote: VoteReference = dummy_hash_from_number(42);
        assert!(
            create_inclusion_exclusion_proof(&tree, &vote, &Validation::Strict)
                .unwrap()
                .0
                .is_some()
        );
        b.iter(|| create_inclusion_exclusion_proof(&tree, &vote, &Validation::Strict).unwrap())
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_inclusion_proof_5k_cache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(5000);
        let tree = generate_tree(votes, true).unwrap();
        let vote: VoteReference = dummy_hash_from_number(42);
        assert!(
            create_inclusion_exclusion_proof(&tree, &vote, &Validation::Strict)
                .unwrap()
                .0
                .is_some()
        );
        b.iter(|| create_inclusion_exclusion_proof(&tree, &vote, &Validation::Strict).unwrap())
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_exclusion_proof_1k_nocache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(1000);
        let tree = generate_tree(votes, false).unwrap();
        let excluded_vote: VoteReference = dummy_hash_from_number(43);
        assert!(
            create_inclusion_exclusion_proof(&tree, &excluded_vote, &Validation::Strict)
                .unwrap()
                .1
                .is_some()
        );
        b.iter(|| {
            create_inclusion_exclusion_proof(&tree, &excluded_vote, &Validation::Strict).unwrap()
        })
    }
    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_exclusion_proof_1k_cache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(1000);
        let tree = generate_tree(votes, true).unwrap();
        let excluded_vote: VoteReference = dummy_hash_from_number(43);
        assert!(
            create_inclusion_exclusion_proof(&tree, &excluded_vote, &Validation::Strict)
                .unwrap()
                .1
                .is_some()
        );
        b.iter(|| {
            create_inclusion_exclusion_proof(&tree, &excluded_vote, &Validation::Strict).unwrap()
        })
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_verify_inclusion_proof_100k(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(100_000);
        let tree = generate_tree(votes, true).unwrap();
        let vote: VoteReference = dummy_hash_from_number(42);
        let proof = create_inclusion_exclusion_proof(&tree, &vote, &Validation::Strict).unwrap();

        b.iter(|| verify_inclusion_proof(&proof.0, &vote, &Validation::Strict).unwrap())
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_verify_exclusion_proof_100k(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(100_000);
        let tree = generate_tree(votes, true).unwrap();
        let excluded_vote: VoteReference = dummy_hash_from_number(43);
        let proof =
            create_inclusion_exclusion_proof(&tree, &excluded_vote, &Validation::Strict).unwrap();

        b.iter(|| verify_exclusion_proof(&proof.1, &excluded_vote, &Validation::Strict).unwrap())
    }
}