use std::collections::HashMap;
use std::sync::Arc;
use rusqlite::{Connection, params};
use crate::db::connection::DbError;
use super::TrackRow;
pub(crate) fn sanitize_fts_query(query: &str) -> String {
let trimmed = query.trim();
if trimmed.is_empty() {
return String::new();
}
let escaped = trimmed.replace('"', "\"\"");
format!("\"{}\"*", escaped)
}
pub fn search_tracks(conn: &Connection, query: &str) -> Result<Vec<TrackRow>, DbError> {
search_tracks_paged(conn, query, 100, 0)
}
pub fn search_tracks_paged(
conn: &Connection,
query: &str,
limit: u32,
offset: u32,
) -> Result<Vec<TrackRow>, DbError> {
let fts_query = sanitize_fts_query(query);
let mut stmt = conn.prepare(
"SELECT t.id, t.album_id, t.artist_id, a.name, aa.name, al.title,
t.disc, t.track_number, t.title, t.duration_ms, t.path,
t.codec, t.sample_rate, t.bit_depth, t.channels, t.bitrate,
t.genre, t.source, t.remote_id, t.cached_path
FROM tracks_fts f
JOIN tracks t ON t.id = f.rowid
LEFT JOIN artists a ON t.artist_id = a.id
LEFT JOIN albums al ON t.album_id = al.id
LEFT JOIN artists aa ON al.artist_id = aa.id
WHERE tracks_fts MATCH ?1
ORDER BY a.name COLLATE LIBRARY, al.date, al.title COLLATE LIBRARY, t.disc, t.track_number
LIMIT ?2 OFFSET ?3",
)?;
let rows = stmt
.query_map(params![fts_query, limit, offset], super::row_to_track_row)?
.collect::<Result<Vec<_>, _>>()?;
Ok(rows)
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum CorpusKind {
Track,
Album,
Artist,
}
pub type Corpus = Arc<Vec<(i64, String)>>;
pub fn fuzzy_corpus(conn: &Connection, kind: CorpusKind) -> Result<Vec<(i64, String)>, DbError> {
let mut stmt = conn.prepare_cached(match kind {
CorpusKind::Track => {
"SELECT t.id, COALESCE(a.name, '') || ' — ' || COALESCE(al.title, '') || ' — ' || t.title
FROM tracks t
LEFT JOIN artists a ON t.artist_id = a.id
LEFT JOIN albums al ON t.album_id = al.id"
}
CorpusKind::Album => {
"SELECT al.id, COALESCE(a.name, '') || ' — ' || al.title
FROM albums al
LEFT JOIN artists a ON al.artist_id = a.id"
}
CorpusKind::Artist => {
"SELECT a.id, a.name FROM artists a
WHERE EXISTS (SELECT 1 FROM albums al WHERE al.artist_id = a.id)"
}
})?;
let rows = stmt.query_map([], |row| Ok((row.get(0)?, row.get(1)?)))?;
rows.collect::<Result<Vec<_>, _>>().map_err(Into::into)
}
pub(crate) fn match_rank(column: &str) -> String {
format!(
"CASE WHEN {column} LIKE ? ESCAPE '\\' THEN 0
WHEN {column} LIKE ? ESCAPE '\\' THEN 1
WHEN {column} LIKE ? ESCAPE '\\' THEN 2 ELSE 3 END"
)
}
pub(crate) fn match_rank_binds(query: &str) -> [String; 3] {
let q = super::artists::escape_like(query.trim());
[q.clone(), format!("{q}%"), format!("% {q}%")]
}
const FUZZY_FLOOR: f64 = 0.6;
const FUZZY_LIMIT: usize = 200;
pub fn fuzzy_ids(conn: &Connection, kind: CorpusKind, query: &str) -> Result<Vec<i64>, DbError> {
use nucleo::pattern::{AtomKind, CaseMatching, Normalization, Pattern};
use nucleo::{Config, Matcher, Utf32Str};
let stamp = Stamp::of(conn)?;
let corpus = {
let slots = FALLBACK.lock();
match slots.get(&kind) {
Some(slot) if slot.stamp == stamp => {
if let Some((q, ids)) = &slot.last
&& q == query
{
return Ok(ids.clone());
}
Some(Arc::clone(&slot.corpus))
}
_ => None,
}
};
let corpus = match corpus {
Some(corpus) => corpus,
None => Arc::new(fuzzy_corpus(conn, kind)?),
};
let pattern = Pattern::new(
query,
CaseMatching::Ignore,
Normalization::Smart,
AtomKind::Fuzzy,
);
let mut matcher = Matcher::new(Config::DEFAULT);
let mut buf = Vec::new();
let mut scored: Vec<(u32, usize, i64)> = corpus
.iter()
.filter_map(|(id, text)| {
pattern
.score(Utf32Str::new(text, &mut buf), &mut matcher)
.map(|score| (score, text.len(), *id))
})
.collect();
scored.sort_by_key(|&(score, len, id)| (std::cmp::Reverse(score), len, id));
let ids: Vec<i64> = match scored.first() {
Some(&(best, ..)) => {
let floor = (best as f64 * FUZZY_FLOOR) as u32;
scored
.into_iter()
.take_while(|&(score, ..)| score >= floor)
.take(FUZZY_LIMIT)
.map(|(.., id)| id)
.collect()
}
None => Vec::new(),
};
FALLBACK.lock().insert(
kind,
Fallback {
stamp,
corpus,
last: Some((query.to_owned(), ids.clone())),
},
);
Ok(ids)
}
struct Fallback {
stamp: Stamp,
corpus: Corpus,
last: Option<(String, Vec<i64>)>,
}
static FALLBACK: std::sync::LazyLock<parking_lot::Mutex<HashMap<CorpusKind, Fallback>>> =
std::sync::LazyLock::new(Default::default);
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct Stamp {
connection: i64,
data_version: i64,
changes: u64,
}
impl Stamp {
fn of(conn: &Connection) -> Result<Self, DbError> {
conn.execute_batch(
"CREATE TEMP TABLE IF NOT EXISTS koan_connection (id INTEGER NOT NULL);
INSERT INTO koan_connection
SELECT random() WHERE NOT EXISTS (SELECT 1 FROM koan_connection);",
)?;
let connection = conn.query_row("SELECT id FROM koan_connection", [], |r| r.get(0))?;
let data_version = conn.query_row("PRAGMA data_version", [], |r| r.get(0))?;
Ok(Self {
connection,
data_version,
changes: conn.total_changes(),
})
}
}
#[derive(Default)]
pub struct CorpusCache {
slots: parking_lot::Mutex<HashMap<CorpusKind, (u64, Corpus)>>,
}
impl CorpusCache {
pub fn get(
&self,
conn: &Connection,
kind: CorpusKind,
version: u64,
) -> Result<Corpus, DbError> {
if let Some((at, corpus)) = self.slots.lock().get(&kind)
&& *at == version
{
return Ok(Arc::clone(corpus));
}
let corpus: Corpus = Arc::new(fuzzy_corpus(conn, kind)?);
self.slots
.lock()
.insert(kind, (version, Arc::clone(&corpus)));
Ok(corpus)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::db::connection::Database;
use crate::db::queries::{sample_meta, upsert_track};
fn test_db() -> Database {
let conn = rusqlite::Connection::open_in_memory().unwrap();
conn.pragma_update(None, "foreign_keys", "on").unwrap();
crate::db::schema::create_tables(&conn).unwrap();
Database { conn }
}
#[test]
fn test_search_fts() {
let db = test_db();
upsert_track(&db.conn, &sample_meta("Vordhosbn", "Aphex Twin", "Drukqs")).unwrap();
upsert_track(
&db.conn,
&sample_meta("Roygbiv", "Boards of Canada", "MHTRTC"),
)
.unwrap();
upsert_track(
&db.conn,
&sample_meta("Tha", "Aphex Twin", "Selected Ambient Works"),
)
.unwrap();
let results = search_tracks(&db.conn, "Aphex").unwrap();
assert_eq!(results.len(), 2);
let results = search_tracks(&db.conn, "Roygbiv").unwrap();
assert_eq!(results.len(), 1);
assert_eq!(results[0].title, "Roygbiv");
let results = search_tracks(&db.conn, "Drukqs").unwrap();
assert_eq!(results.len(), 1);
}
}