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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
|
use std::path::Path;
use crate::atuin_client::{database::ClientSqlite, history::History, settings::FilterMode};
use async_trait::async_trait;
use eyre::Result;
use fuzzy_matcher::{FuzzyMatcher, skim::SkimMatcherV2};
use itertools::Itertools;
use time::OffsetDateTime;
use tokio::task::yield_now;
use tracing::{Level, instrument, warn};
use super::{SearchEngine, SearchState};
pub(crate) struct Search {
all_history: Vec<(History, i32)>,
engine: SkimMatcherV2,
}
impl Search {
pub(crate) fn new() -> Self {
Self {
all_history: vec![],
engine: SkimMatcherV2::default(),
}
}
}
#[async_trait]
impl SearchEngine for Search {
#[instrument(skip_all, level = Level::TRACE, name = "skim_search", fields(query = %state.input.as_str()))]
async fn full_query(
&mut self,
state: &SearchState,
db: &mut ClientSqlite,
) -> Result<Vec<History>> {
if self.all_history.is_empty() {
self.all_history = load_all_history(db).await;
}
Ok(fuzzy_search(&self.engine, state, &self.all_history).await)
}
#[instrument(skip_all, level = Level::TRACE, name = "skim_highlight")]
fn get_highlight_indices(&self, command: &str, search_input: &str) -> Vec<usize> {
let (_, indices) = self
.engine
.fuzzy_indices(command, search_input)
.unwrap_or_default();
indices
}
}
#[instrument(skip_all, level = Level::TRACE, name = "load_all_history")]
async fn load_all_history(db: &ClientSqlite) -> Vec<(History, i32)> {
db.all_with_count().await.unwrap()
}
#[expect(clippy::too_many_lines)]
#[instrument(skip_all, level = Level::TRACE, name = "fuzzy_match", fields(history_count = all_history.len()))]
async fn fuzzy_search(
engine: &SkimMatcherV2,
state: &SearchState,
all_history: &[(History, i32)],
) -> Vec<History> {
let mut set = Vec::with_capacity(200);
let mut ranks = Vec::with_capacity(200);
let query = state.input.as_str();
let now = OffsetDateTime::now_utc();
for (i, (history, count)) in all_history.iter().enumerate() {
if i % 256 == 0 {
yield_now().await;
}
let context = &state.context;
let git_root = context
.git_root
.as_ref()
.and_then(|git_root| git_root.to_str())
.unwrap_or(&context.cwd);
match state.filter_mode {
FilterMode::Global => {}
// we aggregate host by ',' separating them
FilterMode::Host
if history
.hostname
.split(',')
.contains(&context.hostname.as_str()) => {}
// we aggregate session by concattenating them.
// sessions are 32 byte simple uuid formats
FilterMode::Session
if history
.session
.as_bytes()
.chunks(32)
.contains(&context.session.as_bytes()) => {}
// SessionPreload: include current session + global history from before session start
FilterMode::SessionPreload => {
let is_current_session = {
history
.session
.as_bytes()
.chunks(32)
.any(|chunk| chunk == context.session.as_bytes())
};
if !is_current_session {
let Ok(uuid) = uuid::Uuid::parse_str(&context.session) else {
warn!("failed to parse session id '{}'", context.session);
continue;
};
let Some(timestamp) = uuid.get_timestamp() else {
warn!(
"failed to get timestamp from uuid '{}'",
uuid.as_hyphenated()
);
continue;
};
let (seconds, nanos) = timestamp.to_unix();
let Ok(session_start) = OffsetDateTime::from_unix_timestamp_nanos(
i128::from(seconds) * 1_000_000_000 + i128::from(nanos),
) else {
warn!(
"failed to create OffsetDateTime from second: {seconds}, nanosecond: {nanos}"
);
continue;
};
if history.timestamp >= session_start {
continue;
}
}
}
// we aggregate directory by ':' separating them
FilterMode::Directory if history.cwd.split(':').contains(&context.cwd.as_str()) => {}
FilterMode::Workspace if history.cwd.split(':').contains(&git_root) => {}
_ => continue,
}
#[expect(clippy::cast_lossless, clippy::cast_precision_loss)]
if let Some((score, indices)) = engine.fuzzy_indices(&history.command, query) {
let begin = indices.first().copied().unwrap_or_default();
let mut duration = (now - history.timestamp).as_seconds_f64().log2();
if !duration.is_finite() || duration <= 1.0 {
duration = 1.0;
}
// these + X.0 just make the log result a bit smoother.
// log is very spiky towards 1-4, but I want a gradual decay.
// eg:
// log2(4) = 2, log2(5) = 2.3 (16% increase)
// log2(8) = 3, log2(9) = 3.16 (5% increase)
// log2(16) = 4, log2(17) = 4.08 (2% increase)
let count = (*count as f64 + 8.0).log2();
let begin = (begin as f64 + 16.0).log2();
let path = path_dist(history.cwd.as_ref(), state.context.cwd.as_ref());
let path = (path as f64 + 8.0).log2();
// reduce longer durations, raise higher counts, raise matches close to the start
let score = (-score as f64) * count / path / duration / begin;
'insert: {
// algorithm:
// 1. find either the position that this command ranks
// 2. find the same command positioned better than our rank.
for i in 0..set.len() {
// do we out score the current position?
if ranks[i] > score {
ranks.insert(i, score);
set.insert(i, history.clone());
let mut j = i + 1;
while j < set.len() {
// remove duplicates that have a worse score
if set[j].command == history.command {
ranks.remove(j);
set.remove(j);
// break this while loop because there won't be any other
// duplicates.
break;
}
j += 1;
}
// keep it limited
if ranks.len() > 200 {
ranks.pop();
set.pop();
}
break 'insert;
}
// don't continue if this command has a better score already
if set[i].command == history.command {
break 'insert;
}
}
if set.len() < 200 {
ranks.push(score);
set.push(history.clone());
}
}
}
}
set
}
fn path_dist(a: &Path, b: &Path) -> usize {
let mut a: Vec<_> = a.components().collect();
let b: Vec<_> = b.components().collect();
let mut dist = 0;
// pop a until there's a common ancestor
while !b.starts_with(&a) {
dist += 1;
a.pop();
}
b.len() - a.len() + dist
}
|
