tallytree 0.3.4

use crate::hash::NodeHash;
use crate::hash::{hash_node, ERROR_HASH};
use crate::tally::TallyList;
use crate::utilstring::to_hex;
use crate::{Validation, VoteReference};
#[cfg(feature = "with-serde")]
use serde::{Deserialize, Serialize};
use std::fmt;
use std::sync::Arc;

/// A reference to a node in the merkle tally tree.
pub type NodeRef = Option<Arc<Node>>;

/// Tally tree node.
#[derive(PartialEq)]
#[cfg_attr(feature = "with-serde", derive(Serialize, Deserialize))]
pub struct Node {
    /// Reference to the cast vote and the vote it cast.
    pub vote: Option<(VoteReference, Vec<u32>)>,
    /// First child node
    pub left: NodeRef,
    /// Second child node
    pub right: NodeRef,

    #[cfg_attr(feature = "with-serde", serde(skip_serializing))]
    /// (Optional) cache of the node hash
    pub hash: Option<(NodeHash, Option<TallyList>)>,
}

impl Node {
    /// Create a new tally tree node.
    ///
    /// If `use_cache` is true, the nodes will
    /// be hashed as they are created. This is slower, but improves
    /// performance on calls that require many hash operations, such as
    /// creating proofs.
    pub fn new(
        vote: Option<(VoteReference, Vec<u32>)>,
        left_child: NodeRef,
        right_child: NodeRef,
        use_cache: bool,
    ) -> Result<Node, String> {
        let mut n = Node {
            vote,
            left: left_child,
            right: right_child,
            hash: None,
        };
        if use_cache {
            n.hash = Some(hash_node(&n, &Validation::Relaxed)?);
        }
        Ok(n)
    }
}

fn pretty_format_node(node: &Node, f: &mut fmt::Formatter<'_>) -> fmt::Result {
    let (reference_or_hex, tally): (String, Option<Vec<u32>>) =
        if let Some((reference, t)) = &node.vote {
            // It's a leaf node. Represent it with vote reference.
            (to_hex(reference).unwrap(), Some(t.to_vec()))
        } else {
            // Represent the node with its hash value.
            let (h, t) = hash_node(node, &Validation::Strict).unwrap_or((ERROR_HASH, None));
            (to_hex(h.as_slice()).unwrap(), t)
        };
    let tally = if let Some(t) = tally {
        format!(":{:?}", t)
    } else {
        "".to_string()
    };
    write!(f, "Node {{ {}{} }}", reference_or_hex, tally)
}

impl fmt::Display for Node {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        pretty_format_node(self, f)
    }
}

impl fmt::Debug for Node {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        pretty_format_node(self, f)
    }
}

/// Check if node is a leaf node. Leaf nodes contain the original vote reference.
///
/// Example:
/// ```
/// use tallytree::node::is_leaf_node;
/// use tallytree::generate::generate_tree;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
/// ], false).unwrap();
/// assert!(is_leaf_node(&tree.unwrap().left.as_ref().unwrap().left))
/// ```
pub fn is_leaf_node(node: &NodeRef) -> bool {
    match node {
        Some(n) => n.vote.is_some(),
        None => false,
    }
}

/// If this is a `leaf wrapper`. It's a node above a leaf node, that only
/// points to said node.
///
/// Example:
/// ```
/// use tallytree::node::is_wrapper_node;
/// use tallytree::generate::generate_tree;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
/// ], false).unwrap();
/// assert!(is_wrapper_node(&tree.unwrap().left))
/// ```
pub fn is_wrapper_node(node: &NodeRef) -> bool {
    match node {
        Some(n) => n.left.is_some() && n.right.is_none() && n.vote.is_none(),
        None => false,
    }
}

/// Checks if node is a 'Ø-node'. A 'Ø-node' are dummy nodes used to balance
/// the tree, so that all votes are at the same height in the tree.
///
/// Example:
/// ```
/// use tallytree::node::is_null_node_ref;
/// use tallytree::generate::generate_tree;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
/// ], false).unwrap();
/// assert!(is_null_node_ref(&tree.unwrap().right))
/// ```
pub fn is_null_node_ref(node: &NodeRef) -> bool {
    match node {
        Some(n) => is_null_node(n),
        None => false,
    }
}

/// Checks if node is a 'Ø-node'. A 'Ø-node' are dummy nodes used to balance
/// the tree, so that all votes are at the same height in the tree.
///
/// Example:
/// ```
/// use tallytree::node::{is_null_node, null_node};
/// assert!(is_null_node(&null_node()))
/// ```
pub fn is_null_node(node: &Node) -> bool {
    node.right.is_none() && node.left.is_none() && node.vote.is_none()
}

/// Initialize a 'Ø-node'
///
/// Example:
/// ```
/// use tallytree::node::{is_null_node, null_node};
/// assert!(is_null_node(&null_node()))
/// ```
pub fn null_node() -> Node {
    Node {
        vote: None,
        left: None,
        right: None,
        hash: None,
    }
}

#[cfg(test)]
mod tests {
    use crate::node::*;

    #[test]
    fn test_is_null_node() {
        let n = null_node();
        assert!(is_null_node(&n));

        let mut n = null_node();
        n.left = Some(Arc::new(null_node()));
        assert!(!is_null_node(&n));
    }

    #[test]
    fn test_is_leaf_node() {
        let n = null_node();
        assert!(!is_leaf_node(&Some(Arc::new(n))));

        let mut n = null_node();
        n.vote = Some(([0x0; 32], vec![1]));
        assert!(is_leaf_node(&Some(Arc::new(n))));
    }

    #[test]
    fn test_is_wrapper_node() {
        // null node
        let n = null_node();
        assert!(!is_wrapper_node(&Some(Arc::new(n))));

        // leaf node
        let mut n = null_node();
        n.vote = Some(([0x0; 32], vec![1]));
        assert!(!is_wrapper_node(&Some(Arc::new(n))));

        // inner node
        let mut n = null_node();
        n.left = Some(Arc::new(null_node()));
        n.right = Some(Arc::new(null_node()));
        assert!(!is_wrapper_node(&Some(Arc::new(n))));

        // wrapper node
        let mut n = null_node();
        n.left = Some(Arc::new(null_node()));
        assert!(is_wrapper_node(&Some(Arc::new(n))));
    }
}