Skip to main content

ltrait_sorter_frecency/
lib.rs

1//! This crate implement NewFrecency algorithm for ltrait.
2//! See also [User:Jesse/NewFrecency on mozilla wiki](https://wiki.mozilla.org/User:Jesse/NewFrecency?title=User:Jesse/NewFrecency)
3//!
4//! Create a database on `<XDG_DATA_HOME>/ltrait/frecency/frecency.sqlite`
5//!
6//! # Example Usage
7//! ```rust
8//! # use ltrait_sorter_frecency::{FrecencyConfig, Frecency};
9//! # use ltrait::{color_eyre::Result, Launcher};
10//! # use std::time::Duration;
11//! #
12//! # struct DummyUI;
13//! #
14//! # impl<'a> ltrait::UI<'a> for DummyUI {
15//! #     type Context = ();
16//! #
17//! #     async fn run<Cushion: 'a + Send>(
18//! #         &self,
19//! #         _: ltrait::launcher::batcher::Batcher<'a, Cushion, Self::Context>,
20//! #     ) -> Result<Option<Cushion>> {
21//! #         unimplemented!()
22//! #     }
23//! # }
24//! #
25//! # fn main() -> Result<()> {
26//! #
27//! let config = FrecencyConfig {
28//!     half_life: Duration::from_secs(30 * 60 * 60 * 24), // One month
29//!     type_ident: "simple-test".into(),
30//! };
31//!
32//! let launcher = Launcher::default()
33//!     .set_ui(DummyUI, |c| unimplemented!())
34//!     // Cushion is ltrait_sorter_frecency::Context
35//!     .add_raw_sorter(Frecency::new(config.clone())?)
36//!     .add_action(Frecency::new(config)?, |c|
37//!         ltrait_sorter_frecency::Context {
38//!             ident: c.ident.clone(),
39//!             bonus: 15.,
40//!         },
41//!     );
42//! #
43//! # Ok(()) }
44//! ```
45
46use std::collections::HashMap;
47use std::path::PathBuf;
48use std::time::Duration;
49
50use chrono::{DateTime, Utc};
51use ltrait::color_eyre::eyre::{OptionExt, Result, WrapErr};
52use ltrait::{Action, Sorter};
53use rusqlite::{Connection, params};
54
55/// The context of ltrait-sorter-frecency
56/// The `ident` must be unique within the same type_ident (if you want it to be judged as different)
57///
58/// If `bonus` is 0 and it is the first visit, the final score will also be 0 and will not increase. Set the `bonus` appropriately
59/// I don't know how much is optimal, so you'll have to try different things for a while.
60#[derive(Debug, Clone)]
61pub struct Context {
62    pub ident: String,
63    pub bonus: f64,
64}
65
66#[derive(Debug, Clone)]
67pub struct FrecencyConfig {
68    pub half_life: Duration,
69    pub type_ident: String,
70}
71
72#[derive(Debug, Clone)]
73struct Entry {
74    ident: String,
75    pub(crate) score: f64,
76
77    date: DateTime<Utc>,
78}
79
80impl Entry {
81    fn new(ident: String) -> Self {
82        Self {
83            ident,
84            score: 0.,
85            date: Utc::now(),
86        }
87    }
88
89    fn update(mut self, ctx: &Context, config: &FrecencyConfig) -> Self {
90        let ln2 = (2f64).ln();
91        let now = Utc::now();
92        let diff = now.signed_duration_since(self.date);
93
94        self.score = {
95            self.score
96                * (-(ln2 / (config.half_life.as_secs_f64() / 3600f64)) // as hour_f64
97                    * (diff.num_minutes() as f64 / (60f64)))
98                    .exp()
99                + ctx.bonus
100        };
101        self.date = now;
102
103        self
104    }
105}
106
107pub struct Frecency {
108    entries: HashMap<String, Entry>,
109    config: FrecencyConfig,
110}
111
112impl Frecency {
113    pub fn new(config: FrecencyConfig) -> Result<Self> {
114        Ok(Self {
115            entries: {
116                let conn = new_conn()?;
117                // type_ident
118                let mut stmt = conn
119                    .prepare(
120                        "SELECT ident, score, date FROM frecency_entries where type_ident = ?1",
121                    )
122                    .unwrap();
123
124                let entries = stmt
125                    .query_map([&config.type_ident], |row| {
126                        Ok(Entry {
127                            ident: row.get(0).unwrap(),
128                            score: row.get(1).unwrap(),
129                            date: row.get(2).unwrap(),
130                        })
131                    })
132                    .unwrap();
133
134                let mut v = HashMap::new();
135
136                for ei in entries {
137                    let ei: Entry = ei.unwrap();
138                    v.insert(ei.ident.clone(), ei);
139                }
140
141                v
142            },
143            config,
144        })
145    }
146}
147
148impl Sorter for Frecency {
149    type Context = Context;
150
151    fn compare(&self, lhs: &Self::Context, rhs: &Self::Context, _: &str) -> std::cmp::Ordering {
152        ((self.entries.get(&lhs.ident))
153            .map(|e| e.score)
154            .unwrap_or_default())
155        .partial_cmp(
156            &(self.entries.get(&rhs.ident))
157                .map(|e| e.score)
158                .unwrap_or_default(),
159        )
160        .unwrap_or(std::cmp::Ordering::Equal)
161    }
162}
163
164impl Action for Frecency {
165    type Context = Context;
166
167    fn act(&self, ctx: &Self::Context) -> Result<()> {
168        // self.entries
169        let new_entry = if self.entries.contains_key(&ctx.ident) {
170            self.entries
171                .get(&ctx.ident)
172                .unwrap()
173                .clone()
174                .update(ctx, &self.config)
175        } else {
176            Entry::new(ctx.ident.clone())
177        };
178
179        let conn = new_conn()?;
180
181        conn.execute(
182            "INSERT INTO frecency_entries (type_ident, ident, score, date)
183            VALUES (?1, ?2, ?3, ?4)
184            ON CONFLICT(type_ident, ident) DO UPDATE SET
185              score = EXCLUDED.score,
186              date = EXCLUDED.date;",
187            params![
188                &self.config.type_ident,
189                &new_entry.ident,
190                new_entry.score,
191                new_entry.date,
192            ],
193        )
194        .wrap_err("Failed to insert or update")?;
195
196        Ok(())
197    }
198}
199
200fn new_conn() -> Result<Connection> {
201    fn db_dir() -> Option<PathBuf> {
202        let path = dirs::data_dir().map(|p| p.join("ltrait/frecency/frecency.sqlite"));
203
204        if let Some(parent) = path.as_ref().and_then(|p| p.parent()) {
205            std::fs::create_dir_all(parent).ok()?;
206        }
207
208        path
209    }
210
211    let conn = Connection::open(db_dir().ok_or_eyre("Failed to get the path to store db")?)
212        .wrap_err("Failed to open a connectioin")?;
213
214    conn.execute(
215        r"CREATE TABLE IF NOT EXISTS frecency_entries (
216            id INTEGER PRIMARY KEY,
217            type_ident TEXT NOT NULL,
218            ident TEXT NOT NULL,
219            score REAL NOT NULL,
220            date TEXT NOT NULL,
221            UNIQUE(type_ident, ident)
222        );",
223        [],
224    )
225    .wrap_err("Failed to create new table")?;
226
227    Ok(conn)
228}