pgm_extra/index/
segment.rs

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

pgm_extra/index/
segment.rs
