tallytree 0.3.4

use crate::node::{null_node, Node, NodeRef};
use crate::tally::TallyList;
use crate::utilstring::to_hex;
use crate::VoteReference;
use std::collections::VecDeque;
use std::sync::Arc;

/// Build a 'level' in the merkle tally tree with given nodes. Returns the new parent nodes.
fn build_level(mut queue: VecDeque<Node>, nodes_use_cache: bool) -> Result<Vec<Node>, String> {
    assert!(queue.len() % 2 == 0);
    let mut nodes: Vec<Node> = vec![];
    loop {
        let left = queue.pop_front();
        if left.is_none() {
            return Ok(nodes);
        }
        let right = queue.pop_front();
        nodes.push(Node::new(
            None,
            Some(Arc::new(left.unwrap())),
            Some(Arc::new(right.unwrap())),
            nodes_use_cache,
        )?);
    }
}

/// Generate a merkle tally tree from votes.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// // Generate a tree where vote reference where:
/// // - 0xaa votes for option 0,
/// // - 0xbb votes for option 1,
/// // - 0xcc votes for option 0
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
///     ([0xcc; 32], vec![1, 0]),
/// ], false);
/// ```
///
/// The above example results in a tree like this:
/// ```text
///            A
///          /  \
///         B    C
///        / \   | \
///       D   E  F  Ø
///       |   |  |
///       aa  bb cc
/// ```
pub fn generate_tree(
    mut votes: Vec<(VoteReference, TallyList)>,
    nodes_use_cache: bool,
) -> Result<NodeRef, String> {
    if votes.is_empty() {
        return Ok(None);
    }

    votes.sort_by(|a, b| a.partial_cmp(b).unwrap());

    // Check for duplicates
    let mut last: Option<&VoteReference> = None;
    for (v, _) in &votes {
        if last == Some(v) {
            return Err(format!(
                "Cannot have duplicate vote references (found duplciate of '{}')",
                to_hex(v)?
            ));
        }
        last = Some(v);
    }

    let leafs = votes
        .into_iter()
        .map(|v| Node::new(Some(v), None, None, nodes_use_cache));

    // Wrap the leaf to secure against the issue bitcoin was vulnerable to.
    let nodes: Result<Vec<Node>, String> = leafs
        .into_iter()
        .map(|l| Node::new(None, Some(Arc::new(l?)), None, nodes_use_cache))
        .collect();
    let nodes = nodes?;

    let mut queue = VecDeque::from(nodes);

    // Special case: If there is only 1 vote leaf, we need to pair
    // it with a Ø-node.
    //
    // Without this, the node would be the merkle root and we would
    // not be able to generate "proof of number of participants".
    if queue.len() == 1 {
        queue.push_back(null_node());
    }

    loop {
        if queue.len() == 1 {
            return Ok(Some(Arc::new(queue.pop_front().unwrap())));
        }
        if queue.len() % 2 != 0 {
            // Push a NULL node to balance level.
            queue.push_back(null_node());
        }
        queue = VecDeque::from(build_level(queue, nodes_use_cache)?);
    }
}

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

    #[test]
    fn test_generate_tree_one_vote() {
        //  Should look like this:
        //
        // ```text
        //          N
        //        /   \
        //       N     Ø
        //       |
        //       V
        // ```
        //
        //   Where V is the vote, N is a node and Ø is a null node.
        //
        let voter = [0xaa; 32];
        let vote = vec![1, 0];
        let root = generate_tree(vec![(voter, vote.clone())], false).unwrap();

        let root = root.unwrap();
        assert!(is_null_node_ref(&root.right));
        let (found_voter, found_vote) = root
            .left
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .vote
            .as_ref()
            .unwrap();

        assert_eq!(&voter, found_voter);
        assert_eq!(&vote, found_vote);
    }

    #[test]
    fn test_generate_tree_is_sorted() {
        //  Should look like this:
        //
        // ```text
        //          N
        //        /   \
        //       N     N
        //      / \   / \
        //     N  N  N   Ø
        //     |  |  |
        //     a  b  c
        // ```

        let root = generate_tree(
            vec![
                ([0xcc; 32], vec![1, 0]),
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
            ],
            false,
        )
        .unwrap()
        .unwrap();
        let (found_voter, _) = root
            .left
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .vote
            .as_ref()
            .unwrap();
        assert_eq!(&[0xaa; 32], found_voter);

        let (found_voter, _) = root
            .right
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .left
            .as_ref()
            .unwrap()
            .vote
            .as_ref()
            .unwrap();
        assert_eq!(&[0xcc; 32], found_voter);
    }

    #[test]
    fn test_generate_tree_rejects_dupes() {
        let error = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xaa; 32], vec![0, 1]), // duplicate
            ],
            false,
        );

        let error_msg = "Cannot have duplicate vote references (found duplciate of 'aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa')";
        assert_eq!(error.err().unwrap(), error_msg);
    }

    #[test]
    fn test_with_and_without_cache_match() {
        let cached = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xcc; 32], vec![0, 1]),
            ],
            true,
        )
        .unwrap();
        let not_cached = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xcc; 32], vec![0, 1]),
            ],
            false,
        )
        .unwrap();

        assert_eq!(
            hash_node_ref(&cached, &Validation::Strict).unwrap(),
            hash_node_ref(&not_cached, &Validation::Strict).unwrap()
        );
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_generate_tree_10k_no_cache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(10000);
        // TODO: How much of the time is `clone`?
        b.iter(|| generate_tree(votes.clone(), false))
    }

    #[cfg(feature = "with-benchmarks")]
    #[bench]
    fn bench_generate_tree_10k_cache(b: &mut Bencher) {
        let votes = crate::utiltest::create_votes(10000);
        // TODO: How much of the time is `clone`?
        b.iter(|| generate_tree(votes.clone(), true))
    }
}