ltrait_sorter_frecency/
lib.rs1use 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#[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)) * (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 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 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}