pgm_extra/index/
segment.rs1use crate::index::Key;
2
3#[repr(C)]
4#[derive(Clone, Copy, Debug)]
5#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
6pub struct Segment<K: Key> {
7 pub key: K,
8 pub slope: f64,
9 pub intercept: f64,
10}
11
12impl<K: Key> Segment<K> {
13 #[inline]
14 pub fn new(key: K, slope: f64, intercept: f64) -> Self {
15 Self {
16 key,
17 slope,
18 intercept,
19 }
20 }
21
22 #[inline(always)]
23 pub fn predict(&self, key: K) -> usize {
24 let diff = key.to_f64_fast() - self.key.to_f64_fast();
25 let pos = self.intercept + self.slope * diff;
26 pos.max(0.0) as usize
27 }
28
29 #[inline(always)]
30 pub fn predict_f64(&self, key: K) -> f64 {
31 let diff = key.to_f64_fast() - self.key.to_f64_fast();
32 self.intercept + self.slope * diff
33 }
34}
35
36impl<K: Key> Default for Segment<K> {
37 fn default() -> Self {
38 Self {
39 key: K::default(),
40 slope: 0.0,
41 intercept: 0.0,
42 }
43 }
44}
45
46#[cfg(test)]
47mod tests {
48 use super::*;
49
50 #[test]
51 fn test_segment_predict() {
52 let seg = Segment::new(0u64, 1.0, 0.0);
53 assert_eq!(seg.predict(0), 0);
54 assert_eq!(seg.predict(10), 10);
55 assert_eq!(seg.predict(100), 100);
56 }
57
58 #[test]
59 fn test_segment_with_intercept() {
60 let seg = Segment::new(10u64, 0.5, 5.0);
61 assert_eq!(seg.predict(10), 5);
62 assert_eq!(seg.predict(20), 10);
63 }
64
65 #[test]
66 fn test_segment_size() {
67 assert_eq!(core::mem::size_of::<Segment<u64>>(), 24);
68 assert_eq!(core::mem::size_of::<Segment<u32>>(), 24);
69 }
70}