tallytree/

tally.rs

use crate::hash::hash_node;
use crate::node::{is_null_node, is_null_node_ref, is_wrapper_node, Node, NodeRef};
use crate::Validation;

/// List of tallies at a given height.
pub type TallyList = Vec<u32>;

/// Add two tally lists together.
///
/// Example:
/// ```
/// use tallytree::tally::{tally_node_ref, combine_tally};
/// use tallytree::generate::generate_tree;
/// use tallytree::Validation;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
/// ], false).unwrap().unwrap();
/// assert_eq!(Ok([1, 1].to_vec()), combine_tally(
///     &tally_node_ref(&tree.left, &Validation::Strict).unwrap(),
///     &tally_node_ref(&tree.right, &Validation::Strict).unwrap(),
/// ));
/// ```
pub fn combine_tally(
    left: &Option<TallyList>,
    right: &Option<TallyList>,
) -> Result<TallyList, String> {
    if left.is_none() {
        // Left node should never be Ø-node and always have a tally.
        return Err("Left node does not contain a tally".to_string());
    }
    if right.is_none() {
        return Ok(left.as_ref().unwrap().clone());
    }

    let left = left.as_ref().unwrap();
    let right = right.as_ref().unwrap();

    if left.len() != right.len() {
        return Err(format!(
            "Left tally length is not equal to right tally \
            length ({} != {})",
            left.len(),
            right.len()
        ));
    }
    Ok(left.iter().zip(right.iter()).map(|(l, r)| l + r).collect())
}

/// Checks if tally has a single vote
///
/// Example:
/// ```
/// use tallytree::tally::has_one_vote;
/// assert!(has_one_vote(&[1, 0, 0]));
/// assert!(!has_one_vote(&[1, 0, 1]));
/// ```
pub fn has_one_vote(tally: &[u32]) -> bool {
    tally.iter().sum::<u32>() == 1
}

/// Given a node in the merkle tally tree, prints the tree to stdout. Pass `0`
/// for `level`argument.
///
/// Leaf nodes are prefixed with their `VoteReference`. Parent nodes are
/// prefixed with their hash.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::tally::pretty_print_tally;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
/// ], false).unwrap();
/// pretty_print_tally(&tree, 0);
/// ```
pub fn pretty_print_tally(node: &NodeRef, level: usize) {
    match node {
        Some(n) => {
            if is_null_node(n) {
                println!("{:indent$}Ø", "", indent = level * 4);
                return;
            }
            let prefix = match n.vote {
                Some((v, _)) => [v[0], v[1]],
                None => {
                    let h = hash_node(n, &Validation::Strict).unwrap().0;
                    [h.as_slice()[0], h.as_slice()[1]]
                }
            };
            if is_wrapper_node(node) {
                print!(
                    "{:indent$}{:x?}=>{:?} --> ",
                    "",
                    prefix,
                    tally_node(n, &Validation::Strict).unwrap().unwrap(),
                    indent = level * 4
                );
                pretty_print_tally(&n.left, 0);
                assert!(n.right.is_none());
                return;
            }

            pretty_print_tally(&n.right, level + 1);
            println!(
                "{:indent$}{:x?}=>{:?}",
                "",
                prefix,
                tally_node(n, &Validation::Strict).unwrap().unwrap(),
                indent = level * 4
            );
            pretty_print_tally(&n.left, level + 1);
        }
        None => {}
    }
}

/// Tally all the votes in a tree.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::tally::tally_node;
/// use tallytree::Validation;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
/// ], false).unwrap();
/// assert_eq!(Some([1, 1].to_vec()), tally_node(&tree.unwrap(), &Validation::Strict).unwrap());
/// ```
pub fn tally_node(node: &Node, v: &Validation) -> Result<Option<TallyList>, String> {
    if is_null_node(node) {
        return Ok(None);
    }
    if let Some((_, tally)) = &node.vote {
        if node.left.is_some() || node.right.is_some() {
            return Err("Leaf has a child".to_string());
        }
        if !matches!(v, Validation::Relaxed) && !has_one_vote(tally) {
            return Err("Leaf casts more than 1 vote".to_string());
        }
        return Ok(Some(tally.to_vec()));
    }
    if node.right.is_none() || is_null_node_ref(&node.right) {
        return tally_node_ref(&node.left, v);
    }
    Ok(Some(combine_tally(
        &tally_node_ref(&node.left, v)?,
        &tally_node_ref(&node.right, v)?,
    )?))
}

/// Tally all the votes in a tree.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::tally::tally_node_ref;
/// use tallytree::Validation;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
/// ], false).unwrap();
/// assert_eq!(Some([1, 1].to_vec()), tally_node_ref(&tree, &Validation::Strict).unwrap());
/// ```
pub fn tally_node_ref(node: &NodeRef, v: &Validation) -> Result<Option<TallyList>, String> {
    if let Some(n) = node {
        tally_node(n, v)
    } else {
        Ok(None)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::generate::generate_tree;

    #[test]
    fn test_pretty_print() {
        let head = generate_tree(
            vec![
                ([0x11; 32], vec![1, 0, 0]),
                ([0x22; 32], vec![0, 0, 1]),
                ([0x33; 32], vec![0, 0, 1]),
                ([0x44; 32], vec![0, 0, 1]),
                ([0x55; 32], vec![0, 0, 1]),
            ],
            false,
        )
        .unwrap();
        pretty_print_tally(&head, 0);
        assert_eq!(1, 1);
    }

    #[test]
    fn test_combine_tally() {
        assert_eq!(
            Ok(vec![2, 0, 1]),
            combine_tally(&Some(vec![1, 0, 0]), &Some(vec![1, 0, 1]))
        );

        assert_eq!(Ok(vec![2, 0]), combine_tally(&Some(vec![2, 0]), &None,));

        // Left node must contain a tally.
        assert!(combine_tally(&None, &Some(vec![2, 0])).is_err());

        // Tally list must be equal length
        assert!(combine_tally(&Some(vec![0, 0, 1]), &Some(vec![2, 0])).is_err());
    }
}