extern crate tallytree;
use tallytree::generate::generate_tree;
use tallytree::hash::hash_node_ref;
use tallytree::proof::{
    create_inclusion_exclusion_proof, create_proof_of_vote_count, verify_exclusion_proof,
    verify_inclusion_proof, verify_proof_of_vote_count,
};
use tallytree::Validation;

fn main() {
    // A vote with 3 voters where:
    //
    // - 0xaa votes for the first option
    // - 0xcc votes for the first option
    // - 0xdd votes for the second option.
    let tree = generate_tree(
        vec![
            ([0xaa; 32], vec![1, 0]),
            ([0xcc; 32], vec![1, 0]),
            ([0xdd; 32], vec![0, 1]),
        ],
        true,
    )
    .unwrap();
    let v = &Validation::Strict;
    let merkle_root_hash = hash_node_ref(&tree, v).unwrap().unwrap().0;

    // Prove that 0xaa's vote was tallied.
    let (inclusion, _) = create_inclusion_exclusion_proof(&tree, &[0xaa; 32], v).unwrap();
    let (proof_hash, _, vote) = verify_inclusion_proof(&inclusion, &[0xaa; 32], v).unwrap();
    assert_eq!(proof_hash, merkle_root_hash);
    println!("Voter 0xaa has provably voted {:?}.", vote);

    // Proof that 0xbb did not cast a vote.
    let (_, exclusion) = create_inclusion_exclusion_proof(&tree, &[0xbb; 32], v).unwrap();
    let (proof_hash, _) = verify_exclusion_proof(&exclusion, &[0xbb; 32], v).unwrap();
    assert_eq!(proof_hash, merkle_root_hash);
    println!("Voter 0xbb has provably NOT voted.");

    // Proof that there are votes in the merkle tally tree.
    let votes_proof = create_proof_of_vote_count(&tree, v).unwrap();
    let (votes, proof_hash, _) = verify_proof_of_vote_count(&votes_proof).unwrap();
    assert_eq!(proof_hash, merkle_root_hash);
    assert_eq!(votes, 3);
    println!("There are provably 3 votes in the merkle tally tree.");
}