Tests for tree layer sizes (#545)

consensus-engine/src/overlay/tree_overlay/tree.rs

@@ -149,6 +149,8 @@ impl Tree {
 
 #[cfg(test)]
 mod tests {
+    use std::collections::{HashSet, VecDeque};
+
     use super::*;
 
     #[test]
@@ -186,6 +188,38 @@ mod tests {
         assert_eq!(tree.child_committees(root), (Some(one), Some(two)));
     }
 
+    #[test]
+    fn test_carnot_tree_leaf_committees() {
+        let test_cases = [
+            (3, 2),
+            (4, 2),
+            (5, 3),
+            (6, 3),
+            (7, 4),
+            (8, 4),
+            (9, 5),
+            (10, 5),
+            (11, 6),
+            (12, 6),
+            (13, 7),
+            (14, 7),
+            (15, 8),
+            (17, 9),
+            (18, 9),
+            (31, 16),
+        ];
+        for (committees, expected_leaves) in test_cases {
+            let nodes: Vec<_> = (0..committees)
+                .map(|i| NodeId::new([i as u8; 32]))
+                .collect();
+            let tree = Tree::new(&nodes, committees);
+
+            let leaves = tree.leaf_committees();
+
+            assert_eq!(leaves.len(), expected_leaves);
+        }
+    }
+
     #[test]
     fn test_carnot_tree_leaf_parents() {
         let nodes: Vec<_> = (0..10).map(|i| NodeId::new([i as u8; 32])).collect();
@@ -239,4 +273,162 @@ mod tests {
             test_committee_parent(&tree, tcase.1 .1.unwrap(), tcase.0);
         });
     }
+
+    fn carnot_tree_level_sizes_from_top(tree: Tree) -> HashMap<i32, i32> {
+        let root = tree.root_committee();
+        let root_id = root.id::<blake2::Blake2b<digest::typenum::U32>>();
+
+        let mut queue = VecDeque::new();
+        queue.push_back((root_id, 0));
+
+        let mut level_sizes = HashMap::new();
+
+        while let Some((current_id, level)) = queue.pop_front() {
+            *level_sizes.entry(level).or_insert(0) += 1;
+
+            let children = tree.child_committees(&current_id);
+            if let Some(id) = children.0 {
+                queue.push_back((*id, level + 1));
+            }
+            if let Some(id) = children.1 {
+                queue.push_back((*id, level + 1));
+            }
+        }
+
+        level_sizes
+    }
+
+    fn next_pow2(n: usize) -> usize {
+        let mut n = n - 1;
+        n |= n >> 1;
+        n |= n >> 2;
+        n |= n >> 4;
+        n |= n >> 8;
+        n |= n >> 16;
+        n |= n >> 32;
+        n + 1
+    }
+
+    fn is_full_balanced_binary_tree(committee_count: usize) -> bool {
+        let mut count = 1;
+        while count < committee_count + 1 {
+            count *= 2;
+        }
+        count == committee_count + 1
+    }
+
+    fn carnot_tree_level_sizes_from_bottom(tree: Tree) -> HashMap<i32, i32> {
+        // TODO: Find a more elegant way to visit parents of leaves from different levels.
+        let leaves: HashMap<CommitteeId, i32> =
+            if is_full_balanced_binary_tree(tree.inner_committees.len()) {
+                tree.leaf_committees().keys().map(|c| (**c, 0)).collect()
+            } else {
+                let next_full_tree_committees = next_pow2(tree.inner_committees.len()) - 1;
+                let next_full_tree_leafs = (next_full_tree_committees + 1) / 2;
+                let prev_full_tree_committees = next_full_tree_committees - next_full_tree_leafs;
+
+                let mut leaf_committees: HashMap<CommitteeId, i32> =
+                    tree.leaf_committees().keys().map(|c| (**c, 1)).collect();
+
+                for i in prev_full_tree_committees..tree.inner_committees.len() {
+                    leaf_committees.insert(tree.inner_committees[i], 0);
+                }
+
+                leaf_committees
+            };
+
+        let mut queue = VecDeque::new();
+        let mut level_sizes = HashMap::new();
+        let mut visited = HashSet::new();
+
+        for (leaf_id, level) in leaves {
+            queue.push_back((leaf_id, level));
+            visited.insert(leaf_id);
+        }
+
+        while let Some((current_id, level)) = queue.pop_front() {
+            *level_sizes.entry(level).or_insert(0) += 1;
+
+            if let Some(parent_id) = tree.parent_committee(&current_id) {
+                if !visited.contains(parent_id) {
+                    queue.push_back((*parent_id, level + 1));
+                    visited.insert(*parent_id);
+                }
+            }
+        }
+
+        level_sizes
+    }
+
+    #[test]
+    fn test_carnot_tree_level_sizes() {
+        struct TestCase {
+            num_committees: usize,
+            max_depth: usize,
+            expected_level_sizes: HashMap<i32, i32>,
+        }
+
+        let test_cases = vec![
+            TestCase {
+                num_committees: 1,
+                max_depth: 1,
+                expected_level_sizes: [(0, 1)].iter().cloned().collect(),
+            },
+            TestCase {
+                num_committees: 3,
+                max_depth: 2,
+                expected_level_sizes: [(0, 1), (1, 2)].iter().cloned().collect(),
+            },
+            TestCase {
+                num_committees: 7,
+                max_depth: 3,
+                expected_level_sizes: [(0, 1), (1, 2), (2, 4)].iter().cloned().collect(),
+            },
+            TestCase {
+                num_committees: 15,
+                max_depth: 4,
+                expected_level_sizes: [(0, 1), (1, 2), (2, 4), (3, 8)].iter().cloned().collect(),
+            },
+            TestCase {
+                num_committees: 11,
+                max_depth: 4,
+                expected_level_sizes: [(0, 1), (1, 2), (2, 4), (3, 4)].iter().cloned().collect(),
+            },
+            TestCase {
+                num_committees: 25,
+                max_depth: 5,
+                expected_level_sizes: [(0, 1), (1, 2), (2, 4), (3, 8), (4, 10)]
+                    .iter()
+                    .cloned()
+                    .collect(),
+            },
+            TestCase {
+                num_committees: 40,
+                max_depth: 6,
+                expected_level_sizes: [(0, 1), (1, 2), (2, 4), (3, 8), (4, 16), (5, 9)]
+                    .iter()
+                    .cloned()
+                    .collect(),
+            },
+        ];
+
+        for case in test_cases {
+            let nodes: Vec<_> = (0..case.num_committees)
+                .map(|i| NodeId::new([i as u8; 32]))
+                .collect();
+
+            let tree = Tree::new(&nodes, case.num_committees);
+
+            let level_sizes_top = carnot_tree_level_sizes_from_top(tree.clone());
+            let level_sizes_bottom = carnot_tree_level_sizes_from_bottom(tree);
+
+            for (top_level, top_size) in level_sizes_top.iter() {
+                let bottom_level = case.max_depth as i32 - top_level - 1;
+                let bottom_size = level_sizes_bottom.get(&bottom_level).unwrap();
+                assert_eq!(top_size, bottom_size);
+            }
+
+            assert_eq!(level_sizes_top, case.expected_level_sizes);
+        }
+    }
 }