extsort_iter/merge/
array_node.rs

1use super::previous_power_of_two;
2
3#[derive(Clone, Copy)]
4pub struct Winner {
5    pub idx: u32,
6}
7
8/// This is a wrapper around the tree indices in
9/// the classic implicit array tree.
10/// it mainly exists for clarity, as well as to make
11/// implementing the tree navigation easier.
12#[derive(Copy, Clone)]
13pub struct TreeNode {
14    pub idx: usize,
15}
16impl TreeNode {
17    pub fn left(self) -> Self {
18        Self {
19            idx: self.idx * 2 + 1,
20        }
21    }
22    pub fn right(self) -> Self {
23        Self {
24            idx: self.idx * 2 + 2,
25        }
26    }
27    pub fn parent(self) -> Self {
28        Self {
29            idx: self.idx.saturating_sub(1) / 2,
30        }
31    }
32    pub fn root() -> Self {
33        Self { idx: 0 }
34    }
35    pub fn is_root(&self) -> bool {
36        self.idx == 0
37    }
38
39    /// computes the tree node the leaf _would_ occuppy if we did store the leaves
40    /// tree_size should be the total number of leaves in the tree.
41    pub fn leaf_for_winner(leaf: Winner, tree_size: usize) -> Self {
42        let full_level_size = previous_power_of_two(tree_size);
43        let overhang = (tree_size - full_level_size) * 2;
44
45        let leaf_idx = leaf.idx as usize;
46
47        let tree_idx = if leaf_idx < overhang {
48            // we are in the lowest level of the tree... the indexes here start at
49            // 2*the filled level - 1
50            full_level_size * 2 - 1 + leaf_idx
51        } else {
52            // we are in the filled part of the tree. the indexes of this one start at
53            // the filled level - 1
54
55            // for every two leaves in the overhang, there is an internal node
56            // that can not be used.
57            // because we would be counting them twice, we need to subtract half again.
58            let overhang_compensation = overhang / 2;
59
60            full_level_size - 1 - overhang_compensation + leaf_idx
61        };
62
63        Self { idx: tree_idx }
64    }
65}
66
67#[cfg(test)]
68mod test {
69    use crate::merge::array_node::{TreeNode, Winner};
70
71    #[test]
72    fn test_leaf_calculation() {
73        fn run_test(expected_values: &[usize]) {
74            let tree_size = expected_values.len() as u32;
75            let values: Vec<_> = (0..tree_size)
76                .map(|leaf| TreeNode::leaf_for_winner(Winner { idx: leaf }, tree_size as usize).idx)
77                .collect();
78            assert_eq!(expected_values, values);
79        }
80        run_test(&[1, 2]);
81        run_test(&[3, 4, 2]);
82        run_test(&[3, 4, 5, 6]);
83        run_test(&[7, 8, 4, 5, 6]);
84        run_test(&[7, 8, 9, 10, 5, 6]);
85        run_test(&[7, 8, 9, 10, 11, 12, 6]);
86        run_test(&[7, 8, 9, 10, 11, 12, 13, 14]);
87        run_test(&[15, 16, 8, 9, 10, 11, 12, 13, 14]);
88    }
89
90    #[test]
91    fn test_clone_to_make_coverage_happy() {
92        let _ = Winner { idx: 0 };
93        let _ = TreeNode { idx: 0 };
94    }
95}
extsort_iter/merge/array_node.rs

extsort_iter/merge/
array_node.rs
