extsort_iter/merge/
treebuilder.rs

1use std::{cmp::Ordering, ops::Range};
2
3use super::{TreeNode, Winner};
4
5/// This module contains the code to construct a complete loser tree
6/// in an implicit array representation.
7
8/// This is a convenience struct to move the tree construction code out from the main merge
9/// code
10pub(super) struct LoserTreeBuilder<'a, C> {
11    winner_comparer: C,
12    loser_tree: &'a mut Vec<u32>,
13}
14
15/// split a range at the specified point
16/// the split point shall denote the size of the left
17/// range piece after the split.
18fn split_range(range: Range<u32>, split_point: usize) -> (Range<u32>, Range<u32>) {
19    let midpoint = range.start + split_point as u32;
20    let mut left = range.clone();
21    let mut right = range;
22    left.end = midpoint;
23    right.start = midpoint;
24    (left, right)
25}
26
27impl<'a, C> LoserTreeBuilder<'a, C>
28where
29    C: FnMut(Winner, Winner) -> Ordering,
30{
31    pub fn new(winner_comparer: C, loser_tree: &'a mut Vec<u32>) -> Self {
32        Self {
33            winner_comparer,
34            loser_tree,
35        }
36    }
37
38    pub(super) fn build(mut self, number_of_tapes: usize) -> Winner {
39        let num_internal_nodes = number_of_tapes - 1;
40        if self.loser_tree.len() < num_internal_nodes {
41            *self.loser_tree = vec![0; num_internal_nodes];
42        }
43        self.build_tree_complete(0..number_of_tapes as u32, TreeNode::root())
44    }
45
46    fn build_tree_perfect(&mut self, range_todo: Range<u32>, root: TreeNode) -> Winner {
47        if range_todo.len() == 1 {
48            Winner {
49                idx: range_todo.start,
50            }
51        } else {
52            let subtree_size = range_todo.len() / 2;
53            let (left, right) = split_range(range_todo, subtree_size);
54            let winner_left = self.build_tree_perfect(left, root.left());
55            let winner_right = self.build_tree_perfect(right, root.right());
56            self.commit_winner(winner_left, winner_right, root)
57        }
58    }
59
60    fn build_tree_complete(&mut self, range_todo: Range<u32>, root: TreeNode) -> Winner {
61        let total_nodes = range_todo.len();
62        if total_nodes.is_power_of_two() {
63            self.build_tree_perfect(range_todo, root)
64        } else {
65            let nodes_if_lowest_level_was_full = total_nodes.next_power_of_two();
66            let nodes_in_lower_level = (total_nodes - nodes_if_lowest_level_was_full / 2) * 2;
67
68            let perfect_tree_left = nodes_in_lower_level >= nodes_if_lowest_level_was_full / 2;
69
70            if perfect_tree_left {
71                let nodes_in_left_tree = nodes_if_lowest_level_was_full / 2;
72                let (left, right) = split_range(range_todo, nodes_in_left_tree);
73                let w_left = self.build_tree_perfect(left, root.left());
74                let w_right = self.build_tree_complete(right, root.right());
75                self.commit_winner(w_left, w_right, root)
76            } else {
77                // there are _not_ enough nodes to fill left side of the tree completely.
78                // therefore the perfect tree must be on the right and have the size of half the upper level
79                let nodes_in_right_tree = nodes_if_lowest_level_was_full / 2 / 2;
80                let nodes_in_left_tree = total_nodes - nodes_in_right_tree;
81                let (left, right) = split_range(range_todo, nodes_in_left_tree);
82                let w_left = self.build_tree_complete(left, root.left());
83                let w_right = self.build_tree_perfect(right, root.right());
84                self.commit_winner(w_left, w_right, root)
85            }
86        }
87    }
88
89    fn commit_winner(
90        &mut self,
91        candidate_a: Winner,
92        candidate_b: Winner,
93        root: TreeNode,
94    ) -> Winner {
95        let (winner, loser) = if (self.winner_comparer)(candidate_a, candidate_b).is_le() {
96            // left side wins !
97            (candidate_a, candidate_b)
98        } else {
99            // right side wins
100            (candidate_b, candidate_a)
101        };
102        self.loser_tree[root.idx] = loser.idx;
103        winner
104    }
105}
106
107#[cfg(test)]
108mod test {
109    use crate::orderer::{OrdOrderer, Orderer};
110
111    use super::LoserTreeBuilder;
112
113    fn assert_winner(tape: &[i64]) {
114        let orderer = OrdOrderer {};
115        let min_value = *tape.iter().min().unwrap();
116        let ref_tape = tape.iter().collect::<Vec<_>>();
117        let mut tree = Vec::new();
118        let winner = LoserTreeBuilder::new(
119            |a, b| orderer.compare(ref_tape[a.idx as usize], ref_tape[b.idx as usize]),
120            &mut tree,
121        )
122        .build(ref_tape.len());
123        assert_eq!(min_value, tape[winner.idx as usize]);
124        if tape.len() > 1 {
125            assert!(min_value <= tape[tree[0] as usize])
126        }
127    }
128
129    #[test]
130    fn test_power_two() {
131        assert!(1usize.is_power_of_two())
132    }
133
134    fn run_tree_construction_test(max_size: usize) {
135        for r in 1..max_size {
136            println!("constructing zero tree with {r} tapes");
137            assert_winner(&vec![0; r]);
138            println!("constructing ordered tree with {r} tapes");
139            let mut tape = (0..r as i64).collect::<Vec<_>>();
140            assert_winner(&tape);
141            println!("constructing reversed tree wtih {r} tapes");
142            let mut reversed_tape = tape.clone();
143            reversed_tape.reverse();
144            assert_winner(&reversed_tape);
145            println!("constructing tree with min in the middle");
146            reversed_tape.append(&mut tape);
147            assert_winner(&reversed_tape);
148        }
149    }
150
151    #[test]
152    fn test_construct_tree_small() {
153        run_tree_construction_test(10);
154    }
155
156    // this would be too slow to execute with miri.
157    // instead, there is a smaller version with less cases.
158    #[test]
159    #[cfg(not(miri))]
160    fn test_construct_tree() {
161        run_tree_construction_test(100);
162    }
163}
extsort_iter/merge/treebuilder.rs

extsort_iter/merge/
treebuilder.rs
