stratifykit_core/
heap.rs

1use std::collections::BinaryHeap;
2
3/// One candidate in a bounded top-K stream. `key` carries every tie-break the equivalent
4/// full-materialize-then-`sort_by` code would apply; `index` is always the final tiebreak
5/// layer, reproducing `sort_by`'s stability -- which a heap has no notion of on its own, since
6/// two records can otherwise agree on every field `key` compares. `record` is left generic (`R`)
7/// rather than fixed to any one domain's row type, so this heap is reusable outside shogiesa.
8pub struct HeapEntry<K: Ord, R> {
9    pub key: K,
10    pub index: usize,
11    pub record: R,
12}
13
14impl<K: Ord, R> PartialEq for HeapEntry<K, R> {
15    fn eq(&self, other: &Self) -> bool {
16        self.key == other.key && self.index == other.index
17    }
18}
19impl<K: Ord, R> Eq for HeapEntry<K, R> {}
20impl<K: Ord, R> PartialOrd for HeapEntry<K, R> {
21    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
22        Some(self.cmp(other))
23    }
24}
25impl<K: Ord, R> Ord for HeapEntry<K, R> {
26    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
27        self.key
28            .cmp(&other.key)
29            .then_with(|| self.index.cmp(&other.index))
30    }
31}
32
33/// Keeps the `count` smallest `HeapEntry`s seen so far -- the standard bounded-heap top-K
34/// algorithm: push while under capacity, otherwise evict the current worst-kept entry if `entry`
35/// ranks ahead of it. Provably identical final set (and, via `BinaryHeap::into_sorted_vec`,
36/// identical order) to "collect everything, sort ascending by the same key, truncate" -- at
37/// O(count) memory instead of O(n).
38pub fn push_bounded<K: Ord, R>(
39    heap: &mut BinaryHeap<HeapEntry<K, R>>,
40    count: usize,
41    entry: HeapEntry<K, R>,
42) {
43    if heap.len() < count {
44        heap.push(entry);
45    } else if heap.peek().is_some_and(|worst| entry.cmp(worst).is_lt()) {
46        heap.pop();
47        heap.push(entry);
48    }
49}
50
51/// `f32` wrapper with a total order (via `total_cmp`) for callers ranking by a plain fraction
52/// (e.g. a 0..=1 quality score) that is never NaN in practice but has no `Ord` impl of its own.
53#[derive(Debug, Clone, Copy, PartialEq)]
54pub struct TotalF32(pub f32);
55impl Eq for TotalF32 {}
56impl PartialOrd for TotalF32 {
57    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
58        Some(self.cmp(other))
59    }
60}
61impl Ord for TotalF32 {
62    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
63        self.0.total_cmp(&other.0)
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70
71    #[test]
72    fn push_bounded_keeps_k_smallest_keys() {
73        let mut heap: BinaryHeap<HeapEntry<u32, &'static str>> = BinaryHeap::new();
74        for (key, record) in [(5, "a"), (1, "b"), (9, "c"), (2, "d"), (7, "e")] {
75            push_bounded(
76                &mut heap,
77                3,
78                HeapEntry {
79                    key,
80                    index: key as usize,
81                    record,
82                },
83            );
84        }
85        let mut kept: Vec<u32> = heap.into_vec().into_iter().map(|e| e.key).collect();
86        kept.sort();
87        assert_eq!(kept, vec![1, 2, 5]);
88    }
89
90    #[test]
91    fn push_bounded_matches_sort_then_truncate() {
92        let items: Vec<(u32, usize)> = vec![(3, 0), (1, 1), (4, 2), (1, 3), (5, 4), (9, 5), (2, 6)];
93        let mut heap: BinaryHeap<HeapEntry<(u32, usize), usize>> = BinaryHeap::new();
94        for &(key, index) in &items {
95            push_bounded(
96                &mut heap,
97                4,
98                HeapEntry {
99                    key: (key, index),
100                    index,
101                    record: index,
102                },
103            );
104        }
105        let mut via_heap: Vec<usize> = heap
106            .into_sorted_vec()
107            .into_iter()
108            .map(|e| e.record)
109            .collect();
110
111        let mut via_sort = items.clone();
112        via_sort.sort_by_key(|&(key, index)| (key, index));
113        via_sort.truncate(4);
114        let via_sort: Vec<usize> = via_sort.into_iter().map(|(_, index)| index).collect();
115
116        via_heap.sort();
117        let mut via_sort_sorted = via_sort.clone();
118        via_sort_sorted.sort();
119        assert_eq!(via_heap, via_sort_sorted);
120    }
121
122    #[test]
123    fn total_f32_orders_like_a_float_comparison() {
124        let mut v = vec![TotalF32(0.5), TotalF32(0.1), TotalF32(0.9)];
125        v.sort();
126        assert_eq!(
127            v.into_iter().map(|f| f.0).collect::<Vec<_>>(),
128            vec![0.1, 0.5, 0.9]
129        );
130    }
131}
stratifykit_core/heap.rs

stratifykit_core/
heap.rs
