pgm_extra/index/
segment.rs

1use 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}