1use std::cmp::Ordering;
7use std::hash::{Hash, Hasher};
8
9#[derive(Copy, Clone, Debug, Eq, PartialEq)]
14pub struct QuantKey {
15 q: i32,
16 id_prefix: u32,
17}
18
19impl Hash for QuantKey {
20 fn hash<H: Hasher>(&self, state: &mut H) {
21 self.q.hash(state);
22 self.id_prefix.hash(state);
23 }
24}
25
26impl QuantKey {
27 const SCALE: f32 = 1_000_000.0;
28
29 #[inline]
30 pub fn new(score: f32, id_prefix: u32) -> Self {
31 let s = if score.is_nan() { 0.0 } else { score };
32 let q = (s * Self::SCALE)
33 .round()
34 .clamp(i32::MIN as f32, i32::MAX as f32) as i32;
35 Self { q, id_prefix }
36 }
37
38 #[inline]
39 pub fn from_f64(score: f64, id_prefix: u32) -> Self {
40 Self::new(score as f32, id_prefix)
41 }
42
43 #[inline]
44 pub fn quantized_score(&self) -> i32 {
45 self.q
46 }
47
48 #[inline]
49 pub fn score(&self) -> f32 {
50 self.q as f32 / Self::SCALE
51 }
52
53 #[inline]
54 pub fn id_prefix(&self) -> u32 {
55 self.id_prefix
56 }
57}
58
59impl Ord for QuantKey {
60 #[inline]
61 fn cmp(&self, other: &Self) -> Ordering {
62 self.q
63 .cmp(&other.q)
64 .then_with(|| other.id_prefix.cmp(&self.id_prefix))
65 }
66}
67
68impl PartialOrd for QuantKey {
69 #[inline]
70 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
71 Some(self.cmp(other))
72 }
73}
74
75#[cfg(test)]
76mod tests {
77 use super::*;
78 use std::collections::BinaryHeap;
79
80 #[test]
81 fn size_is_8_bytes() {
82 assert_eq!(std::mem::size_of::<QuantKey>(), 8);
83 }
84
85 #[test]
86 fn precision() {
87 let a = QuantKey::new(0.123456, 1);
88 let b = QuantKey::new(0.123457, 2);
89 assert_ne!(a.quantized_score(), b.quantized_score());
90 }
91
92 #[test]
93 fn heap_order() {
94 let mut heap: BinaryHeap<QuantKey> = BinaryHeap::new();
95 heap.push(QuantKey::new(0.95, 3));
96 heap.push(QuantKey::new(0.95, 1));
97 heap.push(QuantKey::new(0.95, 2));
98 heap.push(QuantKey::new(0.87, 4));
99
100 assert_eq!(heap.pop().unwrap().id_prefix(), 1);
101 assert_eq!(heap.pop().unwrap().id_prefix(), 2);
102 assert_eq!(heap.pop().unwrap().id_prefix(), 3);
103 assert_eq!(heap.pop().unwrap().id_prefix(), 4);
104 }
105
106 #[test]
107 fn nan_maps_to_zero() {
108 let nan_key = QuantKey::new(f32::NAN, 1);
109 let zero_key = QuantKey::new(0.0, 1);
110 assert_eq!(nan_key.quantized_score(), zero_key.quantized_score());
111 }
112
113 #[test]
114 fn clamp_high_score() {
115 let key = QuantKey::new(f32::MAX, 0);
116 assert_eq!(key.quantized_score(), i32::MAX);
117 }
118
119 #[test]
120 fn clamp_low_score() {
121 let key = QuantKey::new(f32::MIN, 0);
122 assert_eq!(key.quantized_score(), i32::MIN);
123 }
124
125 #[test]
126 fn from_f64_roundtrip_approx() {
127 let key = QuantKey::from_f64(0.5, 7);
128 assert!(
129 (key.score() - 0.5_f32).abs() < 1e-5,
130 "score was {}",
131 key.score()
132 );
133 assert_eq!(key.id_prefix(), 7);
134 }
135}