use crate::{AppendEntry, MicroscopeReader, LAYER_NAMES};
#[derive(Debug, Clone)]
pub struct Query {
pub keywords: Vec<String>,
pub op: BoolOp,
pub layer_filter: Option<u8>,
pub depth_filter: Option<(u8, u8)>, pub spatial_filter: Option<(f32, f32, f32, f32)>, pub limit: usize,
}
#[derive(Debug, Clone, PartialEq)]
pub enum BoolOp {
And,
Or,
}
impl Default for Query {
fn default() -> Self {
Self {
keywords: Vec::new(),
op: BoolOp::And,
layer_filter: None,
depth_filter: None,
spatial_filter: None,
limit: 10,
}
}
}
pub fn parse(input: &str) -> Query {
let mut q = Query::default();
let mut remaining = input.trim();
while !remaining.is_empty() {
remaining = remaining.trim_start();
if remaining.is_empty() {
break;
}
if let Some(rest) = remaining.strip_prefix("layer:") {
let (val, rest2) = take_word(rest);
q.layer_filter = Some(crate::layer_to_id(&val));
remaining = rest2;
continue;
}
if let Some(rest) = remaining.strip_prefix("depth:") {
let (val, rest2) = take_word(rest);
if let Some((a, b)) = val.split_once("..") {
let lo = a.parse::<u8>().unwrap_or(0);
let hi = b.parse::<u8>().unwrap_or(8);
q.depth_filter = Some((lo, hi));
} else if let Ok(d) = val.parse::<u8>() {
q.depth_filter = Some((d, d));
}
remaining = rest2;
continue;
}
if let Some(rest) = remaining.strip_prefix("near:") {
let (val, rest2) = take_word(rest);
let parts: Vec<f32> = val.split(',').filter_map(|s| s.parse().ok()).collect();
if parts.len() >= 3 {
let r = if parts.len() >= 4 { parts[3] } else { 0.1 };
q.spatial_filter = Some((parts[0], parts[1], parts[2], r));
}
remaining = rest2;
continue;
}
if let Some(rest) = remaining.strip_prefix("limit:") {
let (val, rest2) = take_word(rest);
if let Ok(n) = val.parse::<usize>() {
q.limit = n;
}
remaining = rest2;
continue;
}
if let Some(rest) = remaining.strip_prefix("AND") {
if rest.is_empty() || rest.starts_with(' ') {
q.op = BoolOp::And;
remaining = rest;
continue;
}
}
if let Some(rest) = remaining.strip_prefix("OR") {
if rest.is_empty() || rest.starts_with(' ') {
q.op = BoolOp::Or;
remaining = rest;
continue;
}
}
if remaining.starts_with('"') {
if let Some(end) = remaining[1..].find('"') {
q.keywords.push(remaining[1..1 + end].to_lowercase());
remaining = &remaining[2 + end..];
continue;
}
}
let (word, rest2) = take_word(remaining);
if !word.is_empty() {
q.keywords.push(word.to_lowercase());
}
remaining = rest2;
}
q
}
fn take_word(s: &str) -> (String, &str) {
let s = s.trim_start();
let end = s.find(char::is_whitespace).unwrap_or(s.len());
(s[..end].to_string(), &s[end..])
}
#[derive(Debug)]
pub struct QueryResult {
pub score: f32,
pub block_idx: usize,
pub is_main: bool, }
pub fn execute(q: &Query, reader: &MicroscopeReader, appended: &[AppendEntry]) -> Vec<QueryResult> {
let mut results = Vec::new();
for i in 0..reader.block_count {
let h = reader.header(i);
if let Some(lid) = q.layer_filter {
if h.layer_id != lid {
continue;
}
}
if let Some((lo, hi)) = q.depth_filter {
if h.depth < lo || h.depth > hi {
continue;
}
}
if let Some((sx, sy, sz, sr)) = q.spatial_filter {
let dx = h.x - sx;
let dy = h.y - sy;
let dz = h.z - sz;
if dx * dx + dy * dy + dz * dz > sr * sr {
continue;
}
}
let text = reader.text(i).to_lowercase();
let score = keyword_score(&text, &q.keywords, &q.op);
if score > 0.0 {
results.push(QueryResult {
score,
block_idx: i,
is_main: true,
});
}
}
for (ai, entry) in appended.iter().enumerate() {
if let Some(lid) = q.layer_filter {
if entry.layer_id != lid {
continue;
}
}
if let Some((lo, hi)) = q.depth_filter {
if entry.depth < lo || entry.depth > hi {
continue;
}
}
if let Some((sx, sy, sz, sr)) = q.spatial_filter {
let dx = entry.x - sx;
let dy = entry.y - sy;
let dz = entry.z - sz;
if dx * dx + dy * dy + dz * dz > sr * sr {
continue;
}
}
let text = entry.text.to_lowercase();
let score = keyword_score(&text, &q.keywords, &q.op);
if score > 0.0 {
results.push(QueryResult {
score,
block_idx: ai + 1_000_000,
is_main: false,
});
}
}
results.sort_by(|a, b| b.score.partial_cmp(&a.score).unwrap());
results.truncate(q.limit);
results
}
fn keyword_score(text: &str, keywords: &[String], op: &BoolOp) -> f32 {
if keywords.is_empty() {
return 1.0;
}
let hits: Vec<bool> = keywords
.iter()
.map(|kw| text.contains(kw.as_str()))
.collect();
let hit_count = hits.iter().filter(|&&h| h).count();
match op {
BoolOp::And => {
if hit_count == keywords.len() {
hit_count as f32
} else {
0.0
}
}
BoolOp::Or => hit_count as f32,
}
}
#[allow(dead_code)]
pub fn layer_name(id: u8) -> &'static str {
LAYER_NAMES.get(id as usize).unwrap_or(&"?")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_simple() {
let q = parse("hello");
assert_eq!(q.keywords, vec!["hello"]);
assert_eq!(q.op, BoolOp::And);
}
#[test]
fn test_parse_quoted() {
let q = parse("\"hello world\"");
assert_eq!(q.keywords, vec!["hello world"]);
}
#[test]
fn test_parse_filters() {
let q = parse("layer:long_term depth:2..5 \"memory\"");
assert_eq!(q.layer_filter, Some(1)); assert_eq!(q.depth_filter, Some((2, 5)));
assert_eq!(q.keywords, vec!["memory"]);
}
#[test]
fn test_parse_bool_op() {
let q = parse("\"foo\" OR \"bar\"");
assert_eq!(q.op, BoolOp::Or);
assert_eq!(q.keywords, vec!["foo", "bar"]);
}
#[test]
fn test_parse_limit() {
let q = parse("limit:20 hello");
assert_eq!(q.limit, 20);
}
#[test]
fn test_parse_spatial() {
let q = parse("near:0.2,0.3,0.1,0.05 test");
let (x, y, z, r) = q.spatial_filter.unwrap();
assert!((x - 0.2).abs() < 0.001);
assert!((y - 0.3).abs() < 0.001);
assert!((z - 0.1).abs() < 0.001);
assert!((r - 0.05).abs() < 0.001);
}
#[test]
fn test_keyword_score_and() {
assert_eq!(
keyword_score(
"hello world",
&["hello".into(), "world".into()],
&BoolOp::And
),
2.0
);
assert_eq!(
keyword_score("hello", &["hello".into(), "world".into()], &BoolOp::And),
0.0
);
}
#[test]
fn test_keyword_score_or() {
assert_eq!(
keyword_score("hello", &["hello".into(), "world".into()], &BoolOp::Or),
1.0
);
assert_eq!(
keyword_score("nope", &["hello".into(), "world".into()], &BoolOp::Or),
0.0
);
}
}