use anyhow::Result;
use rusqlite::types::Value;
use rusqlite::{Connection, OptionalExtension, params};
use std::path::Path;
pub(crate) const SCHEMA: &str = r#"
CREATE TABLE IF NOT EXISTS meta (
k TEXT PRIMARY KEY,
v TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS files (
id INTEGER PRIMARY KEY,
path TEXT NOT NULL UNIQUE,
hash TEXT NOT NULL,
bytes INTEGER NOT NULL,
indexed_at INTEGER NOT NULL
);
-- AUTOINCREMENT, not a bare rowid: a sharded index routes a chunk id to a
-- shard by the range its file name claims, so SQLite handing a deleted row's
-- id to a new chunk would put one id inside two shards at once.
CREATE TABLE IF NOT EXISTS chunks (
id INTEGER PRIMARY KEY AUTOINCREMENT,
file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
ord INTEGER NOT NULL,
start_line INTEGER NOT NULL,
end_line INTEGER NOT NULL,
text TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS chunks_file_id ON chunks(file_id);
-- Keyword half of the search. `content='chunks'` means FTS5 keeps only its
-- index, not a second copy of every chunk, so the store barely grows.
CREATE VIRTUAL TABLE IF NOT EXISTS chunks_fts USING fts5(
text,
content='chunks',
content_rowid='id',
tokenize='unicode61'
);
-- External-content FTS5 does not track its source table by itself. These keep
-- the two in step; a delete has to hand back the original text so FTS5 can
-- find the terms it needs to remove. Foreign-key cascades fire them too, which
-- is what keeps `delete_file` correct.
CREATE TRIGGER IF NOT EXISTS chunks_fts_insert AFTER INSERT ON chunks BEGIN
INSERT INTO chunks_fts(rowid, text) VALUES (new.id, new.text);
END;
CREATE TRIGGER IF NOT EXISTS chunks_fts_delete AFTER DELETE ON chunks BEGIN
INSERT INTO chunks_fts(chunks_fts, rowid, text) VALUES('delete', old.id, old.text);
END;
-- Structure half of the store, from 0.12.0. Additive: an older store grows
-- these tables on its next open and they stay empty until something indexes
-- into it, which is why the format version does not move for them. See
-- `docs/compatibility.md`.
--
-- `chunk_id` is nullable on purpose. A symbol's defining line always falls
-- inside some chunk, but a file whose chunks were rewritten between extraction
-- and insertion would otherwise have to fail rather than record the symbol.
CREATE TABLE IF NOT EXISTS symbols (
id INTEGER PRIMARY KEY AUTOINCREMENT,
file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
chunk_id INTEGER,
kind TEXT NOT NULL,
name TEXT NOT NULL,
qualified TEXT NOT NULL,
start_line INTEGER NOT NULL,
end_line INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS symbols_file_id ON symbols(file_id);
CREATE INDEX IF NOT EXISTS symbols_name ON symbols(name);
-- An edge is owned by the file its *source* is in: `src` is a symbol id, and
-- the cascade above deletes an edge when its source file is re-indexed or
-- forgotten.
--
-- `dst` is a symbol NAME, not an id, and that asymmetry is the point. Ids are
-- reissued every time a file is re-extracted, so an id here would mean that
-- re-indexing `b.rs` silently deleted every edge pointing into it from `a.rs`
-- — the graph would rot from the one operation this release exists to make
-- safe. A name is resolved against `symbols_name` at query time instead, so an
-- edge is always as current as both of its ends, and an edge to something not
-- indexed (a standard-library call) is still recorded and simply resolves to
-- nothing.
--
-- `confidence` is `extracted` when an import in the source file named where the
-- target came from, and `inferred` when the target was matched by bare name.
-- Nothing that displays an edge may present the second as the first.
-- `hint` is what the source text said about where `dst` lives, when it said
-- anything: the module of a scoped call, the receiver of a method call, the
-- object of a qualified one. Nullable, and NULL on every row an older binary
-- wrote, which is why the format version does not move for it. It is a lead,
-- not an address, and `edges_out` treats it as one ranking signal among
-- several — see `resolve`.
-- `line` is where the reference was written in the source file, from 0.16.0.
-- Not where either symbol is defined: "where is X invoked" is a question about
-- this column, and until it existed the closest answer was the enclosing
-- symbol's own start line. Nullable and NULL on every row an older binary
-- wrote, so the format version does not move for it either.
CREATE TABLE IF NOT EXISTS edges (
src INTEGER NOT NULL REFERENCES symbols(id) ON DELETE CASCADE,
dst TEXT NOT NULL,
kind TEXT NOT NULL,
confidence TEXT NOT NULL,
hint TEXT,
line INTEGER
);
CREATE INDEX IF NOT EXISTS edges_src ON edges(src);
CREATE INDEX IF NOT EXISTS edges_dst ON edges(dst);
-- Images, from 0.13.0. One row per indexed image file, with the pixel size the
-- Search and Files pages show in place of a line range. The vector itself is in
-- the store's second index, under `images/`, at CLIP's 512 dimensions.
--
-- Additive and `IF NOT EXISTS`, so `format_version` does not move: an older
-- binary opens the store, never looks in here, and searches its text exactly as
-- before. Same reasoning `docs/compatibility.md` records for the graph tables.
--
-- `id` is AUTOINCREMENT for the same reason chunks are: the id addresses a
-- vector in a sharded index, and SQLite reissuing a deleted row's id would put
-- one id inside two shards at once.
CREATE TABLE IF NOT EXISTS images (
id INTEGER PRIMARY KEY AUTOINCREMENT,
file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
width INTEGER NOT NULL,
height INTEGER NOT NULL
);
CREATE INDEX IF NOT EXISTS images_file_id ON images(file_id);
-- The retrieval ledger, from 0.12.0. Additive and empty unless recording is
-- switched on, which it is not by default.
--
-- `prev` is the hash of the row before it and `hash` covers this row's own
-- fields plus `prev`, so the table is a chain: editing or deleting a row
-- breaks every hash after it and `semlith ledger` can say so. That is what
-- makes this an audit record rather than a log file, and it costs one blake3
-- of a short string per query.
--
-- `excerpt_tokens` is what the agent actually read. `whole_file_tokens` is
-- what reading those files whole would have cost — the honest denominator for
-- a savings number, measured rather than claimed.
CREATE TABLE IF NOT EXISTS retrievals (
id INTEGER PRIMARY KEY AUTOINCREMENT,
at INTEGER NOT NULL,
client TEXT NOT NULL,
query TEXT NOT NULL,
hits INTEGER NOT NULL,
micros INTEGER NOT NULL,
excerpt_tokens INTEGER NOT NULL,
whole_file_tokens INTEGER NOT NULL,
prev TEXT NOT NULL,
hash TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS retrievals_at ON retrievals(at);
"#;
#[derive(Debug, Clone)]
pub struct ChunkRow {
pub id: i64,
pub path: String,
pub ord: i64,
pub start_line: u32,
pub end_line: u32,
pub text: String,
}
pub fn open(path: &Path) -> Result<Connection> {
let db = Connection::open(path)?;
db.pragma_update(None, "journal_mode", "WAL")?;
db.pragma_update(None, "synchronous", "NORMAL")?;
db.pragma_update(None, "foreign_keys", "ON")?;
defensive(&db)?;
db.execute_batch(SCHEMA)?;
add_columns(&db)?;
check_format(&db)?;
backfill_fts(&db)?;
read_only(&db, true)?;
Ok(db)
}
fn defensive(db: &Connection) -> Result<()> {
db.pragma_update(None, "trusted_schema", "OFF")?;
db.set_db_config(rusqlite::config::DbConfig::SQLITE_DBCONFIG_DEFENSIVE, true)?;
Ok(())
}
pub fn read_only(db: &Connection, on: bool) -> Result<()> {
db.pragma_update(None, "query_only", if on { "ON" } else { "OFF" })?;
Ok(())
}
pub struct Writing<'a>(&'a Connection);
impl<'a> Writing<'a> {
pub fn begin(db: &'a Connection) -> Result<Self> {
read_only(db, false)?;
Ok(Self(db))
}
}
impl Drop for Writing<'_> {
fn drop(&mut self) {
let _ = read_only(self.0, true);
}
}
pub const FORMAT_VERSION: u32 = 2;
pub const SHARDED_FORMAT: u32 = 2;
pub fn format(db: &Connection) -> Result<u32> {
Ok(get_meta(db, FORMAT_KEY)?
.and_then(|v| v.parse().ok())
.unwrap_or(1))
}
pub const FORMAT_KEY: &str = "format_version";
fn check_format(db: &Connection) -> Result<()> {
let Some(recorded) = get_meta(db, FORMAT_KEY)? else {
return Ok(());
};
let found: u32 = recorded.parse().map_err(|_| {
anyhow::anyhow!("store records {FORMAT_KEY} {recorded:?}, which is not a version number")
})?;
if found > FORMAT_VERSION {
anyhow::bail!(
"store is format {found}, and semlith {} understands format {FORMAT_VERSION}; \
upgrade semlith to read this store",
env!("CARGO_PKG_VERSION"),
);
}
Ok(())
}
const FTS_BUILT: &str = "fts_built";
fn add_columns(db: &Connection) -> Result<()> {
const ADDITIONS: [(&str, &str, &str); 6] = [
("edges", "hint", "TEXT"),
("edges", "line", "INTEGER"),
("retrievals", "session", "TEXT"),
("retrievals", "tool", "TEXT"),
("retrievals", "stale_hits", "INTEGER"),
("retrievals", "tokenizer", "TEXT"),
];
for (table, column, kind) in ADDITIONS {
if has_column(db, table, column)? {
continue;
}
db.execute_batch(&format!("ALTER TABLE {table} ADD COLUMN {column} {kind};"))?;
}
Ok(())
}
fn has_column(db: &Connection, table: &str, column: &str) -> Result<bool> {
let mut stmt = db.prepare(&format!("PRAGMA table_info({table})"))?;
let mut rows = stmt.query([])?;
while let Some(row) = rows.next()? {
let name: String = row.get(1)?;
if name == column {
return Ok(true);
}
}
Ok(false)
}
fn backfill_fts(db: &Connection) -> Result<()> {
if get_meta(db, FTS_BUILT)?.is_some() {
return Ok(());
}
db.execute_batch("INSERT INTO chunks_fts(chunks_fts) VALUES('rebuild');")?;
set_meta(db, FTS_BUILT, "1")?;
Ok(())
}
fn glob_predicate(groups: &[Vec<String>]) -> (String, Vec<String>) {
let mut binds = Vec::new();
let mut clauses = Vec::new();
for group in groups {
let ors: Vec<String> = group
.iter()
.map(|p| {
binds.push(p.clone());
format!("{GLOB_PATH} GLOB ?")
})
.collect();
clauses.push(format!("({})", ors.join(" OR ")));
}
if clauses.is_empty() {
return ("1".to_string(), binds);
}
(clauses.join(" AND "), binds)
}
pub fn filtered_chunk_ids(db: &Connection, groups: &[Vec<String>]) -> Result<Vec<u64>> {
let (predicate, binds) = glob_predicate(groups);
let sql =
format!("SELECT c.id FROM chunks c JOIN files f ON f.id = c.file_id WHERE {predicate}");
let mut stmt = db.prepare(&sql)?;
let args = binds.into_iter().map(Value::Text);
let rows = stmt.query_map(rusqlite::params_from_iter(args), |r| r.get::<_, i64>(0))?;
Ok(rows
.collect::<Result<Vec<i64>, _>>()?
.into_iter()
.map(|i| i as u64)
.collect())
}
pub fn matching_files(db: &Connection, groups: &[Vec<String>]) -> Result<i64> {
let (predicate, binds) = glob_predicate(groups);
let sql = format!("SELECT COUNT(*) FROM files f WHERE {predicate}");
let args = binds.into_iter().map(Value::Text);
Ok(db.query_row(&sql, rusqlite::params_from_iter(args), |r| r.get(0))?)
}
const UNFILTERED_SQL: &str =
"SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ? ORDER BY rank LIMIT ?";
pub fn keyword_search(
db: &Connection,
query: &str,
limit: usize,
groups: &[Vec<String>],
) -> Result<Vec<u64>> {
let terms: Vec<String> = query
.split(|c: char| !c.is_alphanumeric() && c != '_')
.filter(|t| !t.is_empty())
.map(|t| format!("\"{t}\""))
.collect();
if terms.is_empty() {
return Ok(Vec::new());
}
let match_expr = terms.join(" OR ");
let (sql, mut args) = if groups.is_empty() {
(UNFILTERED_SQL.to_string(), vec![Value::Text(match_expr)])
} else {
let (predicate, binds) = glob_predicate(groups);
let mut args = vec![Value::Text(match_expr)];
args.extend(binds.into_iter().map(Value::Text));
(
format!(
"SELECT x.rowid FROM chunks_fts x
JOIN chunks c ON c.id = x.rowid
JOIN files f ON f.id = c.file_id
WHERE x.chunks_fts MATCH ? AND {predicate}
ORDER BY x.rank LIMIT ?"
),
args,
)
};
args.push(Value::Integer(limit as i64));
let mut stmt = db.prepare(&sql)?;
let rows = stmt.query_map(rusqlite::params_from_iter(args), |r| r.get::<_, i64>(0))?;
Ok(rows
.collect::<Result<Vec<i64>, _>>()?
.into_iter()
.map(|i| i as u64)
.collect())
}
pub fn get_meta(db: &Connection, k: &str) -> Result<Option<String>> {
Ok(db
.query_row("SELECT v FROM meta WHERE k = ?1", params![k], |r| r.get(0))
.optional()?)
}
pub fn set_meta(db: &Connection, k: &str, v: &str) -> Result<()> {
db.execute(
"INSERT INTO meta (k, v) VALUES (?1, ?2)
ON CONFLICT(k) DO UPDATE SET v = excluded.v",
params![k, v],
)?;
Ok(())
}
pub fn file_hash(db: &Connection, path: &str) -> Result<Option<String>> {
Ok(db
.query_row(
"SELECT hash FROM files WHERE path = ?1",
params![path],
|r| r.get(0),
)
.optional()?)
}
pub fn delete_file(db: &Connection, path: &str) -> Result<Vec<u64>> {
let ids: Vec<u64> = {
let mut stmt = db.prepare(
"SELECT c.id FROM chunks c JOIN files f ON f.id = c.file_id WHERE f.path = ?1",
)?;
let rows = stmt.query_map(params![path], |r| r.get::<_, i64>(0))?;
rows.collect::<Result<Vec<i64>, _>>()?
.into_iter()
.map(|i| i as u64)
.collect()
};
db.execute("DELETE FROM files WHERE path = ?1", params![path])?;
Ok(ids)
}
pub fn insert_file(db: &Connection, path: &str, hash: &str, bytes: u64, now: i64) -> Result<i64> {
db.execute(
"INSERT INTO files (path, hash, bytes, indexed_at) VALUES (?1, ?2, ?3, ?4)",
params![path, hash, bytes as i64, now],
)?;
Ok(db.last_insert_rowid())
}
pub fn insert_chunk(
db: &Connection,
file_id: i64,
ord: usize,
start_line: u32,
end_line: u32,
text: &str,
) -> Result<i64> {
db.execute(
"INSERT INTO chunks (file_id, ord, start_line, end_line, text)
VALUES (?1, ?2, ?3, ?4, ?5)",
params![file_id, ord as i64, start_line, end_line, text],
)?;
Ok(db.last_insert_rowid())
}
pub fn files_ending_with(db: &Connection, suffix: &str, limit: usize) -> Result<Vec<String>> {
let suffix = suffix
.trim_start_matches(['.', '/', '\\'])
.replace('\\', "/");
let mut stmt = db.prepare(&format!(
"SELECT path FROM files WHERE path = ?1 OR {PATH_AS_TYPED} LIKE '%/' || ?1
ORDER BY LENGTH(path) LIMIT ?2",
))?;
let rows = stmt.query_map(params![suffix, limit as i64], |r| r.get::<_, String>(0))?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
#[cfg(windows)]
const PATH_AS_TYPED: &str = r#"replace(path, '\', '/')"#;
#[cfg(not(windows))]
const PATH_AS_TYPED: &str = "path";
#[cfg(windows)]
const GLOB_PATH: &str = r#"replace(lower(f.path), '\', '/')"#;
#[cfg(not(windows))]
const GLOB_PATH: &str = "lower(f.path)";
pub fn chunks_overlapping(
db: &Connection,
path: &str,
start: u32,
end: u32,
) -> Result<Vec<ChunkRow>> {
let mut stmt = db.prepare(
"SELECT c.id, f.path, c.ord, c.start_line, c.end_line, c.text
FROM chunks c JOIN files f ON f.id = c.file_id
WHERE f.path = ?1 AND c.start_line <= ?2 AND c.end_line >= ?3
ORDER BY c.ord",
)?;
let rows = stmt.query_map(params![path, end, start], |r| {
Ok(ChunkRow {
id: r.get(0)?,
path: r.get(1)?,
ord: r.get(2)?,
start_line: r.get(3)?,
end_line: r.get(4)?,
text: r.get(5)?,
})
})?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn chunk(db: &Connection, id: u64) -> Result<Option<ChunkRow>> {
Ok(db
.query_row(
"SELECT c.id, f.path, c.ord, c.start_line, c.end_line, c.text
FROM chunks c JOIN files f ON f.id = c.file_id
WHERE c.id = ?1",
params![id as i64],
|r| {
Ok(ChunkRow {
id: r.get(0)?,
path: r.get(1)?,
ord: r.get(2)?,
start_line: r.get(3)?,
end_line: r.get(4)?,
text: r.get(5)?,
})
},
)
.optional()?)
}
#[derive(Debug, Clone)]
pub struct FileRow {
pub path: String,
pub bytes: i64,
pub chunks: i64,
pub lines: i64,
pub indexed_at: i64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FileSort {
Path,
Bytes,
Chunks,
Lines,
Indexed,
}
impl FileSort {
pub fn parse(name: &str) -> Option<Self> {
match name {
"path" => Some(Self::Path),
"bytes" => Some(Self::Bytes),
"chunks" => Some(Self::Chunks),
"lines" => Some(Self::Lines),
"indexed" => Some(Self::Indexed),
_ => None,
}
}
fn sql(self) -> &'static str {
match self {
Self::Path => "f.path",
Self::Bytes => "f.bytes",
Self::Chunks => "COUNT(c.id)",
Self::Lines => "COALESCE(MAX(c.end_line), 0)",
Self::Indexed => "f.indexed_at",
}
}
}
pub fn file_count(db: &Connection, groups: &[Vec<String>]) -> Result<i64> {
let (predicate, binds) = glob_predicate(groups);
let sql = format!("SELECT COUNT(*) FROM files f WHERE {predicate}");
let mut stmt = db.prepare(&sql)?;
let args = binds.into_iter().map(Value::Text);
Ok(stmt.query_row(rusqlite::params_from_iter(args), |r| r.get(0))?)
}
pub fn file_rows(
db: &Connection,
groups: &[Vec<String>],
sort: FileSort,
desc: bool,
take: i64,
) -> Result<Vec<FileRow>> {
let (predicate, binds) = glob_predicate(groups);
let direction = if desc { "DESC" } else { "ASC" };
let sql = format!(
"SELECT f.path, f.bytes, COUNT(c.id), COALESCE(MAX(c.end_line), 0), f.indexed_at \
FROM files f LEFT JOIN chunks c ON c.file_id = f.id \
WHERE {predicate} GROUP BY f.id ORDER BY {} {direction}, f.path ASC LIMIT ?",
sort.sql()
);
let mut stmt = db.prepare(&sql)?;
let mut args: Vec<Value> = binds.into_iter().map(Value::Text).collect();
args.push(Value::Integer(take.max(0)));
let rows = stmt.query_map(rusqlite::params_from_iter(args), |r| {
Ok(FileRow {
path: r.get(0)?,
bytes: r.get(1)?,
chunks: r.get(2)?,
lines: r.get(3)?,
indexed_at: r.get(4)?,
})
})?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
#[derive(Debug, Default, Clone)]
pub struct Facets {
pub extensions: Vec<String>,
pub lines: i64,
}
pub fn file_facets(db: &Connection, groups: &[Vec<String>]) -> Result<Facets> {
let (predicate, binds) = glob_predicate(groups);
let sql = format!("SELECT f.path FROM files f WHERE {predicate}");
let mut stmt = db.prepare(&sql)?;
let args = binds.clone().into_iter().map(Value::Text);
let paths = stmt.query_map(rusqlite::params_from_iter(args), |r| r.get::<_, String>(0))?;
let mut extensions: Vec<String> = paths
.collect::<Result<Vec<String>, _>>()?
.iter()
.filter_map(|p| {
std::path::Path::new(p)
.extension()
.and_then(|e| e.to_str())
.map(|e| e.to_ascii_lowercase())
})
.collect();
extensions.sort();
extensions.dedup();
let sql = format!(
"SELECT COALESCE(SUM(n), 0) FROM ( \
SELECT COALESCE(MAX(c.end_line), 0) AS n \
FROM files f LEFT JOIN chunks c ON c.file_id = f.id \
WHERE {predicate} GROUP BY f.id)"
);
let mut stmt = db.prepare(&sql)?;
let args = binds.into_iter().map(Value::Text);
let lines: i64 = stmt.query_row(rusqlite::params_from_iter(args), |r| r.get(0))?;
Ok(Facets { extensions, lines })
}
#[derive(Debug, Clone)]
pub struct ImageRow {
pub id: i64,
pub path: String,
pub width: u32,
pub height: u32,
}
pub fn insert_image(db: &Connection, file_id: i64, width: u32, height: u32) -> Result<i64> {
db.execute(
"INSERT INTO images (file_id, width, height) VALUES (?1, ?2, ?3)",
params![file_id, width, height],
)?;
Ok(db.last_insert_rowid())
}
pub fn image(db: &Connection, id: i64) -> Result<Option<ImageRow>> {
Ok(db
.query_row(
"SELECT i.id, f.path, i.width, i.height FROM images i \
JOIN files f ON f.id = i.file_id WHERE i.id = ?1",
params![id],
|r| {
Ok(ImageRow {
id: r.get(0)?,
path: r.get(1)?,
width: r.get(2)?,
height: r.get(3)?,
})
},
)
.optional()?)
}
pub fn image_ids_of(db: &Connection, path: &str) -> Result<Vec<i64>> {
let mut stmt =
db.prepare("SELECT i.id FROM images i JOIN files f ON f.id = i.file_id WHERE f.path = ?1")?;
let rows = stmt.query_map(params![path], |r| r.get::<_, i64>(0))?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn image_count(db: &Connection) -> Result<i64> {
Ok(db.query_row("SELECT COUNT(*) FROM images", [], |r| r.get(0))?)
}
pub fn image_ids(db: &Connection) -> Result<Vec<i64>> {
let mut stmt = db.prepare("SELECT id FROM images")?;
let rows = stmt.query_map([], |r| r.get::<_, i64>(0))?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn filtered_image_ids(db: &Connection, groups: &[Vec<String>]) -> Result<Vec<i64>> {
let (predicate, binds) = glob_predicate(groups);
let sql =
format!("SELECT i.id FROM images i JOIN files f ON f.id = i.file_id WHERE {predicate}");
let mut stmt = db.prepare(&sql)?;
let args = binds.into_iter().map(Value::Text);
let rows = stmt.query_map(rusqlite::params_from_iter(args), |r| r.get::<_, i64>(0))?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn all_paths(db: &Connection) -> Result<Vec<String>> {
filtered_paths(db, &[])
}
pub fn filtered_paths(db: &Connection, groups: &[Vec<String>]) -> Result<Vec<String>> {
let (predicate, binds) = glob_predicate(groups);
let sql = format!("SELECT f.path FROM files f WHERE {predicate} ORDER BY f.path");
let mut stmt = db.prepare(&sql)?;
let args = binds.into_iter().map(Value::Text);
let rows = stmt.query_map(rusqlite::params_from_iter(args), |r| r.get::<_, String>(0))?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn durable_chunks(db: &Connection) -> Result<i64> {
Ok(db.query_row(
"SELECT COUNT(*) FROM chunks c JOIN files f ON f.id = c.file_id WHERE f.hash != ''",
[],
|r| r.get(0),
)?)
}
#[cfg(test)]
pub fn prepare_for_tests(db: &Connection) {
db.execute_batch(SCHEMA).unwrap();
db.pragma_update(None, "foreign_keys", "ON").unwrap();
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct SymbolRow {
pub id: i64,
#[serde(serialize_with = "crate::serialize_plain")]
pub path: String,
pub kind: String,
pub name: String,
pub qualified: String,
pub start_line: u32,
pub end_line: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub chunk_id: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub store: Option<String>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct EdgeEnd {
#[serde(flatten)]
pub symbol: SymbolRow,
pub kind: String,
pub confidence: String,
pub definitions: usize,
#[serde(skip_serializing_if = "Option::is_none")]
pub from_path: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub from_line: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub line: Option<u32>,
}
const SYMBOL_COLUMNS: &str =
"s.id, f.path, s.kind, s.name, s.qualified, s.start_line, s.end_line, s.chunk_id";
fn symbol_row(r: &rusqlite::Row<'_>) -> rusqlite::Result<SymbolRow> {
Ok(SymbolRow {
id: r.get(0)?,
path: r.get(1)?,
kind: r.get(2)?,
name: r.get(3)?,
qualified: r.get(4)?,
start_line: r.get(5)?,
end_line: r.get(6)?,
chunk_id: r.get(7)?,
store: None,
})
}
pub fn insert_symbol(
db: &Connection,
file_id: i64,
chunk_id: Option<i64>,
symbol: &crate::graph::Symbol,
) -> Result<i64> {
db.execute(
"INSERT INTO symbols (file_id, chunk_id, kind, name, qualified, start_line, end_line)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7)",
params![
file_id,
chunk_id,
symbol.kind,
symbol.name,
symbol.qualified,
symbol.start_line,
symbol.end_line
],
)?;
Ok(db.last_insert_rowid())
}
pub fn insert_edge(
db: &Connection,
src: i64,
dst: &str,
kind: &str,
confidence: &str,
hint: Option<&str>,
line: Option<u32>,
) -> Result<()> {
db.execute(
"INSERT INTO edges (src, dst, kind, confidence, hint, line)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
params![src, dst, kind, confidence, hint, line],
)?;
Ok(())
}
pub fn symbols_named(db: &Connection, name: &str, limit: usize) -> Result<Vec<SymbolRow>> {
let sql = format!(
"SELECT {SYMBOL_COLUMNS} FROM symbols s JOIN files f ON f.id = s.file_id
WHERE s.name = ?1 ORDER BY f.path, s.start_line LIMIT ?2"
);
let mut stmt = db.prepare(&sql)?;
let rows = stmt.query_map(params![name, limit as i64], symbol_row)?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn symbols_scoped(
db: &Connection,
prefix: Option<&str>,
limit: usize,
) -> Result<Vec<SymbolRow>> {
let (predicate, binds) = match prefix {
Some(p) => {
let pattern = format!("*{}*", p.to_lowercase());
("lower(f.path) GLOB ?".to_string(), vec![pattern])
}
None => ("1".to_string(), Vec::new()),
};
let sql = format!(
"SELECT {SYMBOL_COLUMNS} FROM symbols s JOIN files f ON f.id = s.file_id
WHERE {predicate} AND s.kind != 'module'
ORDER BY (
(SELECT COUNT(*) FROM edges e WHERE e.src = s.id)
+ (SELECT COUNT(*) FROM edges e WHERE e.dst = s.name)
) DESC, f.path, s.start_line LIMIT ?"
);
let mut stmt = db.prepare(&sql)?;
let mut args: Vec<Value> = binds.into_iter().map(Value::Text).collect();
args.push(Value::Integer(limit as i64));
let rows = stmt.query_map(rusqlite::params_from_iter(args), symbol_row)?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn symbols_by_names(
db: &Connection,
names: &[String],
groups: &[Vec<String>],
) -> Result<Vec<SymbolRow>> {
if names.is_empty() {
return Ok(Vec::new());
}
let (predicate, mut binds) = glob_predicate(groups);
let holes = vec!["?"; names.len()].join(", ");
let sql = format!(
"SELECT {SYMBOL_COLUMNS} FROM symbols s JOIN files f ON f.id = s.file_id
WHERE {predicate} AND s.name IN ({holes}) ORDER BY f.path, s.start_line"
);
binds.extend(names.iter().cloned());
let mut stmt = db.prepare(&sql)?;
let args = binds.into_iter().map(Value::Text);
let rows = stmt.query_map(rusqlite::params_from_iter(args), symbol_row)?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn edges_out(db: &Connection, name: &str, kinds: &[String]) -> Result<Vec<EdgeEnd>> {
let filter = kind_predicate(kinds, "e.kind");
let target = crate::graph::not_navigational("s.kind");
let sql = format!(
"SELECT {SYMBOL_COLUMNS}, e.kind, e.confidence, e.hint, srcf.path, src.start_line, e.line
FROM symbols src
JOIN files srcf ON srcf.id = src.file_id
JOIN edges e ON e.src = src.id
JOIN symbols s ON s.name = e.dst AND {target}
JOIN files f ON f.id = s.file_id
WHERE src.name = ?1 AND {filter}
ORDER BY f.path, s.start_line"
);
let mut stmt = db.prepare(&sql)?;
let mut binds: Vec<Value> = vec![Value::Text(name.to_string())];
binds.extend(kinds.iter().map(|k| Value::Text(k.clone())));
let rows = stmt.query_map(rusqlite::params_from_iter(binds), |r| {
Ok(Reached {
symbol: symbol_row(r)?,
kind: r.get(8)?,
confidence: r.get(9)?,
hint: r.get(10)?,
src_path: r.get(11)?,
src_line: r.get(12)?,
line: r.get(13)?,
})
})?;
let reached = rows.collect::<Result<Vec<_>, _>>()?;
resolve(db, reached)
}
struct Reached {
symbol: SymbolRow,
kind: String,
confidence: String,
hint: Option<String>,
src_path: String,
src_line: u32,
line: Option<u32>,
}
fn resolve(db: &Connection, reached: Vec<Reached>) -> Result<Vec<EdgeEnd>> {
type Key = (String, String, String, Option<String>);
let mut order: Vec<Key> = Vec::new();
let mut groups: std::collections::HashMap<Key, Vec<Reached>> = std::collections::HashMap::new();
for row in reached {
let key = (
row.src_path.clone(),
row.symbol.name.clone(),
row.kind.clone(),
row.hint.clone(),
);
if !groups.contains_key(&key) {
order.push(key.clone());
}
groups.entry(key).or_default().push(row);
}
let mut imports: std::collections::HashMap<String, Vec<String>> =
std::collections::HashMap::new();
let mut out = Vec::new();
for key in order {
let candidates = groups.remove(&key).unwrap_or_default();
let (src_path, _, _, hint) = &key;
let definitions = candidates.len();
if definitions == 1 {
let row = candidates.into_iter().next().expect("one candidate");
let confidence = if row.confidence == crate::graph::EXTRACTED {
crate::graph::EXTRACTED
} else {
crate::graph::RESOLVED
};
out.push(EdgeEnd {
symbol: row.symbol,
kind: row.kind,
confidence: confidence.to_string(),
definitions,
from_path: Some(row.src_path),
from_line: Some(row.src_line),
line: row.line,
});
continue;
}
let mut ranks: Vec<u8> = candidates
.iter()
.map(|c| rank(&c.symbol.path, src_path, hint.as_deref(), &[]))
.collect();
if survivors(&ranks) != 1 && ranks.contains(&UNRANKED) {
let of_file = match imports.get(src_path) {
Some(found) => found.clone(),
None => {
let found = file_imports(db, src_path)?;
imports.insert(src_path.clone(), found.clone());
found
}
};
ranks = candidates
.iter()
.map(|c| rank(&c.symbol.path, src_path, hint.as_deref(), &of_file))
.collect();
}
let best = ranks.iter().copied().min().unwrap_or(UNRANKED);
let settled = survivors(&ranks) == 1;
for (row, rank) in candidates.into_iter().zip(&ranks) {
if settled && *rank != best {
continue;
}
let confidence = if row.confidence == crate::graph::EXTRACTED && settled {
crate::graph::EXTRACTED
} else if settled {
crate::graph::RESOLVED
} else {
crate::graph::AMBIGUOUS
};
out.push(EdgeEnd {
symbol: row.symbol,
kind: row.kind,
confidence: confidence.to_string(),
definitions,
from_path: Some(row.src_path),
from_line: Some(row.src_line),
line: row.line,
});
}
}
Ok(out)
}
const UNRANKED: u8 = 3;
fn survivors(ranks: &[u8]) -> usize {
match ranks.iter().min() {
Some(best) => ranks.iter().filter(|r| *r == best).count(),
None => 0,
}
}
fn rank(candidate: &str, src_path: &str, hint: Option<&str>, imports: &[String]) -> u8 {
if candidate == src_path {
return 0;
}
if hint.is_some_and(|h| names_file(h, candidate)) {
return 1;
}
if imports.iter().any(|i| {
i.split(['/', '.', ':', '\\'])
.any(|segment| !segment.is_empty() && names_file(segment, candidate))
}) {
return 2;
}
UNRANKED
}
fn names_file(word: &str, path: &str) -> bool {
let path = std::path::Path::new(path);
if path.file_stem().is_some_and(|stem| stem == word) {
return true;
}
path.parent()
.into_iter()
.flat_map(|parent| parent.components())
.any(|component| component.as_os_str() == word)
}
fn file_imports(db: &Connection, path: &str) -> Result<Vec<String>> {
let mut stmt = db.prepare(
"SELECT DISTINCT e.dst FROM edges e
JOIN symbols s ON s.id = e.src
JOIN files f ON f.id = s.file_id
WHERE f.path = ?1 AND e.kind = 'imports'",
)?;
let rows = stmt.query_map(params![path], |r| r.get(0))?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn unresolved_out(db: &Connection, name: &str, kinds: &[String]) -> Result<Vec<Unresolved>> {
let filter = kind_predicate(kinds, "e.kind");
let sql = format!(
"SELECT DISTINCT e.dst, e.kind, e.confidence
FROM symbols src
JOIN edges e ON e.src = src.id
WHERE src.name = ?1 AND {filter}
AND NOT EXISTS (SELECT 1 FROM symbols s WHERE s.name = e.dst)
ORDER BY e.dst"
);
let mut stmt = db.prepare(&sql)?;
let mut binds: Vec<Value> = vec![Value::Text(name.to_string())];
binds.extend(kinds.iter().map(|k| Value::Text(k.clone())));
let rows = stmt.query_map(rusqlite::params_from_iter(binds), |r| {
Ok(Unresolved {
name: r.get(0)?,
kind: r.get(1)?,
confidence: r.get(2)?,
})
})?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct Unresolved {
pub name: String,
pub kind: String,
pub confidence: String,
}
pub fn edges_in(db: &Connection, name: &str, kinds: &[String]) -> Result<Vec<EdgeEnd>> {
let filter = kind_predicate(kinds, "e.kind");
let sql = format!(
"SELECT {SYMBOL_COLUMNS}, e.kind, e.confidence, e.line
FROM edges e
JOIN symbols s ON s.id = e.src
JOIN files f ON f.id = s.file_id
WHERE e.dst = ?1 AND {filter}
ORDER BY f.path, s.start_line"
);
let mut stmt = db.prepare(&sql)?;
let mut binds: Vec<Value> = vec![Value::Text(name.to_string())];
binds.extend(kinds.iter().map(|k| Value::Text(k.clone())));
let rows = stmt.query_map(rusqlite::params_from_iter(binds), |r| {
Ok(EdgeEnd {
symbol: symbol_row(r)?,
kind: r.get(8)?,
confidence: r.get(9)?,
definitions: 1,
from_path: None,
from_line: None,
line: r.get(10)?,
})
})?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
fn kind_predicate(kinds: &[String], column: &str) -> String {
if kinds.is_empty() {
return "1".to_string();
}
let holes = vec!["?"; kinds.len()].join(", ");
format!("{column} IN ({holes})")
}
pub fn symbols_in_chunks(db: &Connection, chunk_ids: &[u64]) -> Result<Vec<String>> {
if chunk_ids.is_empty() {
return Ok(Vec::new());
}
let holes = vec!["?"; chunk_ids.len()].join(", ");
let navigational = crate::graph::not_navigational("kind");
let sql =
format!("SELECT DISTINCT name FROM symbols WHERE chunk_id IN ({holes}) AND {navigational}");
let mut stmt = db.prepare(&sql)?;
let args = chunk_ids.iter().map(|i| Value::Integer(*i as i64));
let rows = stmt.query_map(rusqlite::params_from_iter(args), |r| r.get::<_, String>(0))?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct Retrieval {
pub id: i64,
pub at: i64,
pub client: String,
pub query: String,
pub hits: i64,
pub micros: i64,
pub excerpt_tokens: i64,
pub whole_file_tokens: i64,
pub hash: String,
}
pub fn record_retrieval(db: &Connection, row: &NewRetrieval<'_>) -> Result<()> {
let (client, query) = (row.client, row.query);
let (hits, micros) = (row.hits, row.micros);
let (excerpt_tokens, whole_file_tokens) = (row.excerpt_tokens, row.whole_file_tokens);
let _writing = Writing::begin(db)?;
let at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0);
let prev: String = db
.query_row(
"SELECT hash FROM retrievals ORDER BY id DESC LIMIT 1",
[],
|r| r.get(0),
)
.optional()?
.unwrap_or_default();
let hash = chain_hash(&prev, at, row);
db.execute(
"INSERT INTO retrievals
(at, client, query, hits, micros, excerpt_tokens, whole_file_tokens, prev, hash,
session, tool, stale_hits, tokenizer)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)",
params![
at,
client,
query,
hits,
micros,
excerpt_tokens,
whole_file_tokens,
prev,
hash,
row.session,
row.tool,
row.stale_hits,
row.tokenizer
],
)?;
Ok(())
}
#[derive(Debug, Clone)]
pub struct NewRetrieval<'a> {
pub client: &'a str,
pub session: &'a str,
pub tool: &'a str,
pub query: &'a str,
pub hits: i64,
pub micros: i64,
pub excerpt_tokens: i64,
pub whole_file_tokens: i64,
pub stale_hits: i64,
pub tokenizer: &'a str,
}
fn chain_hash(prev: &str, at: i64, row: &NewRetrieval<'_>) -> String {
let NewRetrieval {
client,
query,
hits,
micros,
excerpt_tokens: excerpt,
whole_file_tokens: whole,
..
} = row;
let mut payload = format!(
"{prev}\u{1f}{at}\u{1f}{client}\u{1f}{query}\u{1f}{hits}\u{1f}{micros}\u{1f}{excerpt}\u{1f}{whole}"
);
payload.push_str(&format!(
"\u{1f}{}\u{1f}{}\u{1f}{}\u{1f}{}",
row.session, row.tool, row.stale_hits, row.tokenizer
));
blake3::hash(payload.as_bytes()).to_hex().to_string()
}
pub fn record_legacy_retrieval(
db: &Connection,
client: &str,
query: &str,
excerpt_tokens: i64,
whole_file_tokens: i64,
) -> Result<()> {
let _writing = Writing::begin(db)?;
let at = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs() as i64)
.unwrap_or(0);
let prev: String = db
.query_row(
"SELECT hash FROM retrievals ORDER BY id DESC LIMIT 1",
[],
|r| r.get(0),
)
.optional()?
.unwrap_or_default();
let (hits, micros) = (1, 1000);
let hash = legacy_chain_hash(
&prev,
at,
client,
query,
hits,
micros,
excerpt_tokens,
whole_file_tokens,
);
db.execute(
"INSERT INTO retrievals
(at, client, query, hits, micros, excerpt_tokens, whole_file_tokens, prev, hash)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
params![
at,
client,
query,
hits,
micros,
excerpt_tokens,
whole_file_tokens,
prev,
hash
],
)?;
Ok(())
}
#[allow(clippy::too_many_arguments)]
fn legacy_chain_hash(
prev: &str,
at: i64,
client: &str,
query: &str,
hits: i64,
micros: i64,
excerpt: i64,
whole: i64,
) -> String {
let payload = format!(
"{prev}\u{1f}{at}\u{1f}{client}\u{1f}{query}\u{1f}{hits}\u{1f}{micros}\u{1f}{excerpt}\u{1f}{whole}"
);
blake3::hash(payload.as_bytes()).to_hex().to_string()
}
pub fn retrievals(db: &Connection, limit: usize) -> Result<Vec<Retrieval>> {
let mut stmt = db.prepare(
"SELECT id, at, client, query, hits, micros, excerpt_tokens, whole_file_tokens, hash
FROM retrievals ORDER BY id DESC LIMIT ?1",
)?;
let rows = stmt.query_map(params![limit as i64], |r| {
Ok(Retrieval {
id: r.get(0)?,
at: r.get(1)?,
client: r.get(2)?,
query: r.get(3)?,
hits: r.get(4)?,
micros: r.get(5)?,
excerpt_tokens: r.get(6)?,
whole_file_tokens: r.get(7)?,
hash: r.get(8)?,
})
})?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn ledger_totals(db: &Connection) -> Result<(i64, i64, i64, i64)> {
Ok(db.query_row(
"SELECT COUNT(*), COUNT(DISTINCT client), COALESCE(SUM(excerpt_tokens), 0),
COALESCE(SUM(whole_file_tokens), 0) FROM retrievals",
[],
|r| Ok((r.get(0)?, r.get(1)?, r.get(2)?, r.get(3)?)),
)?)
}
pub fn ledger_savings(db: &Connection) -> Result<Savings> {
let (credited, net): (i64, i64) = db.query_row(
"SELECT COUNT(*), COALESCE(SUM(whole_file_tokens - excerpt_tokens), 0)
FROM retrievals WHERE hits > 0",
[],
|r| Ok((r.get(0)?, r.get(1)?)),
)?;
let total: i64 = db.query_row("SELECT COUNT(*) FROM retrievals", [], |r| r.get(0))?;
let estimated: i64 = db.query_row(
"SELECT COUNT(*) FROM retrievals
WHERE hits > 0 AND (tokenizer IS NULL OR tokenizer != 'model')",
[],
|r| r.get(0),
)?;
Ok(Savings {
net: net.max(0),
credited,
total,
measured: credited > 0 && estimated == 0,
})
}
#[derive(Debug, Clone, Copy, serde::Serialize)]
pub struct Savings {
pub net: i64,
pub credited: i64,
pub total: i64,
pub measured: bool,
}
impl Savings {
pub fn tier(&self) -> &'static str {
if self.measured {
"measured"
} else {
"modelled"
}
}
pub fn coverage(&self) -> i64 {
if self.total == 0 {
return 0;
}
self.credited * 100 / self.total
}
}
pub fn ledger_clients(db: &Connection) -> Result<Vec<(String, i64)>> {
let mut stmt = db.prepare(
"SELECT client, COUNT(*) FROM retrievals GROUP BY client ORDER BY COUNT(*) DESC, client",
)?;
let rows = stmt.query_map([], |r| Ok((r.get(0)?, r.get(1)?)))?;
Ok(rows.collect::<Result<Vec<_>, _>>()?)
}
pub fn ledger_break(db: &Connection) -> Result<Option<i64>> {
let mut stmt = db.prepare(
"SELECT id, at, client, query, hits, micros, excerpt_tokens, whole_file_tokens, prev, hash,
session, tool, stale_hits, tokenizer
FROM retrievals ORDER BY id",
)?;
let mut rows = stmt.query([])?;
let mut expected = String::new();
while let Some(r) = rows.next()? {
let (id, at): (i64, i64) = (r.get(0)?, r.get(1)?);
let client: String = r.get(2)?;
let query: String = r.get(3)?;
let (hits, micros, excerpt, whole): (i64, i64, i64, i64) =
(r.get(4)?, r.get(5)?, r.get(6)?, r.get(7)?);
let prev: String = r.get(8)?;
let hash: String = r.get(9)?;
if prev != expected {
return Ok(Some(id));
}
let tool: Option<String> = r.get(11)?;
let recomputed = match tool {
Some(tool) => {
let session: String = r.get(10)?;
let stale: i64 = r.get(12)?;
let tokenizer: String = r.get(13)?;
chain_hash(
&prev,
at,
&NewRetrieval {
client: &client,
session: &session,
tool: &tool,
query: &query,
hits,
micros,
excerpt_tokens: excerpt,
whole_file_tokens: whole,
stale_hits: stale,
tokenizer: &tokenizer,
},
)
}
None => legacy_chain_hash(&prev, at, &client, &query, hits, micros, excerpt, whole),
};
if recomputed != hash {
return Ok(Some(id));
}
expected = hash;
}
Ok(None)
}
pub fn grep_cost(db: &Connection, name: &str) -> Result<i64> {
Ok(db.query_row(
"SELECT COALESCE(SUM(bytes), 0) FROM files WHERE id IN (
SELECT s.file_id FROM symbols s WHERE s.name = ?1
UNION
SELECT s.file_id FROM symbols s JOIN edges e ON e.src = s.id WHERE e.dst = ?1
)",
params![name],
|r| r.get(0),
)?)
}
pub type SymbolSpan = (u32, u32, String, String);
pub fn symbols_in_files(
db: &Connection,
paths: &[String],
) -> Result<std::collections::HashMap<String, Vec<SymbolSpan>>> {
let mut out: std::collections::HashMap<String, Vec<SymbolSpan>> = Default::default();
if paths.is_empty() {
return Ok(out);
}
let holes = vec!["?"; paths.len()].join(", ");
let sql = format!(
"SELECT f.path, s.start_line, s.end_line, s.name, s.kind
FROM symbols s JOIN files f ON f.id = s.file_id
WHERE f.path IN ({holes}) AND s.kind != 'module'
ORDER BY f.path, s.start_line"
);
let mut stmt = db.prepare(&sql)?;
let args = paths.iter().cloned().map(Value::Text);
let mut rows = stmt.query(rusqlite::params_from_iter(args))?;
while let Some(r) = rows.next()? {
let path: String = r.get(0)?;
out.entry(path)
.or_default()
.push((r.get(1)?, r.get(2)?, r.get(3)?, r.get(4)?));
}
Ok(out)
}
pub fn file_stamps(
db: &Connection,
paths: &[String],
) -> Result<std::collections::HashMap<String, (i64, i64)>> {
if paths.is_empty() {
return Ok(Default::default());
}
let holes = vec!["?"; paths.len()].join(", ");
let sql = format!("SELECT path, bytes, indexed_at FROM files WHERE path IN ({holes})");
let mut stmt = db.prepare(&sql)?;
let args = paths.iter().cloned().map(Value::Text);
let rows = stmt.query_map(rusqlite::params_from_iter(args), |r| {
Ok((r.get::<_, String>(0)?, (r.get(1)?, r.get(2)?)))
})?;
Ok(rows.collect::<Result<std::collections::HashMap<_, _>, _>>()?)
}
pub fn graph_stats(db: &Connection) -> Result<(i64, i64)> {
let symbols: i64 = db.query_row("SELECT COUNT(*) FROM symbols", [], |r| r.get(0))?;
let edges: i64 = db.query_row("SELECT COUNT(*) FROM edges", [], |r| r.get(0))?;
Ok((symbols, edges))
}
pub fn stats(db: &Connection) -> Result<(i64, i64, i64)> {
let files: i64 = db.query_row("SELECT COUNT(*) FROM files", [], |r| r.get(0))?;
let chunks: i64 = db.query_row("SELECT COUNT(*) FROM chunks", [], |r| r.get(0))?;
let bytes: i64 = db.query_row("SELECT COALESCE(SUM(bytes), 0) FROM files", [], |r| {
r.get(0)
})?;
Ok((files, chunks, bytes))
}
#[cfg(test)]
mod tests {
use super::*;
fn sym(name: &str) -> crate::graph::Symbol {
crate::graph::Symbol {
kind: "function".to_string(),
name: name.to_string(),
qualified: name.to_string(),
start_line: 1,
end_line: 2,
}
}
fn at(mut s: crate::graph::Symbol, start: u32, end: u32) -> crate::graph::Symbol {
s.start_line = start;
s.end_line = end;
s
}
fn one_symbol(db: &Connection, path: &str, name: &str) -> i64 {
db.execute_batch(SCHEMA).unwrap();
db.pragma_update(None, "foreign_keys", "ON").unwrap();
let file_id = insert_file(db, path, "h", 1, 0).unwrap();
let chunk_id = insert_chunk(db, file_id, 0, 1, 2, name).unwrap();
insert_symbol(db, file_id, Some(chunk_id), &sym(name)).unwrap()
}
#[test]
fn deleting_a_file_deletes_its_symbols_and_its_outgoing_edges() {
let db = Connection::open_in_memory().unwrap();
let caller = one_symbol(&db, "a.rs", "caller");
insert_edge(&db, caller, "callee", "calls", "inferred", None, None).unwrap();
assert_eq!(graph_stats(&db).unwrap(), (1, 1));
delete_file(&db, "a.rs").unwrap();
assert_eq!(
graph_stats(&db).unwrap(),
(0, 0),
"the file's symbols and the edges leaving them both cascade"
);
}
#[test]
fn re_indexing_the_target_file_leaves_edges_into_it_intact() {
let db = Connection::open_in_memory().unwrap();
let caller = one_symbol(&db, "a.rs", "caller");
insert_edge(&db, caller, "callee", "calls", "inferred", None, None).unwrap();
let b = insert_file(&db, "b.rs", "h", 1, 0).unwrap();
insert_symbol(&db, b, None, &sym("callee")).unwrap();
assert_eq!(edges_out(&db, "caller", &[]).unwrap().len(), 1);
delete_file(&db, "b.rs").unwrap();
let b = insert_file(&db, "b.rs", "h2", 1, 0).unwrap();
insert_symbol(&db, b, None, &at(sym("callee"), 9, 10)).unwrap();
let out = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(
out.len(),
1,
"the edge survived its target being re-indexed"
);
assert_eq!(
out[0].symbol.start_line, 9,
"and now points at the new rows"
);
}
fn many_definitions(paths: &[&str]) -> Connection {
let db = Connection::open_in_memory().unwrap();
one_symbol(&db, "src/caller.rs", "caller");
for (i, path) in paths.iter().enumerate() {
let file = db
.query_row("SELECT id FROM files WHERE path = ?1", params![path], |r| {
r.get(0)
})
.optional()
.unwrap()
.unwrap_or_else(|| insert_file(&db, path, &format!("h{i}"), 1, 0).unwrap());
insert_symbol(&db, file, None, &sym("callee")).unwrap();
}
db
}
fn call(db: &Connection, hint: Option<&str>) {
let src: i64 = db
.query_row("SELECT id FROM symbols WHERE name = 'caller'", [], |r| {
r.get(0)
})
.unwrap();
insert_edge(db, src, "callee", "calls", "inferred", hint, None).unwrap();
}
#[test]
fn a_name_with_one_definition_resolves() {
let db = many_definitions(&["src/other.rs"]);
call(&db, None);
let out = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(out.len(), 1);
assert_eq!(out[0].confidence, crate::graph::RESOLVED);
assert_eq!(out[0].definitions, 1);
}
#[test]
fn a_name_nothing_narrows_is_ambiguous_and_keeps_every_candidate() {
let db = many_definitions(&["src/a.rs", "src/b.rs", "src/c.rs", "src/d.rs"]);
call(&db, None);
let out = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(out.len(), 4);
assert!(out.iter().all(|e| e.confidence == crate::graph::AMBIGUOUS));
assert!(out.iter().all(|e| e.definitions == 4));
}
#[test]
fn a_hint_that_names_a_file_resolves_to_that_definition() {
let db = many_definitions(&["src/a.rs", "src/store.rs", "src/c.rs"]);
call(&db, Some("store"));
let out = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(out.len(), 1, "{out:?}");
assert_eq!(out[0].symbol.path, "src/store.rs");
assert_eq!(out[0].confidence, crate::graph::RESOLVED);
assert_eq!(
out[0].definitions, 3,
"the count is what the ranking chose between, not what survived"
);
}
#[test]
fn a_hint_names_a_directory_as_readily_as_a_file() {
let db = many_definitions(&["src/a.rs", "src/store/mod.rs"]);
call(&db, Some("store"));
let out = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(out.len(), 1);
assert_eq!(out[0].symbol.path, "src/store/mod.rs");
}
#[test]
fn the_calling_file_outranks_the_hint() {
let db = many_definitions(&["src/caller.rs", "src/store.rs"]);
call(&db, Some("store"));
let out = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(out.len(), 1);
assert_eq!(out[0].symbol.path, "src/caller.rs");
}
#[test]
fn an_import_resolves_what_the_hint_cannot() {
let db = many_definitions(&["src/a.rs", "src/b.rs"]);
let src: i64 = db
.query_row("SELECT id FROM symbols WHERE name = 'caller'", [], |r| {
r.get(0)
})
.unwrap();
insert_edge(
&db,
src,
"crate::b::callee",
"imports",
"extracted",
None,
None,
)
.unwrap();
call(&db, None);
let out = edges_out(&db, "caller", &["calls".to_string()]).unwrap();
assert_eq!(out.len(), 1, "{out:?}");
assert_eq!(out[0].symbol.path, "src/b.rs");
assert_eq!(out[0].confidence, crate::graph::RESOLVED);
}
#[test]
fn the_hint_outranks_an_import() {
let db = many_definitions(&["src/a.rs", "src/b.rs"]);
let src: i64 = db
.query_row("SELECT id FROM symbols WHERE name = 'caller'", [], |r| {
r.get(0)
})
.unwrap();
insert_edge(
&db,
src,
"crate::b::callee",
"imports",
"extracted",
None,
None,
)
.unwrap();
call(&db, Some("a"));
let out = edges_out(&db, "caller", &["calls".to_string()]).unwrap();
assert_eq!(out.len(), 1, "{out:?}");
assert_eq!(out[0].symbol.path, "src/a.rs");
}
#[test]
fn an_extracted_edge_stays_extracted() {
let db = many_definitions(&["src/store.rs"]);
let src: i64 = db
.query_row("SELECT id FROM symbols WHERE name = 'caller'", [], |r| {
r.get(0)
})
.unwrap();
insert_edge(
&db,
src,
"callee",
"calls",
"extracted",
Some("store"),
None,
)
.unwrap();
let out = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(out[0].confidence, crate::graph::EXTRACTED);
}
#[test]
fn a_null_hint_behaves_as_it_did_before_hints_existed() {
let db = many_definitions(&["src/a.rs"]);
call(&db, None);
let one = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(one.len(), 1);
assert_eq!(one[0].symbol.path, "src/a.rs");
let db = many_definitions(&["src/a.rs", "src/b.rs"]);
call(&db, None);
let two = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(
two.len(),
2,
"both candidates are still reachable, now labelled rather than asserted"
);
}
#[test]
fn callers_keep_the_stored_confidence() {
let db = many_definitions(&["src/a.rs", "src/b.rs"]);
call(&db, None);
let inbound = edges_in(&db, "callee", &[]).unwrap();
assert_eq!(inbound.len(), 1);
assert_eq!(inbound[0].confidence, crate::graph::INFERRED);
assert_eq!(inbound[0].definitions, 1);
}
#[test]
fn edges_resolve_in_both_directions_and_filter_by_kind() {
let db = Connection::open_in_memory().unwrap();
let caller = one_symbol(&db, "a.rs", "caller");
let b = insert_file(&db, "b.rs", "h", 1, 0).unwrap();
insert_symbol(&db, b, None, &sym("callee")).unwrap();
insert_edge(&db, caller, "callee", "calls", "extracted", None, None).unwrap();
insert_edge(&db, caller, "callee", "references", "inferred", None, None).unwrap();
assert_eq!(edges_out(&db, "caller", &[]).unwrap().len(), 2);
let calls = edges_out(&db, "caller", &["calls".to_string()]).unwrap();
assert_eq!(calls.len(), 1);
assert_eq!(calls[0].confidence, "extracted");
let inbound = edges_in(&db, "callee", &["calls".to_string()]).unwrap();
assert_eq!(inbound.len(), 1);
assert_eq!(
inbound[0].symbol.name, "caller",
"edges_in reports the source"
);
}
#[test]
fn an_edge_to_an_unindexed_target_resolves_to_nothing_without_erroring() {
let db = Connection::open_in_memory().unwrap();
let caller = one_symbol(&db, "a.rs", "caller");
insert_edge(&db, caller, "println", "calls", "inferred", None, None).unwrap();
assert!(edges_out(&db, "caller", &[]).unwrap().is_empty());
assert_eq!(graph_stats(&db).unwrap().1, 1, "but the edge row is there");
}
#[test]
fn a_store_without_the_key_is_the_format_we_understand() {
let db = Connection::open_in_memory().unwrap();
db.execute_batch(SCHEMA).unwrap();
assert!(get_meta(&db, FORMAT_KEY).unwrap().is_none());
check_format(&db).expect("a pre-0.6.0 store is readable");
assert!(get_meta(&db, FORMAT_KEY).unwrap().is_none());
}
#[test]
fn a_newer_format_is_refused_naming_both_numbers() {
let db = Connection::open_in_memory().unwrap();
db.execute_batch(SCHEMA).unwrap();
let newer = FORMAT_VERSION + 1;
set_meta(&db, FORMAT_KEY, &newer.to_string()).unwrap();
let refused = check_format(&db).expect_err("a newer format must not be read");
let said = refused.to_string();
assert!(said.contains(&newer.to_string()), "{said}");
assert!(
said.contains(&FORMAT_VERSION.to_string()),
"the error must name what this binary understands: {said}"
);
set_meta(&db, FORMAT_KEY, &FORMAT_VERSION.to_string()).unwrap();
check_format(&db).unwrap();
}
fn seeded() -> Connection {
let db = Connection::open_in_memory().unwrap();
db.pragma_update(None, "foreign_keys", "ON").unwrap();
db.execute_batch(SCHEMA).unwrap();
let f = insert_file(&db, "/a/lib.rs", "h", 10, 0).unwrap();
insert_chunk(&db, f, 0, 1, 5, "const EMBED_BATCH: usize = 32;").unwrap();
insert_chunk(
&db,
f,
1,
6,
9,
"a paragraph about retry backoff and jitter",
)
.unwrap();
db
}
fn native(path: &str) -> String {
path.replace('/', std::path::MAIN_SEPARATOR_STR)
}
fn mixed() -> Connection {
let db = Connection::open_in_memory().unwrap();
db.pragma_update(None, "foreign_keys", "ON").unwrap();
db.execute_batch(SCHEMA).unwrap();
for path in [
"/proj/src/lib.rs",
"/proj/src/notes.md",
"/proj/vendor/other.rs",
"/proj/README.MD",
] {
let f = insert_file(&db, &native(path), "h", 10, 0).unwrap();
insert_chunk(&db, f, 0, 1, 5, "retry backoff and jitter").unwrap();
}
db
}
fn filter(paths: &[&str], exts: &[&str], langs: &[&str]) -> Vec<Vec<String>> {
let own = |v: &[&str]| v.iter().map(|s| s.to_string()).collect::<Vec<_>>();
crate::filter::Filter::new(&own(paths), &own(exts), &own(langs))
.unwrap()
.groups()
.to_vec()
}
fn paths_of(db: &Connection, ids: &[u64]) -> Vec<String> {
let mut out: Vec<String> = ids
.iter()
.map(|&i| chunk(db, i).unwrap().unwrap().path)
.collect();
out.sort();
out
}
#[test]
fn keyword_search_finds_an_exact_identifier() {
let db = seeded();
let hits = keyword_search(&db, "EMBED_BATCH", 10, &[]).unwrap();
assert_eq!(hits, vec![1], "expected the defining chunk, got {hits:?}");
}
#[test]
fn a_path_glob_selects_only_what_is_under_it() {
let db = mixed();
let ids = filtered_chunk_ids(&db, &filter(&["src/**"], &[], &[])).unwrap();
assert_eq!(
paths_of(&db, &ids),
[native("/proj/src/lib.rs"), native("/proj/src/notes.md")]
);
}
#[test]
fn an_absolute_glob_matches_only_what_it_literally_covers() {
let db = mixed();
let ids = filtered_chunk_ids(&db, &filter(&["/proj/vendor/*"], &[], &[])).unwrap();
assert_eq!(paths_of(&db, &ids), [native("/proj/vendor/other.rs")]);
}
#[test]
fn extensions_union_and_a_language_resolves_to_the_same_set() {
let db = mixed();
let by_ext = filtered_chunk_ids(&db, &filter(&[], &["rs"], &[])).unwrap();
let by_lang = filtered_chunk_ids(&db, &filter(&[], &[], &["rust"])).unwrap();
assert_eq!(paths_of(&db, &by_ext), paths_of(&db, &by_lang));
assert_eq!(
paths_of(&db, &by_ext),
[native("/proj/src/lib.rs"), native("/proj/vendor/other.rs")]
);
let both = filtered_chunk_ids(&db, &filter(&[], &["rs", "md"], &[])).unwrap();
assert_eq!(both.len(), 4, "repeating --ext must union, not intersect");
}
#[test]
fn extension_matching_is_case_insensitive() {
let db = mixed();
let ids = filtered_chunk_ids(&db, &filter(&[], &["md"], &[])).unwrap();
assert_eq!(
paths_of(&db, &ids),
[native("/proj/README.MD"), native("/proj/src/notes.md")]
);
}
#[test]
fn a_path_and_an_extension_intersect() {
let db = mixed();
let ids = filtered_chunk_ids(&db, &filter(&["src/**"], &["md"], &[])).unwrap();
assert_eq!(
paths_of(&db, &ids),
[native("/proj/src/notes.md")],
"the README is Markdown but is not under src"
);
}
#[test]
fn the_keyword_half_honours_the_same_filter() {
let db = mixed();
let unfiltered = keyword_search(&db, "retry backoff", 10, &[]).unwrap();
assert_eq!(unfiltered.len(), 4);
let scoped =
keyword_search(&db, "retry backoff", 10, &filter(&["src/**"], &[], &[])).unwrap();
assert_eq!(
paths_of(&db, &scoped),
[native("/proj/src/lib.rs"), native("/proj/src/notes.md")]
);
}
#[test]
fn a_filter_matching_nothing_selects_nothing() {
let db = mixed();
let groups = filter(&["nowhere/**"], &[], &[]);
assert_eq!(matching_files(&db, &groups).unwrap(), 0);
assert!(filtered_chunk_ids(&db, &groups).unwrap().is_empty());
assert!(
keyword_search(&db, "retry", 10, &groups)
.unwrap()
.is_empty()
);
}
#[test]
fn an_unfiltered_search_issues_the_0_2_0_statement_verbatim() {
assert_eq!(
UNFILTERED_SQL,
"SELECT rowid FROM chunks_fts WHERE chunks_fts MATCH ? ORDER BY rank LIMIT ?"
);
assert_eq!(glob_predicate(&[]), ("1".to_string(), vec![]));
}
#[test]
fn no_filter_counts_every_file() {
let db = mixed();
assert_eq!(matching_files(&db, &[]).unwrap(), 4);
}
#[test]
fn fts_syntax_in_a_query_is_not_executed_as_syntax() {
let db = seeded();
for query in [
"EMBED_BATCH AND retry",
"unbalanced \" quote",
"call_me(",
"prefix*",
"NEAR(a b)",
"-negated",
] {
let result = keyword_search(&db, query, 10, &[]);
assert!(
result.is_ok(),
"query {query:?} errored: {:?}",
result.err()
);
}
}
#[test]
fn a_query_with_no_terms_matches_nothing() {
let db = seeded();
assert!(keyword_search(&db, "!!! ???", 10, &[]).unwrap().is_empty());
}
#[test]
fn deleting_a_file_removes_its_chunks_from_the_keyword_index() {
let db = seeded();
assert!(
!keyword_search(&db, "EMBED_BATCH", 10, &[])
.unwrap()
.is_empty()
);
delete_file(&db, "/a/lib.rs").unwrap();
assert!(
keyword_search(&db, "EMBED_BATCH", 10, &[])
.unwrap()
.is_empty(),
"cascade delete left the keyword index stale"
);
}
#[test]
fn a_pre_0_16_store_opens_and_its_edges_report_no_call_site() {
let db = Connection::open_in_memory().unwrap();
db.execute_batch(
"CREATE TABLE files (id INTEGER PRIMARY KEY, path TEXT NOT NULL UNIQUE,
hash TEXT NOT NULL, bytes INTEGER NOT NULL, indexed_at INTEGER NOT NULL);
CREATE TABLE symbols (id INTEGER PRIMARY KEY AUTOINCREMENT,
file_id INTEGER NOT NULL REFERENCES files(id) ON DELETE CASCADE,
chunk_id INTEGER, kind TEXT NOT NULL, name TEXT NOT NULL,
qualified TEXT NOT NULL, start_line INTEGER NOT NULL, end_line INTEGER NOT NULL);
CREATE TABLE edges (src INTEGER NOT NULL REFERENCES symbols(id) ON DELETE CASCADE,
dst TEXT NOT NULL, kind TEXT NOT NULL, confidence TEXT NOT NULL, hint TEXT);
INSERT INTO files VALUES (1, '/a/lib.rs', 'h', 10, 0);
INSERT INTO symbols VALUES (1, 1, NULL, 'function', 'caller', 'm::caller', 1, 9);
INSERT INTO symbols VALUES (2, 1, NULL, 'function', 'callee', 'm::callee', 20, 24);
CREATE TABLE retrievals (id INTEGER PRIMARY KEY AUTOINCREMENT, at INTEGER NOT NULL,
client TEXT NOT NULL, query TEXT NOT NULL, hits INTEGER NOT NULL,
micros INTEGER NOT NULL, excerpt_tokens INTEGER NOT NULL,
whole_file_tokens INTEGER NOT NULL);
INSERT INTO edges VALUES (1, 'callee', 'calls', 'inferred', NULL);",
)
.unwrap();
add_columns(&db).unwrap();
let out = edges_out(&db, "caller", &[]).unwrap();
assert_eq!(out.len(), 1, "{out:?}");
assert_eq!(
out[0].line, None,
"a 0.15.0 row must not invent a call site"
);
assert_eq!(
crate::graph::call_site(&out[0], &crate::graph::verbatim),
"",
"nothing is rendered for an edge with no line"
);
let into = edges_in(&db, "callee", &[]).unwrap();
assert_eq!(into.len(), 1, "{into:?}");
assert_eq!(into[0].line, None);
}
#[test]
fn an_old_store_gets_its_keyword_index_backfilled() {
let db = Connection::open_in_memory().unwrap();
db.execute_batch(
"CREATE TABLE meta (k TEXT PRIMARY KEY, v TEXT NOT NULL);
CREATE TABLE files (id INTEGER PRIMARY KEY, path TEXT NOT NULL UNIQUE,
hash TEXT NOT NULL, bytes INTEGER NOT NULL, indexed_at INTEGER NOT NULL);
CREATE TABLE chunks (id INTEGER PRIMARY KEY, file_id INTEGER NOT NULL
REFERENCES files(id) ON DELETE CASCADE, ord INTEGER NOT NULL,
start_line INTEGER NOT NULL, end_line INTEGER NOT NULL, text TEXT NOT NULL);
INSERT INTO files VALUES (1, '/a/lib.rs', 'h', 10, 0);
INSERT INTO chunks VALUES (1, 1, 0, 1, 5, 'const EMBED_BATCH: usize = 32;');",
)
.unwrap();
db.execute_batch(SCHEMA).unwrap();
backfill_fts(&db).unwrap();
assert_eq!(
keyword_search(&db, "EMBED_BATCH", 10, &[]).unwrap(),
vec![1]
);
}
}