tallytree 0.3.4

use crate::node::{is_null_node_ref, is_wrapper_node, Node};
use crate::{NodeRef, VoteReference};
use std::collections::VecDeque;
use std::sync::Arc;

/// A pair of vote references
pub type VoteReferencePair = (Option<VoteReference>, Option<VoteReference>);

/// Find the path to the inner node furthest down to the right.
/// Returns all nodes in the path.
/// ```text
///            A
///          /  \
///         B    C <- We want the path to this one
///        / \   | \
///       D   E  F  Ø
///       |   |  |
///       V1  V2 V3
/// ```
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::navigate::find_down_right_most_branch;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![0, 1]),
///     ([0xcc; 32], vec![1, 0]),
/// ], false).unwrap();
/// let path = find_down_right_most_branch(&tree);
/// ```
pub fn find_down_right_most_branch(node: &NodeRef) -> Option<Vec<NodeRef>> {
    if node.is_none() {
        return None;
    }
    if is_null_node_ref(node) {
        return None;
    }
    let node_ref = node.as_ref().unwrap();
    let node_vec = vec![node.clone()];
    if is_wrapper_node(&node_ref.left) {
        // We're at the last branch, success.
        return Some(node_vec);
    }
    if is_null_node_ref(&node_ref.right) {
        return match find_down_right_most_branch(&node_ref.left) {
            Some(path) => Some([node_vec, path].concat()),
            None => Some(node_vec),
        };
    }
    match find_down_right_most_branch(&node_ref.right) {
        Some(path) => Some([node_vec, path].concat()),
        None => Some(node_vec),
    }
}

/// Find a node in a merkle tally tree given a vote reference.
/// Returns all nodes in the path.
///
/// Example:
/// ```
/// use tallytree::generate::generate_tree;
/// use tallytree::navigate::find_node_by_vote_reference;
/// let tree = generate_tree(vec![
///     ([0xaa; 32], vec![1, 0]),
///     ([0xbb; 32], vec![1, 0]),
///     ([0xcc; 32], vec![0, 1]),
/// ], false).unwrap();
/// let path = find_node_by_vote_reference(&tree, &[0xbb; 32]);
/// ```
/// The above example returns A, B, E, bb
/// ```text
///            A
///          /  \
///         B    C
///        / \   | \
///       D   E  F  Ø
///       |   |  |
///       aa  bb cc
/// ```
pub fn find_node_by_vote_reference(
    node: &NodeRef,
    reference: &VoteReference,
) -> Option<Vec<NodeRef>> {
    // TODO: This can probably be optimized by using the fact that the leafs are sorted.

    if node.is_none() {
        return None;
    }

    if let Some((r, _)) = node.as_ref().unwrap().vote {
        return if &r == reference {
            // Found it
            Some(vec![node.clone()])
        } else {
            None
        };
    }
    let node_vec = vec![node.clone()];
    let node = node.as_ref().unwrap();
    if let Some(path) = find_node_by_vote_reference(&node.left, reference) {
        return Some([node_vec, path].concat());
    }
    if let Some(path) = find_node_by_vote_reference(&node.right, reference) {
        return Some([node_vec, path].concat());
    }
    None
}

/**
 * Finds the two nearest sibling vote references of a given reference. The
 * given reference does not need to exist. A boolean representing if it exists
 * or not is returned.
 *
 * Example:
 * ```
 * use tallytree::generate::generate_tree;
 * use tallytree::navigate::find_closest_siblings;
 * let tree = generate_tree(vec![
 *     ([0xaa; 32], vec![1, 0]),
 *     ([0xbb; 32], vec![1, 0]),
 *     ([0xdd; 32], vec![0, 1]),
 * ], false).unwrap();
 * let ((left, right), exists) = find_closest_siblings(&tree, &[0xcc; 32]).unwrap();
 * assert!(!exists);
 * assert_eq!(left, Some([0xbb; 32]));
 * assert_eq!(right, Some([0xdd; 32]));
 * ```
 */
pub fn find_closest_siblings(
    root: &NodeRef,
    reference: &VoteReference,
) -> Result<(VoteReferencePair, bool), String> {
    // TODO: This can probably be optimized by using the fact that the leafs are sorted.

    let root: &Arc<Node> = root
        .as_ref()
        .ok_or_else(|| "Invalid root node".to_string())?;

    let mut queue = VecDeque::from([root]);
    let mut left_sibling: Option<VoteReference> = None;
    let mut right_sibling: Option<VoteReference> = None;
    let mut reference_exists = false;

    while let Some(node) = queue.pop_front() {
        if let Some(n) = &node.left {
            queue.push_back(n);
        }
        if let Some(n) = &node.right {
            queue.push_back(n);
        }

        if let Some((r, _)) = &node.vote {
            if r == reference {
                reference_exists = true;
            }
            if r < reference {
                if let Some(l) = left_sibling {
                    if r > &l {
                        left_sibling = Some(*r);
                    }
                } else {
                    left_sibling = Some(*r);
                }
            }
            if r > reference {
                if let Some(l) = right_sibling {
                    if r < &l {
                        right_sibling = Some(*r);
                    }
                } else {
                    right_sibling = Some(*r);
                }
            }
        }
    }
    Ok(((left_sibling, right_sibling), reference_exists))
}

/// Find the depth of the merkle tally tree at this node, including leaf wrapper nodes.
///
/// For example, a tree with 3 nodes will have a depth of 5.
/// ```text
///            A
///          /  \
///         B    C
///        / \   | \
///       D   E  F  Ø
///       |   |  |
///       V1  V2 V3
/// ```
pub fn tree_depth(node: &NodeRef) -> usize {
    match node {
        Some(n) => 1 + tree_depth(&n.left),
        None => 0,
    }
}

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

    #[test]
    fn test_find_down_right_most_branch() {
        // We want to find A:
        //
        // ```text
        //          A
        //        /   \
        //       B     Ø
        //       |
        //       V
        // ```
        let root = generate_tree(vec![([0xaa; 32], vec![1, 0])], true).unwrap();
        let path = find_down_right_most_branch(&root).unwrap();
        assert_eq!(path.len(), 1);
        assert_eq!(
            hash_node_ref(&path[0], &Validation::Strict),
            hash_node_ref(&root, &Validation::Strict)
        );

        // We want to find C
        // ```text
        //            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 path = find_down_right_most_branch(&root).unwrap();
        assert_eq!(path.len(), 2);
        assert_eq!(
            hash_node_ref(&path[1], &Validation::Strict),
            hash_node_ref(&root.unwrap().right, &Validation::Strict)
        );
    }

    #[test]
    fn test_find_node_by_vote_reference() {
        // ```text
        //            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();

        // Find V2
        let path = find_node_by_vote_reference(&root, &[0xbb; 32]).unwrap();
        assert_eq!(path.len(), 4);
        assert_eq!(
            hash_node_ref(&path[0], &Validation::Strict),
            hash_node_ref(&root, &Validation::Strict)
        );
        assert_eq!(
            hash_node_ref(&path[1], &Validation::Strict),
            hash_node_ref(&root.as_ref().unwrap().left, &Validation::Strict)
        );
        assert_eq!(
            hash_node_ref(&path[2], &Validation::Strict),
            hash_node_ref(
                &root.as_ref().unwrap().left.as_ref().unwrap().right,
                &Validation::Strict
            )
        );
        assert_eq!(
            hash_node_ref(&path[3], &Validation::Strict),
            hash_node_ref(
                &root
                    .as_ref()
                    .unwrap()
                    .left
                    .as_ref()
                    .unwrap()
                    .right
                    .as_ref()
                    .unwrap()
                    .left,
                &Validation::Strict
            )
        );
    }

    #[test]
    fn test_find_closest_siblings() {
        // ```text
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        // ```
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xdd; 32], vec![0, 1]),
            ],
            false,
        )
        .unwrap();

        // V1 has no sibling to the left
        let ((left, right), exists) = find_closest_siblings(&root, &[0xaa; 32]).unwrap();
        assert!(exists);
        assert!(left.is_none());
        assert_eq!(right, Some([0xbb; 32]));

        // V2 has V1 to the left, and V3 to the right.
        let ((left, right), exists) = find_closest_siblings(&root, &[0xbb; 32]).unwrap();
        assert!(exists);
        assert_eq!(left, Some([0xaa; 32]));
        assert_eq!(right, Some([0xdd; 32]));

        // V3 has V2 to the left, and none to the right.
        let ((left, right), exists) = find_closest_siblings(&root, &[0xdd; 32]).unwrap();
        assert!(exists);
        assert_eq!(left, Some([0xbb; 32]));
        assert!(right.is_none());

        // 0xcc does not exist, but would have V2 to the left and V3 to the right.
        let ((left, right), exists) = find_closest_siblings(&root, &[0xcc; 32]).unwrap();
        assert!(!exists);
        assert_eq!(left, Some([0xbb; 32]));
        assert_eq!(right, Some([0xdd; 32]));
    }

    #[test]
    fn test_tree_depth() {
        // ```text
        //          A
        //        /   \
        //       B     Ø
        //       |
        //       V
        // ```
        let root = generate_tree(vec![([0xaa; 32], vec![1, 0])], false).unwrap();
        assert_eq!(3, tree_depth(&root));

        // ```text
        //            A
        //          /  \
        //         B    C
        //        / \   | \
        //       D   E  F  Ø
        //       |   |  |
        //       V1  V2 V3
        // ```
        let root = generate_tree(
            vec![
                ([0xaa; 32], vec![1, 0]),
                ([0xbb; 32], vec![1, 0]),
                ([0xdd; 32], vec![0, 1]),
            ],
            false,
        )
        .unwrap();

        assert_eq!(4, tree_depth(&root));
    }
}