moros/src/chuva.rs

Rename chuva::Model to chuva::Chuva 💬 by Caio 7 months ago (log)

Consumed my persim berry

Blob moros/src/chuva.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131

use std::path::Path;

use fst::{Automaton, IntoStreamer, Streamer};
use jiff::Timestamp;

use chuva::{MAX_OFFSET, Model, Projector, STEPS};

type Result<T> = crate::Result<T>;

pub type Prediction<'a> = &'a [f32; STEPS];

pub struct Moros {
    model: Model,
    proj: Projector,
    fst: fst::Map<&'static [u8]>,
}

impl Moros {
    pub fn load_from_dir<P: AsRef<Path>>(dir: P) -> Result<Self> {
        let model = Model::load_from_dir(dir)?;
        let fst = fst::Map::new(FST_STATE)?;

        Ok(Self {
            proj: Projector::new(),
            fst,
            model,
        })
    }

    pub fn by_lat_lon(&self, lat: f64, lon: f64) -> Option<Prediction<'_>> {
        let offset = self.proj.to_offset(lat, lon)?;
        self.by_offset(offset)
    }

    pub fn by_postcode(&self, code: &str) -> Option<Prediction<'_>> {
        let mut stream = self
            .fst
            .search(AsciiUpperCase::new(code).starts_with())
            .into_stream();
        let (_, offset) = stream.next()?;
        self.by_offset(offset as usize)
    }

    pub fn by_postcode4(&self, code: &str) -> Option<Prediction<'_>> {
        let mut stream = self.fst.range().gt(code).into_stream();
        let (key, offset) = stream.next()?;
        assert_eq!(6, key.len(), "key is pc6");
        if &key[..4] == code.as_bytes() {
            self.by_offset(offset as usize)
        } else {
            None
        }
    }

    #[inline]
    pub(crate) fn by_offset(&self, offset: usize) -> Option<Prediction<'_>> {
        assert!(offset <= MAX_OFFSET);
        Some(
            self.model.data[offset..(offset + STEPS)]
                .try_into()
                .unwrap(),
        )
    }

    pub fn created_at(&self) -> Timestamp {
        self.model.created_at
    }

    pub fn filename(&self) -> &str {
        &self.model.filename
    }

    pub fn kind(&self) -> chuva::ModelKind {
        self.model.kind
    }

    pub fn get_time_slot(&self, now: Timestamp) -> Result<usize> {
        get_time_slot(self.model.created_at, now).map_err(|_| "Dataset too old".into())
    }
}

static FST_STATE: &[u8] = include_bytes!("../asset/postcodes.fst").as_slice();

fn get_time_slot(created_at: Timestamp, now: Timestamp) -> std::result::Result<usize, i64> {
    let age = (now - created_at)
        .total(jiff::Unit::Minute)
        .map_err(|_| 420)?; // irrelevant

    if !(0.0..=120.0).contains(&age) {
        Err(age as i64)
    } else {
        let slot = (age / 5.0) as usize;
        assert!(slot < STEPS);
        Ok(slot)
    }
}

#[cfg(test)]
mod tests {
    use super::{AsciiUpperCase, FST_STATE, get_time_slot};

    use fst::{Automaton, IntoStreamer, Streamer};
    use jiff::{Timestamp, ToSpan};

    #[test]
    fn slot_works() {
        let now = Timestamp::now();

        assert_eq!(Err(-1), get_time_slot(now, now - 1.minute()));
        assert_eq!(Ok(0), get_time_slot(now, now));
        assert_eq!(Ok(24), get_time_slot(now, now + 2.hours()));
        assert_eq!(Err(121), get_time_slot(now, now + 121.minutes()));
    }

    #[test]
    fn case_insensitive_postcode_search() {
        let fst = fst::Map::new(FST_STATE).expect("valid fst state");

        let _ = fst.get("1017CE").expect("Key 1017CE exists in the fst");

        let mut stream = fst
            .search(AsciiUpperCase::new("1017ce").starts_with())
            .into_stream();
        let (key, _) = stream.next().expect("lower case search matches");

        assert_eq!(
            b"1017CE", key,
            "lower case search should match upper case key"
        );
    }
}