use std::collections::HashMap;
use std::fmt;
use std::ops::Range;
use nucleo_matcher::pattern::{CaseMatching, Normalization, Pattern};
use nucleo_matcher::{Config, Matcher, Utf32Str, Utf32String};
use serde::{Deserialize, Serialize};
use crate::common::{Block, Message, Meta, Role, Tool, ToolOutput};
use crate::{Common, HarnessId, Span, Transcript};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Mode {
Fuzzy,
Substring,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Case {
Smart,
Insensitive,
Sensitive,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Origin {
User,
Assistant,
Thinking,
ToolUse,
ToolResult,
Meta,
}
impl Origin {
pub const ALL: [Origin; 6] = [
Origin::User,
Origin::Assistant,
Origin::Thinking,
Origin::ToolUse,
Origin::ToolResult,
Origin::Meta,
];
pub const DEFAULT: [Origin; 5] = [
Origin::User,
Origin::Assistant,
Origin::Thinking,
Origin::ToolUse,
Origin::Meta,
];
const fn bit(self) -> u8 {
1 << (self as u8)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(default)]
pub struct Query {
pub pattern: String,
pub mode: Mode,
pub case: Case,
pub origins: Vec<Origin>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub harnesses: Option<Vec<HarnessId>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub limit: Option<usize>,
#[serde(
default = "default_hits_per_doc",
skip_serializing_if = "Option::is_none"
)]
pub hits_per_doc: Option<usize>,
}
#[allow(clippy::unnecessary_wraps)] fn default_hits_per_doc() -> Option<usize> {
Some(8)
}
impl Default for Query {
fn default() -> Query {
Query::fuzzy("")
}
}
impl Query {
#[must_use]
pub fn fuzzy(pattern: impl Into<String>) -> Query {
Query::new(pattern, Mode::Fuzzy)
}
#[must_use]
pub fn substring(pattern: impl Into<String>) -> Query {
Query::new(pattern, Mode::Substring)
}
fn new(pattern: impl Into<String>, mode: Mode) -> Query {
Query {
pattern: pattern.into(),
mode,
case: Case::Smart,
origins: Origin::DEFAULT.to_vec(),
harnesses: None,
limit: None,
hits_per_doc: default_hits_per_doc(),
}
}
fn compile(&self) -> Compiled {
match self.mode {
Mode::Fuzzy => {
let case = match self.case {
Case::Smart => CaseMatching::Smart,
Case::Insensitive => CaseMatching::Ignore,
Case::Sensitive => CaseMatching::Respect,
};
Compiled::Fuzzy {
pattern: Pattern::parse(&self.pattern, case, Normalization::Smart),
exact: Needle::new(self.pattern.trim(), self.case),
}
}
Mode::Substring => Compiled::Substring(Needle::new(&self.pattern, self.case)),
}
}
fn origin_mask(&self) -> u8 {
self.origins.iter().fold(0, |m, o| m | o.bit())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Hit {
pub span: Span,
pub block: usize,
pub origin: Origin,
pub line: String,
pub highlights: Vec<Range<u32>>,
pub score: u32,
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct DocKey {
pub harness: HarnessId,
pub id: String,
}
impl fmt::Display for DocKey {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}:{}", self.harness, self.id)
}
}
#[derive(Debug)]
pub struct DocMatch<'i> {
pub key: &'i DocKey,
pub meta: &'i Meta,
pub score: u32,
pub hits: Vec<Hit>,
}
#[must_use]
pub fn search(transcript: &Transcript<Common>, query: &Query) -> Vec<Hit> {
let lines = extract(&transcript.meta, &transcript.body);
let pattern = query.compile();
if pattern_is_empty(&pattern, query) {
Vec::new()
} else {
let mask = query.origin_mask();
let mut matcher = Matcher::new(Config::DEFAULT);
let mut indices = Vec::new();
lines
.iter()
.filter(|line| line.origin.bit() & mask != 0)
.filter_map(|line| line.hit(&pattern, &mut matcher, &mut indices))
.take(query.limit.unwrap_or(usize::MAX))
.collect()
}
}
#[derive(Default)]
pub struct Index {
docs: Vec<Doc>,
by_key: HashMap<DocKey, usize>,
}
struct Doc {
key: DocKey,
meta: Meta,
lines: Vec<Line>,
chars: usize,
}
pub struct Extracted(Doc);
impl Extracted {
#[must_use]
pub fn new(key: DocKey, transcript: &Transcript<Common>) -> Extracted {
let lines = extract(&transcript.meta, &transcript.body);
Extracted(Doc {
key,
meta: transcript.meta.clone(),
chars: lines.iter().map(|l| l.text.len()).sum(),
lines,
})
}
}
impl Index {
#[must_use]
pub fn new() -> Index {
Index::default()
}
#[must_use]
pub fn len(&self) -> usize {
self.docs.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.docs.is_empty()
}
#[must_use]
pub fn lines(&self) -> usize {
self.docs.iter().map(|d| d.lines.len()).sum()
}
#[must_use]
pub fn chars(&self) -> usize {
self.docs.iter().map(|d| d.chars).sum()
}
pub fn insert(&mut self, key: DocKey, transcript: &Transcript<Common>) {
self.insert_extracted(Extracted::new(key, transcript));
}
pub fn insert_extracted(&mut self, extracted: Extracted) {
let Extracted(doc) = extracted;
if let Some(&i) = self.by_key.get(&doc.key) {
self.docs[i] = doc;
} else {
self.by_key.insert(doc.key.clone(), self.docs.len());
self.docs.push(doc);
}
}
pub fn remove(&mut self, key: &DocKey) -> bool {
match self.by_key.remove(key) {
None => false,
Some(i) => {
self.docs.swap_remove(i);
if let Some(moved) = self.docs.get(i) {
self.by_key.insert(moved.key.clone(), i);
}
true
}
}
}
#[must_use]
pub fn query(&self, query: &Query) -> Vec<DocMatch<'_>> {
let pattern = query.compile();
if pattern_is_empty(&pattern, query) {
self.all_docs(query)
} else {
let mask = query.origin_mask();
let mut scored = self.score_all(&pattern, mask, query);
scored.sort_by(|a, b| {
b.1.cmp(&a.1).then_with(|| {
self.docs[b.0]
.meta
.timestamp
.cmp(&self.docs[a.0].meta.timestamp)
})
});
if let Some(limit) = query.limit {
scored.truncate(limit);
}
self.materialize(scored, &pattern, query)
}
}
#[cfg(not(target_arch = "wasm32"))]
fn materialize(
&self,
mut scored: Vec<Scored>,
pattern: &Compiled,
query: &Query,
) -> Vec<DocMatch<'_>> {
let workers = std::thread::available_parallelism().map_or(1, std::num::NonZero::get);
if workers > 1 && scored.len() >= 64 {
let chunk = scored.len().div_ceil(workers);
std::thread::scope(|scope| {
let handles: Vec<_> = scored
.chunks_mut(chunk)
.map(|part| scope.spawn(move || self.materialize_chunk(part, pattern, query)))
.collect();
handles
.into_iter()
.flat_map(|h| h.join().unwrap_or_default())
.collect()
})
} else {
self.materialize_chunk(&mut scored, pattern, query)
}
}
#[cfg(target_arch = "wasm32")]
fn materialize(
&self,
mut scored: Vec<Scored>,
pattern: &Compiled,
query: &Query,
) -> Vec<DocMatch<'_>> {
self.materialize_chunk(&mut scored, pattern, query)
}
fn materialize_chunk(
&self,
scored: &mut [Scored],
pattern: &Compiled,
query: &Query,
) -> Vec<DocMatch<'_>> {
let mut matcher = Matcher::new(Config::DEFAULT);
let mut indices = Vec::new();
scored
.iter_mut()
.map(|(d, best, hit_lines)| {
let doc = &self.docs[*d];
let hits = select_hits(std::mem::take(hit_lines), query.hits_per_doc)
.into_iter()
.filter_map(|(l, _)| doc.lines[l].hit(pattern, &mut matcher, &mut indices))
.collect();
DocMatch {
key: &doc.key,
meta: &doc.meta,
score: *best,
hits,
}
})
.collect()
}
#[cfg(not(target_arch = "wasm32"))]
fn score_all(&self, pattern: &Compiled, mask: u8, query: &Query) -> Vec<Scored> {
let workers = std::thread::available_parallelism().map_or(1, std::num::NonZero::get);
if workers > 1 && self.lines() >= 64 * 1024 {
let shards = self.shards(workers);
std::thread::scope(|scope| {
let handles: Vec<_> = shards
.into_iter()
.map(|shard| {
let docs = &self.docs[shard.clone()];
scope.spawn(move || score_docs(docs, shard.start, pattern, mask, query))
})
.collect();
handles
.into_iter()
.flat_map(|h| h.join().unwrap_or_default())
.collect()
})
} else {
score_docs(&self.docs, 0, pattern, mask, query)
}
}
#[cfg(target_arch = "wasm32")]
fn score_all(&self, pattern: &Compiled, mask: u8, query: &Query) -> Vec<Scored> {
score_docs(&self.docs, 0, pattern, mask, query)
}
#[cfg(not(target_arch = "wasm32"))]
fn shards(&self, workers: usize) -> Vec<std::ops::Range<usize>> {
let target = (self.chars() / workers).max(1);
let mut shards = Vec::with_capacity(workers);
let mut start = 0;
let mut acc = 0;
for (i, doc) in self.docs.iter().enumerate() {
acc += doc.chars;
if acc >= target && shards.len() + 1 < workers {
shards.push(start..i + 1);
start = i + 1;
acc = 0;
}
}
if start < self.docs.len() {
shards.push(start..self.docs.len());
}
shards
}
fn all_docs(&self, query: &Query) -> Vec<DocMatch<'_>> {
let mut all: Vec<&Doc> = self.docs.iter().filter(|d| d.selected(query)).collect();
all.sort_by_key(|doc| std::cmp::Reverse(doc.meta.timestamp));
if let Some(limit) = query.limit {
all.truncate(limit);
}
all.into_iter()
.map(|doc| DocMatch {
key: &doc.key,
meta: &doc.meta,
score: 0,
hits: Vec::new(),
})
.collect()
}
}
impl Doc {
fn selected(&self, query: &Query) -> bool {
query
.harnesses
.as_ref()
.is_none_or(|hs| hs.contains(&self.key.harness))
}
}
type Scored = (usize, u32, Vec<(usize, u32)>);
fn score_docs(
docs: &[Doc],
base: usize,
pattern: &Compiled,
mask: u8,
query: &Query,
) -> Vec<Scored> {
let mut matcher = Matcher::new(Config::DEFAULT);
docs.iter()
.enumerate()
.filter(|(_, doc)| doc.selected(query))
.filter_map(|(d, doc)| {
let hit_lines: Vec<(usize, u32)> = doc
.lines
.iter()
.enumerate()
.filter(|(_, line)| line.origin.bit() & mask != 0)
.filter_map(|(l, line)| {
pattern
.score(line.text.slice(..), &mut matcher)
.map(|score| (l, score))
})
.collect();
let best = hit_lines.iter().map(|&(_, score)| score).max();
best.map(|best| (base + d, best, hit_lines))
})
.collect()
}
fn select_hits(mut hit_lines: Vec<(usize, u32)>, cap: Option<usize>) -> Vec<(usize, u32)> {
if let Some(cap) = cap
&& hit_lines.len() > cap
{
hit_lines.sort_unstable_by_key(|&(_, score)| std::cmp::Reverse(score));
hit_lines.truncate(cap);
hit_lines.sort_unstable_by_key(|&(line, _)| line);
}
hit_lines
}
fn pattern_is_empty(_: &Compiled, query: &Query) -> bool {
query.pattern.trim().is_empty()
}
enum Compiled {
Fuzzy { pattern: Pattern, exact: Needle },
Substring(Needle),
}
const EXACT_BONUS: u32 = 1 << 16;
impl Compiled {
fn score(&self, hay: Utf32Str<'_>, matcher: &mut Matcher) -> Option<u32> {
match self {
Compiled::Fuzzy { pattern, exact } => pattern
.score(hay, matcher)
.map(|score| score + exact.find(hay).map_or(0, |_| EXACT_BONUS)),
Compiled::Substring(needle) => needle
.find(hay)
.map(|start| needle.score_at(hay, start, matcher)),
}
}
fn indices(&self, hay: Utf32Str<'_>, matcher: &mut Matcher, out: &mut Vec<u32>) -> Option<u32> {
match self {
Compiled::Fuzzy { pattern, exact } => match exact.find(hay) {
Some(start) => pattern.score(hay, matcher).map(|score| {
out.extend(start..start + exact.len_u32());
score + EXACT_BONUS
}),
None => pattern.indices(hay, matcher, out),
},
Compiled::Substring(needle) => needle.find(hay).map(|start| {
out.extend(start..start + needle.len_u32());
needle.score_at(hay, start, matcher)
}),
}
}
}
struct Needle {
chars: Vec<char>,
bytes: Option<Vec<u8>>,
utf32: Utf32String,
sensitive: bool,
}
impl Needle {
fn new(pattern: &str, case: Case) -> Needle {
let sensitive = match case {
Case::Sensitive => true,
Case::Insensitive => false,
Case::Smart => pattern.chars().any(char::is_uppercase),
};
let chars: Vec<char> = pattern
.chars()
.map(|c| if sensitive { c } else { fold(c) })
.collect();
let bytes = chars
.iter()
.all(char::is_ascii)
.then(|| chars.iter().map(|&c| c as u8).collect());
Needle {
chars,
bytes,
utf32: Utf32String::from(pattern),
sensitive,
}
}
fn len_u32(&self) -> u32 {
u32::try_from(self.chars.len()).unwrap_or(u32::MAX)
}
#[allow(clippy::cast_possible_truncation)] fn find(&self, hay: Utf32Str<'_>) -> Option<u32> {
match self.chars.len() {
0 => None,
n if n > hay.len() => None,
n => match hay {
Utf32Str::Ascii(bytes) => self.bytes.as_deref().and_then(|needle| {
if self.sensitive {
memchr::memmem::find(bytes, needle).map(|i| i as u32)
} else {
let (first, upper) = (needle[0], needle[0].to_ascii_uppercase());
memchr::memchr2_iter(first, upper, &bytes[..=bytes.len() - n])
.find(|&at| bytes[at..at + n].eq_ignore_ascii_case(needle))
.map(|at| at as u32)
}
}),
Utf32Str::Unicode(chars) => chars
.windows(n)
.position(|w| {
w.iter().zip(&self.chars).all(|(&c, &want)| {
if self.sensitive {
c == want
} else {
fold(c) == want
}
})
})
.map(|i| i as u32),
},
}
}
fn score_at(&self, hay: Utf32Str<'_>, start: u32, matcher: &mut Matcher) -> u32 {
let m = matcher.exact_match(
hay.slice_u32(start..start + self.len_u32()),
self.utf32.slice(..),
);
m.map_or(u32::from(SCORE_FALLBACK), u32::from)
}
}
const SCORE_FALLBACK: u16 = 16;
fn fold(c: char) -> char {
let mut lower = c.to_lowercase();
match (lower.next(), lower.next()) {
(Some(l), None) => l,
(Some(_), Some(_)) | (None, _) => c,
}
}
struct Line {
message: u32,
block: u32,
origin: Origin,
text: Utf32String,
}
impl Line {
fn hit(
&self,
pattern: &Compiled,
matcher: &mut Matcher,
indices: &mut Vec<u32>,
) -> Option<Hit> {
indices.clear();
let score = pattern.indices(self.text.slice(..), matcher, indices)?;
indices.sort_unstable();
indices.dedup();
Some(Hit {
span: match self.origin {
Origin::Meta => Span(0..0),
_ => Span(self.message as usize..self.message as usize + 1),
},
block: self.block as usize,
origin: self.origin,
line: self.text.to_string(),
highlights: merge_spans(indices),
score,
})
}
}
fn merge_spans(indices: &[u32]) -> Vec<Range<u32>> {
let mut spans: Vec<Range<u32>> = Vec::new();
for &i in indices {
match spans.last_mut() {
Some(last) if last.end == i => last.end = i + 1,
Some(_) | None => spans.push(i..i + 1),
}
}
spans
}
fn extract(meta: &Meta, messages: &[Message]) -> Vec<Line> {
let mut lines = Vec::new();
let mut push = |message: u32, block: u32, origin: Origin, text: &str| {
for line in text.lines() {
let line = line.trim_end();
if !line.trim_start().is_empty() {
lines.push(Line {
message,
block,
origin,
text: Utf32String::from(line),
});
}
}
};
for text in [&meta.title, &meta.cwd, &meta.git_branch]
.into_iter()
.flatten()
{
push(0, 0, Origin::Meta, text);
}
for (m, message) in messages.iter().enumerate() {
let m = u32::try_from(m).unwrap_or(u32::MAX);
for (b, block) in message.content.iter().enumerate() {
let b = u32::try_from(b).unwrap_or(u32::MAX);
match block {
Block::Text { text } => {
let origin = match message.role {
Role::User => Origin::User,
Role::Assistant => Origin::Assistant,
};
push(m, b, origin, text);
}
Block::Thinking { text, .. } => push(m, b, Origin::Thinking, text),
Block::ToolUse { tool, .. } => {
extract_tool(tool, |text| push(m, b, Origin::ToolUse, text));
}
Block::ToolResult { content, .. } => match content {
ToolOutput::Text(text) => push(m, b, Origin::ToolResult, text),
ToolOutput::Json(value) => {
push(m, b, Origin::ToolResult, &value.to_string());
}
},
Block::Image { .. } => {}
}
}
}
lines
}
fn extract_tool(tool: &Tool, mut push: impl FnMut(&str)) {
match tool {
Tool::Read { file_path, .. } => push(file_path),
Tool::Write { file_path, content } => {
push(file_path);
push(content);
}
Tool::Edit {
file_path,
old_string,
new_string,
..
} => {
push(file_path);
push(old_string);
push(new_string);
}
Tool::MultiEdit { file_path, edits } => {
push(file_path);
for edit in edits {
push(&edit.old_string);
push(&edit.new_string);
}
}
Tool::Bash {
command,
description,
..
} => {
push(command);
if let Some(description) = description {
push(description);
}
}
Tool::Raw { tool_name, input } => {
push(tool_name);
if !input.is_null() {
push(&input.to_string());
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn index_is_send_and_sync() {
fn assert_send_sync<T: Send + Sync>() {}
assert_send_sync::<Index>();
}
#[test]
fn merge_spans_groups_consecutive_indices() {
assert_eq!(merge_spans(&[0, 1, 2, 5, 6, 9]), vec![0..3, 5..7, 9..10]);
assert_eq!(merge_spans(&[]), Vec::<Range<u32>>::new());
}
}