use std::collections::{BTreeSet, HashMap};
use redb::ReadableDatabase;
use sinter_core::{Node, NodeId};
use crate::error::StoreError;
use crate::store::{INTERN, NAME_NODES, Store, TOKENS_WORDS};
pub(crate) fn trigrams(name: &str) -> Vec<String> {
let lower: Vec<char> = name.to_lowercase().chars().collect();
if lower.len() < 3 {
return vec![lower.iter().collect()];
}
let mut grams: Vec<String> = lower.windows(3).map(|w| w.iter().collect()).collect();
grams.sort();
grams.dedup();
grams
}
pub(crate) fn node_tokens(node: &Node) -> BTreeSet<String> {
let mut words = BTreeSet::new();
for text in [
node.name.as_str(),
node.signature.as_str(),
node.file.as_str(),
node.doc.as_deref().unwrap_or(""),
] {
for ident in text
.split(|c: char| !c.is_alphanumeric())
.filter(|s| !s.is_empty())
{
for sub in camel_split(ident) {
if sub.chars().count() >= 2 {
words.insert(sub);
}
}
let whole = ident.to_lowercase();
if whole.chars().count() >= 2 {
words.insert(whole);
}
}
}
words
}
fn camel_split(ident: &str) -> Vec<String> {
let chars: Vec<char> = ident.chars().collect();
let mut words = Vec::new();
let mut cur = String::new();
for (i, &c) in chars.iter().enumerate() {
let boundary = i > 0
&& c.is_uppercase()
&& (!chars[i - 1].is_uppercase() || chars.get(i + 1).is_some_and(|n| n.is_lowercase()));
if boundary && !cur.is_empty() {
words.push(std::mem::take(&mut cur));
}
cur.extend(c.to_lowercase());
}
if !cur.is_empty() {
words.push(cur);
}
words
}
impl Store {
pub fn nodes_named(&self, name: &str) -> Result<Vec<Node>, StoreError> {
let txn = self.db.begin_read()?;
let table = txn.open_multimap_table(NAME_NODES)?;
let mut interned = BTreeSet::new();
for guard in table.get(name)? {
interned.insert(guard?.value());
}
drop(table);
drop(txn);
self.decode_ids(interned)
}
pub fn nodes_with_token(&self, word: &str) -> Result<Vec<Node>, StoreError> {
self.decode_ids(self.token_ids([word].into_iter())?)
}
pub fn candidates_for_terms(&self, terms: &[String]) -> Result<Vec<Node>, StoreError> {
let words = terms.iter().flat_map(|term| {
let singular = term
.strip_suffix('s')
.filter(|singular| !singular.is_empty());
[Some(term.as_str()), singular].into_iter().flatten()
});
self.decode_ids(self.token_ids(words)?)
}
fn token_ids<'a>(
&self,
words: impl Iterator<Item = &'a str>,
) -> Result<BTreeSet<u32>, StoreError> {
let txn = self.db.begin_read()?;
let table = txn.open_multimap_table(TOKENS_WORDS)?;
let mut ids = BTreeSet::new();
for word in words {
for guard in table.get(word)? {
ids.insert(guard?.value());
}
}
Ok(ids)
}
fn decode_ids(&self, interned: BTreeSet<u32>) -> Result<Vec<Node>, StoreError> {
let txn = self.db.begin_read()?;
let table = txn.open_table(INTERN)?;
let mut ids = Vec::new();
for i in interned {
if let Some(guard) = table.get(i)? {
ids.push(guard.value().to_string());
}
}
drop(table);
drop(txn);
ids.sort();
let mut nodes = Vec::new();
for id in ids {
if let Some(node) = self.node(&NodeId::new(id))? {
nodes.push(node);
}
}
Ok(nodes)
}
pub fn search(&self, query: &str, limit: usize) -> Result<Vec<Node>, StoreError> {
let txn = self.db.begin_read()?;
let table = txn.open_multimap_table(crate::store::TRIGRAMS)?;
let mut hits: HashMap<u32, usize> = HashMap::new();
let query_grams = trigrams(query);
for gram in &query_grams {
for guard in table.get(gram.as_str())? {
*hits.entry(guard?.value()).or_default() += 1;
}
}
drop(table);
let intern = txn.open_table(INTERN)?;
let mut ranked: Vec<(String, usize)> = Vec::new();
for (interned, shared) in hits {
if let Some(guard) = intern.get(interned)? {
ranked.push((guard.value().to_string(), shared));
}
}
drop(intern);
drop(txn);
ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
let mut nodes = Vec::new();
for (id, shared) in ranked.into_iter().take(limit.max(1)) {
if shared * 2 >= query_grams.len()
&& let Some(node) = self.node(&NodeId::new(id))?
{
nodes.push(node);
}
}
Ok(nodes)
}
}
#[cfg(test)]
mod tests {
use sinter_core::{Node, NodeId, Span, SymbolKind};
use super::node_tokens;
fn node(name: &str, file: &str, signature: &str, doc: Option<&str>) -> Node {
Node {
id: NodeId::new(format!("{file}#{name}@0")),
kind: SymbolKind::Function,
name: name.to_string(),
file: file.to_string(),
span: Span { start: 0, end: 1 },
signature: signature.to_string(),
doc: doc.map(str::to_string),
}
}
fn has(words: &std::collections::BTreeSet<String>, expect: &[&str]) {
for w in expect {
assert!(words.contains(*w), "missing {w:?} in {words:?}");
}
}
#[test]
fn camel_case_splits_and_keeps_whole_identifier() {
let words = node_tokens(&node("PlayerCharacterV2", "a.rs", "", None));
has(&words, &["player", "character", "v2", "playercharacterv2"]);
}
#[test]
fn acronym_word_boundary() {
let words = node_tokens(&node("HTTPServer", "a.rs", "", None));
has(&words, &["http", "server", "httpserver"]);
}
#[test]
fn snake_case_and_signature_and_doc() {
let words = node_tokens(&node(
"climb_state",
"a.rs",
"fn climb_state(input: MoveInput)",
Some("Main traversal controller."),
));
has(
&words,
&[
"climb",
"state",
"move",
"input",
"traversal",
"controller",
"fn",
],
);
}
#[test]
fn path_segments_indexed() {
let words = node_tokens(&node("f", "src/player/ClimbComponent.test.ts", "", None));
has(
&words,
&[
"src",
"player",
"climb",
"component",
"climbcomponent",
"test",
"ts",
],
);
}
#[test]
fn short_subwords_dropped_but_short_wholes_kept_at_two_chars() {
let words = node_tokens(&node("aB", "x.rs", "", None));
assert!(words.contains("ab"), "{words:?}");
assert!(!words.contains("a") && !words.contains("b"), "{words:?}");
}
}