tallytree 0.3.4

use crate::hash::NodeHash;
use crate::navigate::find_down_right_most_branch;
use crate::node::NodeRef;
use crate::proof::util::{
    collapse_branches, create_merkle_proof, directions_to_node_index, CollapseBranchConstrains,
};
use crate::proof::Proof;
use crate::tally::TallyList;
use crate::Validation;

/// Create a proof for vote count given a root node. Using this proof, it
/// is possible to derive and verify the voter count.
///
/// Hash and tally of the nodes in (parentheses) are included in the proof.
///
/// * `V3` and `Ø` allow for generating the hash of `C`.
/// * `B` and the derived `C` allow generating the hash of `A`.
/// * A is the root node passed into the function.
/// ```text
///            A
///          /  \
///         (B)  C
///        / \   | \
///       D   E  F  (Ø)
///       |   |  |
///       V1  V2 (V3)
/// ```
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::hash::hash_node_ref;
/// use tallytree::proof::{create_proof_of_vote_count, verify_proof_of_vote_count};
/// use tallytree::Validation;
///
/// let root = generate_tree(vec![([0xaa; 32], vec![1, 0])], true).unwrap();
/// let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
/// let (vote_count, root_hash, tally) = verify_proof_of_vote_count(&proof).unwrap();
/// assert_eq!(vote_count, 1);
/// assert_eq!(vec![1, 0], tally);
/// assert_eq!(
///    hash_node_ref(&root, &Validation::Strict).unwrap(),
///    Some((root_hash, Some(tally)))
/// );
/// ```
pub fn create_proof_of_vote_count(node: &NodeRef, v: &Validation) -> Result<Proof, String> {
    let path = find_down_right_most_branch(node)
        .ok_or_else(|| "Invalid/incomplete merkle tally tree".to_string())?;
    create_merkle_proof(&path, v)
}

/// Parses a proof of vote count, returning vote count, merkle root hash and
/// tally.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::hash::hash_node_ref;
/// use tallytree::proof::{create_proof_of_vote_count, verify_proof_of_vote_count};
/// use tallytree::Validation;
///
/// let root = generate_tree(vec![([0xaa; 32], vec![1, 0])], true).unwrap();
/// let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
/// let (vote_count, root_hash, tally) = verify_proof_of_vote_count(&proof).unwrap();
/// assert_eq!(vote_count, 1);
/// assert_eq!(vec![1, 0], tally);
/// assert_eq!(
///    hash_node_ref(&root, &Validation::Strict).unwrap(),
///    Some((root_hash, Some(tally)))
/// );
/// ```
pub fn verify_proof_of_vote_count(proof: &Proof) -> Result<(usize, NodeHash, TallyList), String> {
    if proof.branches.is_empty() {
        return Err("Invalid proof, proof is empty.".to_string());
    }

    let (node, directions) = collapse_branches(
        proof.branches.clone(),
        &CollapseBranchConstrains::RejectLeftPath,
        &Validation::Strict,
    )?;
    let hash = node.0;
    let tally = node.1.unwrap();

    let tally_count = tally.iter().sum::<u32>();
    let vote_count = directions_to_node_index(&directions) + 1;

    if vote_count != tally_count as usize {
        return Err(format!(
            "Number of votes and tally mismatch ({} votes != {} tallied)",
            vote_count, tally_count
        ));
    }

    Ok((vote_count, NodeHash(hash), tally))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::generate::generate_tree;
    use crate::hash::{hash_node_ref, NULL_HASH};
    use crate::node::null_node;
    use crate::Validation;
    use std::sync::Arc;
    #[cfg(feature = "with-benchmarks")]
    use test::Bencher;

    #[test]
    fn proof_of_vote_count() {
        let node: NodeRef = None;
        let proof = create_proof_of_vote_count(&node, &Validation::Strict);
        assert!(proof.is_err());

        let node = Some(Arc::new(null_node()));
        let proof = create_proof_of_vote_count(&node, &Validation::Strict);
        assert!(proof.is_err());

        // ```text
        // Tree with one Ø-node
        //            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 proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
        // The proof should consist of the following branches:
        // - V3, Ø
        // - B, None
        assert_eq!(proof.branches.len(), 2);

        let (v3, o) = &proof.branches[0];
        assert_eq!(v3.as_ref().unwrap(), &([0xcc; 32], Some(vec![0, 1])));
        assert_eq!(v3.as_ref().unwrap().1, Some(vec![0, 1]));

        assert_eq!(&o.as_ref().unwrap().0, NULL_HASH.as_array());
        assert!(o.as_ref().unwrap().1.is_none());

        let (b, none) = &proof.branches[1];
        let (b_hash, b_tally) = hash_node_ref(&root.as_ref().unwrap().left, &Validation::Strict)
            .unwrap()
            .unwrap();
        assert_eq!(b, &Some((*b_hash.as_array(), b_tally)));
        assert_eq!(b.as_ref().unwrap().1, Some(vec![2, 0]));
        assert!(none.is_none());

        // ```text
        // Tree with no Ø-nodes
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  G
        //       |   |  |  |
        //       V1  V2 V3 V4
        // ```
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xcc; 32], vec![0, 1]),
                ([0xdd; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
        // The proof should consist of the following branches:
        // - V3, V4
        // - B, None
        assert_eq!(proof.branches.len(), 2);
        let (_v3, v4) = &proof.branches[0];
        assert_eq!(v4.as_ref().unwrap(), &([0xdd; 32], Some(vec![0, 1])));
    }

    #[test]
    fn proof_of_vote_count_with_9_votes() {
        // ```text
        // Tree with 9 leafs and 3 Ø-nodes
        //            A
        //          /  \
        //         /    \
        //        B      C
        //       / \     | \
        //      /   \    \  \
        //     D     E    F  Ø
        //    / \   / \   | \
        //   G   H I   J  K  Ø
        //  /| / | |\  |\ \ \
        // L M N O P Q R S T Ø
        // | | | | | | | | |
        // V V V V V V V V V9
        // ```
        let root = generate_tree(
            vec![
                ([0x11; 32], vec![1, 0]),
                ([0x22; 32], vec![1, 0]),
                ([0x33; 32], vec![0, 1]),
                ([0x44; 32], vec![1, 0]),
                ([0x55; 32], vec![1, 0]),
                ([0x66; 32], vec![0, 1]),
                ([0x77; 32], vec![1, 0]),
                ([0x88; 32], vec![1, 0]),
                ([0x99; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();

        // The proof should consist of the following branches:
        // - V9, Ø
        // - None, Ø
        // - None, Ø
        // - B, None
        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
        assert_eq!(proof.branches.len(), 4);

        let (v9, o) = &proof.branches[0];
        assert_eq!(v9.as_ref().unwrap(), &([0x99; 32], Some(vec![0, 1])));
        assert_eq!(o.as_ref().unwrap().0, *NULL_HASH.as_array());
        assert_eq!(o.as_ref().unwrap().1, None);

        let (left, o) = &proof.branches[1];
        assert!(left.is_none());
        assert_eq!(o.as_ref().unwrap().0, *NULL_HASH.as_array());
        assert_eq!(o.as_ref().unwrap().1, None);

        let (left, o) = &proof.branches[2];
        assert!(left.is_none());
        assert_eq!(o.as_ref().unwrap().0, *NULL_HASH.as_array());
        assert_eq!(o.as_ref().unwrap().1, None);

        let (b, right) = &proof.branches[3];
        let (b_hash, b_tally) = hash_node_ref(&root.as_ref().unwrap().left, &Validation::Strict)
            .unwrap()
            .unwrap();
        assert!(right.is_none());
        assert_eq!(b, &Some((*b_hash.as_array(), b_tally)));
        assert_eq!(b.as_ref().unwrap().1, Some(vec![6, 2]));
    }

    #[test]
    fn test_verify_proof_of_vote_count_1_vote() {
        // ```text
        // Proof of a tree with 1 vote.
        //          N
        //        /   \
        //       N     Ø
        //       |
        //       V
        // ```
        let root = generate_tree(vec![([0xaa; 32], vec![1, 0])], true).unwrap();
        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
        let (vote_count, root_hash, tally) = verify_proof_of_vote_count(&proof).unwrap();
        assert_eq!(vote_count, 1);
        assert_eq!(vec![1, 0], tally);
        assert_eq!(
            hash_node_ref(&root, &Validation::Strict).unwrap(),
            Some((root_hash, Some(tally)))
        );
    }

    #[test]
    fn test_verify_proof_of_vote_count_2_votes() {
        // ```text
        // Proof of a tree with 2 votes and 0 Ø-nodes.
        //          N
        //        /   \
        //       N     N
        //       |     |
        //       V1    V2
        // ```
        let root = generate_tree(
            vec![([0xaa; 32], vec![1, 0]), ([0xbb; 32], vec![0, 1])],
            true,
        )
        .unwrap();
        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
        let (vote_count, root_hash, tally) = verify_proof_of_vote_count(&proof).unwrap();
        assert_eq!(vote_count, 2);
        assert_eq!(vec![1, 1], tally);
        assert_eq!(
            hash_node_ref(&root, &Validation::Strict).unwrap(),
            Some((root_hash, Some(tally)))
        );
    }

    #[test]
    fn test_verify_proof_of_vote_count_3_votes() {
        // ```text
        // Proof of a tree with 2 votes and 1 Ø-node.
        //            N
        //          /  \
        //         N    N
        //        / \   | \
        //       N   N  N  Ø
        //       |   |  |
        //       V1  V2 V3
        // ```
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![0, 1]),
                ([0xcc; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        let proof = create_proof_of_vote_count(&root, &Validation::Strict).unwrap();
        let (vote_count, root_hash, tally) = verify_proof_of_vote_count(&proof).unwrap();
        assert_eq!(vote_count, 3);
        assert_eq!(vec![1, 2], tally);
        assert_eq!(
            hash_node_ref(&root, &Validation::Strict).unwrap(),
            Some((root_hash, Some(tally)))
        );
    }

    #[test]
    fn test_verify_proof_of_vote_count_4_votes() {
        // ```text
        // Proof of a tree with 5 votes and 2 Ø-node.
        //           N
        //          / \
        //         /   \
        //        /     \
        //       N       N
        //      / \     / \
        //     N   N   N   Ø
        //    / \  |\  | \
        //   V  V  V V V  Ø
        // ```
        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 (vote_count, root_hash, tally) = verify_proof_of_vote_count(&proof).unwrap();
        assert_eq!(vote_count, 5);
        assert_eq!(vec![3, 2], tally);
        assert_eq!(
            hash_node_ref(&root, &Validation::Strict).unwrap(),
            Some((root_hash, Some(tally)))
        );
    }

    #[test]
    fn test_verify_proof_with_invalid_tally() {
        // ```text
        //          N
        //        /   \
        //       N     N
        //       |     |
        //       V1    V2
        // ```
        // Let V1 cast two votes
        let root = generate_tree(
            vec![([0xaa; 32], vec![2, 0]), ([0xbb; 32], vec![0, 1])],
            true,
        )
        .unwrap();
        let proof = create_proof_of_vote_count(&root, &Validation::Relaxed).unwrap();

        let verified = verify_proof_of_vote_count(&proof);
        assert_eq!(
            verified.err(),
            Some("Invalid proof. Left node should have exactly one vote.".to_string())
        );
    }

    #[test]
    fn test_verify_proof_with_invalid_inner_tally() {
        // ```text
        //            N
        //          /  \
        //         L    N
        //        / \   | \
        //       N   N  N  Ø
        //       |   |  |
        //       V1  V2 V3
        //
        // Node L has too many votes. This should be detected.
        // ```
        let proof = Proof {
            branches: vec![
                // Nodes V3 and Ø
                (
                    Some(([0xaa; 32], Some(vec![1, 0]))),
                    Some((*NULL_HASH.as_array(), None)),
                ),
                // Node L
                (Some(([0xbb; 32], Some(vec![0, 42]))), None),
            ],
        };
        let verified = verify_proof_of_vote_count(&proof);
        let expected_error =
            "Number of votes and tally mismatch (3 votes != 43 tallied)".to_string();
        assert_eq!(verified.err(), Some(expected_error));
    }

    #[test]
    fn test_verify_proof_of_vote_wrong_path() {
        // ```text
        //            N
        //          /  \
        //         L    N
        //        / \   | \
        //       N   N  N  Ø
        //       |   |  |
        //       V1  V2 V3
        //
        // We create a proof for V1, V2. But we should have taken the
        // right-most path and created a proof
        // for (V3, Ø)
        // ```
        let proof = Proof {
            branches: vec![
                // Nodes V1 and V2
                (
                    Some(([0xaa; 32], Some(vec![1, 0]))),
                    Some(([0xbb; 32], Some(vec![0, 1]))),
                ),
                // Node N
                (None, Some(([0xcc; 32], Some(vec![1, 0])))),
            ],
        };
        let verified = verify_proof_of_vote_count(&proof);
        assert_eq!(
            verified.err(),
            Some("Proof is constrained from collapsing to the left".to_string())
        );
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_vote_count_proof_1k_nocache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(1000);
        let tree = generate_tree(votes, false).unwrap();
        b.iter(|| create_proof_of_vote_count(&tree, &Validation::Relaxed))
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_vote_count_proof_1k_cache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(1000);
        let tree = generate_tree(votes, true).unwrap();
        b.iter(|| create_proof_of_vote_count(&tree, &Validation::Relaxed))
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_vote_count_proof_5k_nocache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(5000);
        let tree = generate_tree(votes, false).unwrap();
        b.iter(|| create_proof_of_vote_count(&tree, &Validation::Relaxed))
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_create_vote_count_proof_5k_cache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(5000);
        let tree = generate_tree(votes, true).unwrap();
        b.iter(|| create_proof_of_vote_count(&tree, &Validation::Relaxed))
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_verify_vote_count_proof_10k(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(10_000);
        let tree = generate_tree(votes, true).unwrap();
        let proof = create_proof_of_vote_count(&tree, &Validation::Relaxed).unwrap();
        b.iter(|| verify_proof_of_vote_count(&proof))
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_verify_vote_count_proof_100k(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(100_000);
        let tree = generate_tree(votes, true).unwrap();
        let proof = create_proof_of_vote_count(&tree, &Validation::Relaxed).unwrap();
        b.iter(|| verify_proof_of_vote_count(&proof))
    }
}