use crate::filter::Filter;
use crate::store;
use anyhow::{Result, bail};
use tree_sitter::{Parser, Query, QueryCursor, StreamingIterator};
pub const MAX_MATCHES: usize = 200;
pub const MAX_FILES: usize = 2000;
#[derive(Debug, Clone, serde::Serialize)]
pub struct Match {
pub path: String,
pub capture: String,
pub start_line: u32,
pub end_line: u32,
pub text: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub store: Option<String>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct Matches {
pub language: String,
pub matches: Vec<Match>,
pub files: usize,
pub truncated: bool,
pub skipped: usize,
}
pub fn run(
db: &rusqlite::Connection,
language: &str,
source: &str,
filter: &Filter,
offset: usize,
) -> Result<Matches> {
let language = language.trim().to_ascii_lowercase();
let Some(grammar) = crate::graph::language_grammar(&language) else {
bail!(
"no grammar for {language:?}; the languages with one are {}",
crate::graph::languages().join(", ")
);
};
let query = Query::new(&grammar, source)
.map_err(|e| anyhow::anyhow!("that is not a valid tree-sitter pattern: {e:?}"))?;
let names = query.capture_names().to_vec();
let scoped = filter.and_language(&language)?;
let files = store::file_rows(db, scoped.groups(), store::FileSort::Path, false, i64::MAX)?;
let mut parser = Parser::new();
parser.set_language(&grammar)?;
let mut out = Vec::new();
let mut parsed = 0usize;
let mut truncated = false;
let mut skipped = 0usize;
for file in files.iter().take(MAX_FILES) {
if out.len() >= MAX_MATCHES {
truncated = true;
break;
}
let Some(text) = file_text(db, &file.path)? else {
continue;
};
let Some(tree) = parser.parse(&text, None) else {
continue;
};
parsed += 1;
let mut cursor = QueryCursor::new();
let mut hits = cursor.matches(&query, tree.root_node(), text.as_bytes());
while let Some(m) = hits.next() {
for capture in m.captures() {
if out.len() >= MAX_MATCHES {
truncated = true;
break;
}
if skipped < offset {
skipped += 1;
continue;
}
let node = capture.node;
out.push(Match {
path: file.path.clone(),
capture: names
.get(capture.index as usize)
.map(|n| n.to_string())
.unwrap_or_default(),
start_line: node.start_position().row as u32 + 1,
end_line: node.end_position().row as u32 + 1,
text: first_line(&text[node.byte_range()]),
store: None,
});
}
}
}
if files.len() > MAX_FILES {
truncated = true;
}
Ok(Matches {
language,
matches: out,
files: parsed,
truncated,
skipped,
})
}
pub fn grep(
db: &rusqlite::Connection,
source: &str,
filter: &Filter,
offset: usize,
) -> Result<Matches> {
let regex = regex::Regex::new(source)
.or_else(|_| regex::Regex::new(®ex::escape(source)))
.map_err(|e| anyhow::anyhow!("that query cannot be searched for: {e}"))?;
let files = store::file_rows(db, filter.groups(), store::FileSort::Path, false, i64::MAX)?;
let mut out = Vec::new();
let mut searched = 0usize;
let mut truncated = false;
let mut skipped = 0usize;
'files: for file in &files {
let Some(text) = file_text(db, &file.path)? else {
continue;
};
searched += 1;
let mut symbols = None;
for (i, line) in text.lines().enumerate() {
if !regex.is_match(line) {
continue;
}
if out.len() >= MAX_MATCHES {
truncated = true;
break 'files;
}
if skipped < offset {
skipped += 1;
continue;
}
if symbols.is_none() {
symbols = Some(
store::symbols_in_files(db, std::slice::from_ref(&file.path))?
.remove(&file.path)
.unwrap_or_default(),
);
}
let defs = symbols.as_deref().unwrap_or_default();
let at = i as u32 + 1;
out.push(Match {
path: file.path.clone(),
capture: crate::enclosing_definition(defs, at, at)
.map(|(_, name, _)| name)
.unwrap_or_default(),
start_line: at,
end_line: at,
text: first_line(line),
store: None,
});
}
}
Ok(Matches {
language: String::new(),
matches: out,
files: searched,
truncated,
skipped,
})
}
fn file_text(db: &rusqlite::Connection, path: &str) -> Result<Option<String>> {
let chunks = store::chunks_overlapping(db, path, 1, u32::MAX)?;
if chunks.is_empty() {
return Ok(None);
}
let mut lines: std::collections::BTreeMap<u32, String> = Default::default();
for chunk in &chunks {
for (offset, line) in chunk.text.lines().enumerate() {
lines
.entry(chunk.start_line + offset as u32)
.or_insert_with(|| line.to_string());
}
}
let Some(&last) = lines.keys().next_back() else {
return Ok(None);
};
let text = (1..=last)
.map(|n| lines.get(&n).map_or("", String::as_str))
.collect::<Vec<_>>()
.join("\n");
Ok(Some(text))
}
fn first_line(raw: &str) -> String {
const WIDTH: usize = 120;
let line = raw.lines().next().unwrap_or("").trim();
if line.chars().count() <= WIDTH {
return line.to_string();
}
line.chars().take(WIDTH).collect::<String>() + "…"
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn grep_finds_every_line_and_names_its_definition() {
let db = rusqlite::Connection::open_in_memory().unwrap();
db.execute_batch(store::SCHEMA).unwrap();
db.execute_batch(
"INSERT INTO files (id, path, hash, bytes, indexed_at) VALUES (1, '/r/a.rs', 'h', 1, 0);
INSERT INTO chunks (id, file_id, ord, start_line, end_line, text)
VALUES (1, 1, 0, 1, 4, 'fn one() {\n // NOTE first\n}\nlet h = Row {');
INSERT INTO chunks (id, file_id, ord, start_line, end_line, text)
VALUES (2, 1, 1, 3, 6, '}\nlet h = Row {\nfn two() { // HACK second }\n');
INSERT INTO symbols (file_id, name, qualified, kind, start_line, end_line)
VALUES (1, 'one', 'one', 'function', 1, 3), (1, 'two', 'two', 'function', 5, 5);",
)
.unwrap();
let found = grep(&db, r"\b(NOTE|HACK)\b", &Filter::default(), 0).unwrap();
let rows: Vec<_> = found
.matches
.iter()
.map(|m| (m.start_line, m.capture.as_str()))
.collect();
assert_eq!(rows, [(2, "one"), (5, "two")]);
let literal = grep(&db, "Row {", &Filter::default(), 0).unwrap();
assert_eq!(literal.matches.len(), 1, "{:?}", literal.matches);
assert_eq!(literal.matches[0].start_line, 4);
}
#[test]
fn grep_searches_a_file_that_does_not_start_at_line_one() {
let db = rusqlite::Connection::open_in_memory().unwrap();
db.execute_batch(store::SCHEMA).unwrap();
db.execute_batch(
"INSERT INTO files (id, path, hash, bytes, indexed_at) VALUES (1, '/r/a.md', 'h', 1, 0);
INSERT INTO chunks (id, file_id, ord, start_line, end_line, text)
VALUES (1, 1, 0, 3, 4, '# Title\nNEEDLE on four');
INSERT INTO chunks (id, file_id, ord, start_line, end_line, text)
VALUES (2, 1, 1, 7, 8, 'text\nNEEDLE on eight');",
)
.unwrap();
let found = grep(&db, "NEEDLE", &Filter::default(), 0).unwrap();
let lines: Vec<u32> = found.matches.iter().map(|m| m.start_line).collect();
assert_eq!(lines, [4, 8], "{:?}", found.matches);
}
#[test]
fn an_invalid_pattern_is_an_error_rather_than_an_empty_answer() {
let db = rusqlite::Connection::open_in_memory().unwrap();
db.execute_batch(store::SCHEMA).unwrap();
let err = run(&db, "rust", "(this is not a pattern", &Filter::default(), 0)
.expect_err("an unbalanced pattern does not compile");
assert!(
err.to_string().contains("not a valid tree-sitter pattern"),
"{err}"
);
}
#[test]
fn every_advertised_language_compiles_a_pattern() {
let db = rusqlite::Connection::open_in_memory().unwrap();
db.execute_batch(store::SCHEMA).unwrap();
let mut broken: Vec<String> = Vec::new();
for entry in crate::filter::LANGUAGES {
if !crate::graph::has_graph(entry.name) {
continue;
}
if let Err(e) = run(&db, entry.name, "(_) @node", &Filter::default(), 0) {
broken.push(format!("{}: {e}", entry.name));
}
}
assert!(broken.is_empty(), "{}", broken.join("\n"));
}
#[test]
fn a_language_with_no_grammar_says_which_have_one() {
let db = rusqlite::Connection::open_in_memory().unwrap();
db.execute_batch(store::SCHEMA).unwrap();
let err = run(&db, "cobol", "(identifier) @x", &Filter::default(), 0)
.expect_err("cobol carries no grammar in this release");
assert!(err.to_string().contains("rust"), "{err}");
}
}
#[cfg(test)]
mod offset_tests {
use super::*;
fn store_of(calls: usize) -> rusqlite::Connection {
let db = rusqlite::Connection::open_in_memory().unwrap();
db.execute_batch(store::SCHEMA).unwrap();
let mut body = String::from("fn main() {\n");
for i in 0..calls {
body.push_str(&format!(" call_{i}();\n"));
}
body.push_str("}\n");
let lines = body.lines().count() as u32;
let file = store::insert_file(&db, "/work/api/src/main.rs", "hash", 1, 0).unwrap();
store::insert_chunk(&db, file, 0, 1, lines, &body).unwrap();
db
}
#[test]
fn offset_continues_the_listing_and_repeats_exactly() {
let db = store_of(MAX_MATCHES + 40);
let pattern = "(call_expression function: (identifier) @f)";
let first = run(&db, "rust", pattern, &Filter::default(), 0).unwrap();
assert_eq!(first.matches.len(), MAX_MATCHES);
assert!(first.truncated);
let next = run(&db, "rust", pattern, &Filter::default(), MAX_MATCHES).unwrap();
assert_eq!(next.skipped, MAX_MATCHES);
assert_eq!(next.matches.len(), 40);
assert_eq!(
next.matches[0].text, "call_200",
"the continuation starts at the match after the cap"
);
assert!(
!first.matches.iter().any(|m| m.text == "call_200"),
"the two pages must not overlap"
);
let again = run(&db, "rust", pattern, &Filter::default(), MAX_MATCHES).unwrap();
let texts: Vec<&str> = next.matches.iter().map(|m| m.text.as_str()).collect();
let repeat: Vec<&str> = again.matches.iter().map(|m| m.text.as_str()).collect();
assert_eq!(texts, repeat);
}
#[test]
fn a_path_filter_narrows_which_files_are_parsed() {
let db = store_of(3);
let pattern = "(call_expression function: (identifier) @f)";
let all = run(&db, "rust", pattern, &Filter::default(), 0).unwrap();
assert_eq!(all.files, 1);
let elsewhere = Filter::new(&["src/other/**".to_string()], &[], &[]).unwrap();
let none = run(&db, "rust", pattern, &elsewhere, 0).unwrap();
assert_eq!(none.files, 0);
assert!(none.matches.is_empty());
let here = Filter::new(&["**/src/**".to_string()], &[], &[]).unwrap();
let some = run(&db, "rust", pattern, &here, 0).unwrap();
assert_eq!(some.matches.len(), 3);
}
}