Tests for tree layer sizes (#545)
This commit is contained in:
parent
86ee8ff8c7
commit
09cd539bf2
|
@ -149,6 +149,8 @@ impl Tree {
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
|
use std::collections::{HashSet, VecDeque};
|
||||||
|
|
||||||
use super::*;
|
use super::*;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
|
@ -186,6 +188,38 @@ mod tests {
|
||||||
assert_eq!(tree.child_committees(root), (Some(one), Some(two)));
|
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]
|
#[test]
|
||||||
fn test_carnot_tree_leaf_parents() {
|
fn test_carnot_tree_leaf_parents() {
|
||||||
let nodes: Vec<_> = (0..10).map(|i| NodeId::new([i as u8; 32])).collect();
|
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);
|
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(¤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<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(¤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<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);
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
Loading…
Reference in New Issue