From 09cd539bf29fa159a50723f193aa3d04c7c8daac Mon Sep 17 00:00:00 2001 From: gusto Date: Wed, 3 Jan 2024 15:19:42 +0200 Subject: [PATCH] Tests for tree layer sizes (#545) --- .../src/overlay/tree_overlay/tree.rs | 192 ++++++++++++++++++ 1 file changed, 192 insertions(+) diff --git a/consensus-engine/src/overlay/tree_overlay/tree.rs b/consensus-engine/src/overlay/tree_overlay/tree.rs index c9436e10..d6536590 100644 --- a/consensus-engine/src/overlay/tree_overlay/tree.rs +++ b/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 { + let root = tree.root_committee(); + let root_id = root.id::>(); + + 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(¤t_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 { + // TODO: Find a more elegant way to visit parents of leaves from different levels. + let leaves: HashMap = + 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 = + 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(¤t_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, + } + + 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); + } + } }