diff --git a/.claude/skills/verify/SKILL.md b/.claude/skills/verify/SKILL.md new file mode 100644 index 0000000..8b83830 --- /dev/null +++ b/.claude/skills/verify/SKILL.md @@ -0,0 +1,17 @@ +--- +name: verify +summary: Exercise csp indexing, persisted-index search, and cache reuse through the built CLI. +--- + +# Verify csp + +1. Build the user-facing binary with `cargo build -p code-search-please --bin csp`. +2. Use a scratch directory outside the repository for all generated indexes and HOME state. +3. Drive persisted-index behavior: + - `target/debug/csp index crates/csp/src --out /index` + - `target/debug/csp search "load_or_build_index cache" --index /index --top-k 3` +4. Drive auto-cache behavior twice with an isolated home: + - `HOME=/home target/debug/csp search "DEFAULT_MODEL_NAME" crates/csp/src/indexing --top-k 1` + - Repeat the same command to exercise cache reuse. +5. Probe errors with a missing `--index` path; expect a clear `Index not found` message and exit status 1. +6. Do not use tests as runtime verification evidence; run the project quality gate separately. diff --git a/.please/docs/references/semble.md b/.please/docs/references/semble.md index 54616a4..632238a 100644 --- a/.please/docs/references/semble.md +++ b/.please/docs/references/semble.md @@ -338,7 +338,7 @@ Clean two-layer split: | strong / moderate / mild penalty | `0.3` / `0.5` / `0.7` | same | `ranking/penalties.rs` | | file saturation threshold / decay | `1` / `0.5` | `1` / `0.5` | `ranking/penalties.rs` | | max file bytes | `1_000_000` | `1_000_000` | `index/files.py` / `indexing/create.rs` | -| default model | `minishlab/potion-code-16M` | same (real + stub) | `utils.py` / `indexing/dense.rs` | +| default model | `minishlab/potion-code-16M-v2` (semble#219) | same (real + stub); cache keyed on model_id | `utils.py` / `indexing/dense.rs` | | MCP in-mem LRU | `10` | `10` | `mcp.py` / `csp::mcp` | | cache dir mode | — | `0o700` | `indexing/cache.rs` | diff --git a/README.ko.md b/README.ko.md index c60750b..f82c33e 100644 --- a/README.ko.md +++ b/README.ko.md @@ -497,7 +497,7 @@ for (const { chunk, score } of results) { ## 동작 원리 -`csp`는 [tree-sitter](https://tree-sitter.github.io/)로 각 파일을 코드 인지형 청크로 분할한 뒤, 두 개의 상호 보완적인 검색기로 쿼리를 모든 청크와 점수화합니다. 의미 유사도를 위한 정적 [Model2Vec](https://github.com/MinishLab/model2vec) 임베딩(코드 특화 `potion-code-16M` 모델 사용)과, 식별자/API명 등 어휘 매칭을 위한 BM25입니다. 두 점수 리스트는 Reciprocal Rank Fusion(RRF)으로 결합됩니다. +`csp`는 [tree-sitter](https://tree-sitter.github.io/)로 각 파일을 코드 인지형 청크로 분할한 뒤, 두 개의 상호 보완적인 검색기로 쿼리를 모든 청크와 점수화합니다. 의미 유사도를 위한 정적 [Model2Vec](https://github.com/MinishLab/model2vec) 임베딩(코드 특화 `potion-code-16M-v2` 모델 사용)과, 식별자/API명 등 어휘 매칭을 위한 BM25입니다. 두 점수 리스트는 Reciprocal Rank Fusion(RRF)으로 결합됩니다. 결합 후에는 코드 인지형 신호들로 결과를 재정렬합니다. diff --git a/README.md b/README.md index e5ce79a..e764a79 100644 --- a/README.md +++ b/README.md @@ -497,7 +497,7 @@ for (const { chunk, score } of results) { ## How it works -`csp` splits each file into code-aware chunks using [tree-sitter](https://tree-sitter.github.io/), then scores every query against the chunks with two complementary retrievers: static [Model2Vec](https://github.com/MinishLab/model2vec) embeddings using the code-specialized `potion-code-16M` model for semantic similarity, and BM25 for lexical matches on identifiers and API names. The two score lists are fused with Reciprocal Rank Fusion (RRF). +`csp` splits each file into code-aware chunks using [tree-sitter](https://tree-sitter.github.io/), then scores every query against the chunks with two complementary retrievers: static [Model2Vec](https://github.com/MinishLab/model2vec) embeddings using the code-specialized `potion-code-16M-v2` model for semantic similarity, and BM25 for lexical matches on identifiers and API names. The two score lists are fused with Reciprocal Rank Fusion (RRF). After fusing, results are reranked with a set of code-aware signals: diff --git a/crates/csp/src/indexing/cache.rs b/crates/csp/src/indexing/cache.rs index 5e84c2b..422e073 100644 --- a/crates/csp/src/indexing/cache.rs +++ b/crates/csp/src/indexing/cache.rs @@ -179,13 +179,33 @@ pub fn compute_content_hash(files: &[CacheFile]) -> String { let mut hasher = Sha256::new(); for file in sorted { - let len16 = file.path.encode_utf16().count(); - hasher.update(format!("{len16}:{}", file.path).as_bytes()); - hasher.update(&file.content); + update_content_hash(&mut hasher, &file.path, &file.content); } to_hex(&hasher.finalize()) } +/// Compute the same content hash while reading one file at a time, avoiding +/// retention of every file body for large repositories. Unreadable files are +/// skipped, matching the indexer's existing best-effort collection behavior. +pub(crate) fn compute_content_hash_from_paths(mut files: Vec<(String, PathBuf)>) -> String { + files.sort_by(|a, b| a.0.cmp(&b.0)); + + let mut hasher = Sha256::new(); + for (relative_path, path) in files { + let Ok(content) = std::fs::read(path) else { + continue; + }; + update_content_hash(&mut hasher, &relative_path, &content); + } + to_hex(&hasher.finalize()) +} + +fn update_content_hash(hasher: &mut Sha256, path: &str, content: &[u8]) { + let len16 = path.encode_utf16().count(); + hasher.update(format!("{len16}:{path}").as_bytes()); + hasher.update(content); +} + /// Directories from `home` down to `leaf` (inclusive), home-first. When `leaf` /// is not under `home`, only `leaf` is returned. fn chain_to(leaf: &Path, home: &Path) -> Vec { @@ -380,6 +400,22 @@ mod tests { assert_eq!(a, b); } + #[test] + fn streamed_content_hash_matches_in_memory_hash() { + let dir = tempdir().unwrap(); + let a_path = dir.path().join("a.ts"); + let b_path = dir.path().join("b.ts"); + std::fs::write(&a_path, "one").unwrap(); + std::fs::write(&b_path, "two").unwrap(); + + let streamed = compute_content_hash_from_paths(vec![ + ("b.ts".to_string(), b_path), + ("a.ts".to_string(), a_path), + ]); + let in_memory = compute_content_hash(&[cfile("a.ts", "one"), cfile("b.ts", "two")]); + assert_eq!(streamed, in_memory); + } + #[test] fn content_hash_is_hex_sha256() { let h = compute_content_hash(&[cfile("a.ts", "x")]); diff --git a/crates/csp/src/indexing/cache_orchestrator.rs b/crates/csp/src/indexing/cache_orchestrator.rs new file mode 100644 index 0000000..663e1e7 --- /dev/null +++ b/crates/csp/src/indexing/cache_orchestrator.rs @@ -0,0 +1,139 @@ +//! Build-or-reuse orchestration for the global on-disk index cache. Port of +//! `src/indexing/cache.ts`. + +use std::path::{Path, PathBuf}; + +use crate::chunking::source::DESIRED_CHUNK_LENGTH_CHARS; +use crate::indexing::cache::{ + compute_content_hash_from_paths, ensure_cache_dir, resolve_cache_dir, CacheLocation, +}; +use crate::indexing::create::MAX_FILE_BYTES; +use crate::indexing::dense::DEFAULT_MODEL_NAME; +use crate::indexing::file_walker::walk_files; +use crate::indexing::files::get_extensions; +use crate::indexing::index::{normalize_content, parse_manifest, CspIndex, LoadOptions}; +use crate::types::ContentType; +use crate::utils::is_git_url; + +/// Options for [`load_or_build_index`]. +#[derive(Debug, Clone, Default)] +pub struct LoadOrBuildOptions { + pub base_dir: Option, + pub git_ref: Option, + pub content: Option>, + pub model_path: Option, +} + +fn normalized_hash_path(path: &Path) -> String { + path.to_string_lossy().replace('\\', "/") +} + +/// Collect the source files `from_path` would index as sorted hash inputs. +/// File contents are read later, one at a time, by the hashing helper. +fn collect_source_paths(root: &Path, content: &[ContentType]) -> Vec<(String, PathBuf)> { + let resolved = get_extensions(content, None); + let ext_refs: Vec<&str> = resolved.iter().map(String::as_str).collect(); + let mut files = Vec::new(); + for file_path in walk_files(root, &ext_refs, &[]) { + let Ok(meta) = std::fs::metadata(&file_path) else { + continue; + }; + if meta.len() > MAX_FILE_BYTES { + continue; + } + let rel = file_path.strip_prefix(root).unwrap_or(&file_path); + files.push((normalized_hash_path(rel), file_path)); + } + files +} + +/// Load a cached index for `source` if fresh, else build, persist, and return. +pub fn load_or_build_index(source: &str, options: &LoadOrBuildOptions) -> Result { + let content = normalize_content(options.content.clone()); + let is_git = is_git_url(source); + + let location = CacheLocation { + base_dir: options.base_dir.clone(), + git_ref: options.git_ref.clone(), + }; + let cache_dir = resolve_cache_dir(source, &content, &location); + let base_only = CacheLocation { + base_dir: options.base_dir.clone(), + git_ref: None, + }; + ensure_cache_dir(&cache_dir, &base_only)?; + + // Local sources: the source-file hash is the cache-validity oracle. Git + // sources are URL+ref keyed (no cheap live hash). + let source_hash = if is_git { + None + } else { + Some(compute_content_hash_from_paths(collect_source_paths( + Path::new(source), + &content, + ))) + }; + + // The resolved model name the index would be (re)built with. A cache built + // with a different model must not be reused — its vectors are incompatible + // (mirrors semble#219 bumping the default to the `-v2` weights). + let expected_model = options.model_path.as_deref().unwrap_or(DEFAULT_MODEL_NAME); + + if let Some(cached) = try_reuse(&cache_dir, is_git, source_hash.as_deref(), expected_model) { + return Ok(cached); + } + + let load_options = LoadOptions { + model_path: options.model_path.clone(), + content: Some(content), + }; + let index = if is_git { + CspIndex::from_git(source, &load_options, options.git_ref.as_deref())? + } else { + CspIndex::from_path(Path::new(source), &load_options)? + }; + index.save(&cache_dir, source_hash.as_deref())?; + Ok(index) +} + +/// Reuse a cached index when present and valid, else `None`. +fn try_reuse( + cache_dir: &Path, + is_git: bool, + source_hash: Option<&str>, + expected_model: &str, +) -> Option { + let manifest_path = cache_dir.join("manifest.json"); + if !manifest_path.exists() { + return None; + } + let raw = std::fs::read_to_string(&manifest_path).ok()?; + let value: serde_json::Value = serde_json::from_str(&raw).ok()?; + let manifest = parse_manifest(&value).ok()?; + // A chunk_size change re-chunks every file, so a cache built with a different + // target length is stale even if the source files are byte-identical. + if manifest.chunk_size != Some(DESIRED_CHUNK_LENGTH_CHARS as u32) { + return None; + } + // A model change makes the persisted vectors incompatible with queries + // embedded by the new model, so a cache built with a different model is stale. + if manifest.model_id != expected_model { + return None; + } + // The persisted vectors and live query model must use the same runtime + // implementation. This distinguishes a real Model2Vec cache from an offline + // deterministic-stub cache even when both share the same requested model id. + let (query_model, _) = crate::indexing::dense::load_model(Some(expected_model)); + if manifest.model_kind.as_deref() != Some(query_model.kind()) { + return None; + } + // Local sources additionally validate the live source-file hash; git sources + // are URL+ref keyed (no cheap live hash). + if !is_git && Some(manifest.content_hash.as_str()) != source_hash { + return None; + } + CspIndex::load_from_disk(cache_dir).ok() +} + +#[cfg(test)] +mod tests; diff --git a/crates/csp/src/indexing/cache_orchestrator/tests.rs b/crates/csp/src/indexing/cache_orchestrator/tests.rs new file mode 100644 index 0000000..9c5b440 --- /dev/null +++ b/crates/csp/src/indexing/cache_orchestrator/tests.rs @@ -0,0 +1,147 @@ +use super::*; +use crate::indexing::cache::{resolve_cache_dir, CacheLocation}; +use crate::indexing::dense::{make_stub_model, SelectableBasicBackend}; +use crate::indexing::index::{CspIndexState, DEFAULT_CONTENT}; +use crate::indexing::sparse::Bm25Index; +use crate::types::Chunk; +use tempfile::tempdir; + +fn make_chunk(file_path: &str, content: &str) -> Chunk { + Chunk { + content: content.to_string(), + file_path: file_path.to_string(), + start_line: 1, + end_line: 10, + language: Some("typescript".to_string()), + } +} + +fn build_index(chunks: Vec) -> CspIndex { + let model = make_stub_model(4); + let vectors = vec![vec![1.0, 0.0, 0.0, 0.0]; chunks.len()]; + CspIndex::new(CspIndexState { + model, + bm25_index: Bm25Index::build(&vec![vec!["x".to_string()]; chunks.len()]), + semantic_index: SelectableBasicBackend::from_vectors(vectors).unwrap(), + chunks, + model_path: "test-model".to_string(), + root: None, + content: DEFAULT_CONTENT.to_vec(), + }) +} + +#[test] +fn hash_path_normalizes_platform_separators() { + assert_eq!(normalized_hash_path(Path::new(r"src\lib.rs")), "src/lib.rs"); + assert_eq!(normalized_hash_path(Path::new("src/lib.rs")), "src/lib.rs"); +} + +#[test] +fn try_reuse_rejects_stale_chunk_size() { + let chunks = vec![make_chunk("a.ts", "A")]; + let idx = build_index(chunks); + let dir = tempdir().unwrap(); + idx.save(dir.path(), Some("deadbeef")).unwrap(); + + // Fresh cache (matching hash + current chunk_size) is reused. + assert!(try_reuse(dir.path(), false, Some("deadbeef"), "test-model").is_some()); + + // Rewrite the manifest with a different chunk_size → stale → rebuild. + let manifest_path = dir.path().join("manifest.json"); + let raw = std::fs::read_to_string(&manifest_path).unwrap(); + let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); + value["chunkSize"] = serde_json::json!(9999); + std::fs::write(&manifest_path, value.to_string()).unwrap(); + assert!(try_reuse(dir.path(), false, Some("deadbeef"), "test-model").is_none()); + + // A manifest predating the field (absent chunkSize) is also stale. + let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); + value.as_object_mut().unwrap().remove("chunkSize"); + std::fs::write(&manifest_path, value.to_string()).unwrap(); + assert!(try_reuse(dir.path(), false, Some("deadbeef"), "test-model").is_none()); +} + +#[test] +fn try_reuse_rejects_mismatched_model_kind() { + let idx = build_index(vec![make_chunk("a.ts", "A")]); + let dir = tempdir().unwrap(); + idx.save(dir.path(), Some("deadbeef")).unwrap(); + + let manifest_path = dir.path().join("manifest.json"); + let raw = std::fs::read_to_string(&manifest_path).unwrap(); + let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); + value["modelKind"] = serde_json::json!("static"); + std::fs::write(&manifest_path, value.to_string()).unwrap(); + + assert!(try_reuse(dir.path(), false, Some("deadbeef"), "test-model").is_none()); +} + +#[test] +fn try_reuse_rejects_stale_model() { + let chunks = vec![make_chunk("a.ts", "A")]; + let idx = build_index(chunks); + let dir = tempdir().unwrap(); + idx.save(dir.path(), Some("deadbeef")).unwrap(); + + // Same model → reused; a different model (e.g. after a default bump to + // `-v2`) → stale even with a matching hash + chunk_size. + assert!(try_reuse(dir.path(), false, Some("deadbeef"), "test-model").is_some()); + assert!(try_reuse(dir.path(), false, Some("deadbeef"), "other-model").is_none()); +} + +// --- load_or_build_index (cache.ts loadOrBuildIndex parity) --- + +#[test] +fn load_or_build_miss_then_hit_then_invalidate() { + let home = tempdir().unwrap(); + let src = tempdir().unwrap(); + let base = home.path().join(".csp"); + std::fs::write( + src.path().join("a.ts"), + "export function alpha() { return 1 }\n", + ) + .unwrap(); + let src_str = src.path().to_string_lossy().into_owned(); + let opts = LoadOrBuildOptions { + base_dir: Some(base.clone()), + ..Default::default() + }; + + // Miss: builds and writes a manifest. + let first = load_or_build_index(&src_str, &opts).unwrap(); + assert!(!first.chunks.is_empty()); + let cache_dir = resolve_cache_dir( + &src_str, + DEFAULT_CONTENT, + &CacheLocation { + base_dir: Some(base.clone()), + git_ref: None, + }, + ); + let manifest_path = cache_dir.join("manifest.json"); + assert!(manifest_path.exists()); + + // Add an ignored sentinel to make reuse observable: a rebuild rewrites the + // manifest from the typed struct and removes this extra field. + let raw = std::fs::read_to_string(&manifest_path).unwrap(); + let mut manifest: serde_json::Value = serde_json::from_str(&raw).unwrap(); + manifest["reuseSentinel"] = serde_json::json!(true); + std::fs::write(&manifest_path, manifest.to_string()).unwrap(); + + // Hit: a second call reuses the cache without rewriting the manifest. + let second = load_or_build_index(&src_str, &opts).unwrap(); + assert_eq!(second.chunks.len(), first.chunks.len()); + let after_hit: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(&manifest_path).unwrap()).unwrap(); + assert_eq!(after_hit["reuseSentinel"], serde_json::json!(true)); + + // Invalidation: add a file → content hash changes → rebuild reflects it. + std::fs::write( + src.path().join("b.ts"), + "export function beta() { return 2 }\n", + ) + .unwrap(); + let third = load_or_build_index(&src_str, &opts).unwrap(); + assert!(third.chunks.iter().any(|c| c.file_path == "b.ts")); + assert!(third.chunks.len() >= first.chunks.len()); +} diff --git a/crates/csp/src/indexing/dense.rs b/crates/csp/src/indexing/dense.rs index 334e5b9..c574822 100644 --- a/crates/csp/src/indexing/dense.rs +++ b/crates/csp/src/indexing/dense.rs @@ -13,18 +13,17 @@ //! candidate-selector filtering and a csp-local on-disk format. use std::collections::HashMap; -use std::path::Path; use std::sync::{Arc, LazyLock, Mutex}; use model2vec_rs::model::StaticModel; -use serde::{Deserialize, Serialize}; use crate::types::Chunk; -/// Default Model2Vec model name (kept identical to semble for parity). -pub const DEFAULT_MODEL_NAME: &str = "minishlab/potion-code-16M"; +/// Default Model2Vec model name (kept identical to semble for parity; +/// semble#219 bumped this to the `-v2` weights). +pub const DEFAULT_MODEL_NAME: &str = "minishlab/potion-code-16M-v2"; -/// Stub embedding dimension (the real `potion-code-16M` emits 256-dim vectors). +/// Stub embedding dimension (the real `potion-code-16M-v2` emits 256-dim vectors). const DEFAULT_STUB_DIM: usize = 256; /// Deterministic 32-bit FNV-1a over UTF-16 code units (matches JS `charCodeAt`). @@ -112,6 +111,14 @@ impl Model { } } + /// Stable runtime implementation identifier persisted with cached vectors. + pub fn kind(&self) -> &'static str { + match self { + Model::Static { .. } => "static", + Model::Stub { .. } => "stub", + } + } + /// Embedding dimension. pub fn dim(&self) -> usize { match self { @@ -184,469 +191,9 @@ pub fn embed_chunks(model: &Model, chunks: &[Chunk]) -> Vec> { model.encode(&texts) } -// --------------------------------------------------------------------------- -// SelectableBasicBackend -// --------------------------------------------------------------------------- - -/// Backend arguments. For parity only cosine is supported. -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct BasicArgs { - #[serde(default, skip_serializing_if = "Option::is_none")] - pub metric: Option, -} - -impl Default for BasicArgs { - fn default() -> Self { - Self { - metric: Some("cosine".to_string()), - } - } -} - -/// L2-normalise a vector in place (f64 accumulation, f32 storage — matching TS). -/// Zero vectors stay zero. -fn normalize_in_place(v: &mut [f32]) { - let mut n: f64 = 0.0; - for &x in v.iter() { - n += (x as f64) * (x as f64); - } - n = n.sqrt(); - if n == 0.0 { - return; - } - for x in v.iter_mut() { - *x = ((*x as f64) / n) as f32; - } -} - -fn dot(a: &[f32], b: &[f32]) -> f64 { - let mut s = 0.0; - for i in 0..a.len() { - s += (a[i] as f64) * (b[i] as f64); - } - s -} - -/// In-memory cosine vector backend with optional candidate-selector filtering — -/// port of `SelectableBasicBackend(CosineBasicBackend)`. -#[derive(Debug)] -pub struct SelectableBasicBackend { - /// Pre-normalised row vectors. - pub vectors: Vec>, - pub arguments: BasicArgs, - pub dim: usize, -} - -impl SelectableBasicBackend { - /// Build from raw vectors (defensively copied and L2-normalised so cosine - /// distance reduces to `1 - dot`). Errors on inconsistent dimensions. - pub fn new(vectors: Vec>, arguments: BasicArgs) -> Result { - let dim = vectors.first().map(Vec::len).unwrap_or(0); - let mut normalized = Vec::with_capacity(vectors.len()); - for v in vectors { - if v.len() != dim { - return Err(format!( - "Inconsistent vector dimensions: expected {dim}, got {}", - v.len() - )); - } - let mut copy = v; - normalize_in_place(&mut copy); - normalized.push(copy); - } - Ok(Self { - vectors: normalized, - arguments, - dim, - }) - } - - /// Convenience constructor with default (cosine) arguments. - pub fn from_vectors(vectors: Vec>) -> Result { - Self::new(vectors, BasicArgs::default()) - } +mod backend; - /// Batched k-NN query. Returns, per query, `[(chunk_index, cosine_distance)]` - /// sorted by ascending distance. `selector` constrains results to a pool. - pub fn query( - &self, - query_vectors: &[Vec], - k: usize, - selector: Option<&[u32]>, - ) -> Result>, String> { - if k < 1 { - return Err(format!("k should be >= 1, is now {k}")); - } - - let num_vectors = self.vectors.len(); - let mut effective_k = k.min(num_vectors); - if let Some(sel) = selector { - for &idx in sel { - if idx as usize >= num_vectors { - return Err(format!( - "Selector index out of bounds: {idx} (total vectors: {num_vectors})" - )); - } - } - effective_k = effective_k.min(sel.len()); - } - - let mut out: Vec> = Vec::with_capacity(query_vectors.len()); - if effective_k == 0 { - out.resize(query_vectors.len(), Vec::new()); - return Ok(out); - } - - for raw in query_vectors { - if raw.len() != self.dim { - return Err(format!( - "Query vector dimension mismatch: expected {}, got {}", - self.dim, - raw.len() - )); - } - let mut q = raw.clone(); - normalize_in_place(&mut q); - - let pool_size = selector.map(<[u32]>::len).unwrap_or(num_vectors); - // (pool_idx, distance) pairs, stably sorted by ascending distance. - let mut pairs: Vec<(usize, f64)> = (0..pool_size) - .map(|i| { - let vec_idx = selector.map_or(i, |s| s[i] as usize); - (i, 1.0 - dot(&q, &self.vectors[vec_idx])) - }) - .collect(); - // total_cmp is NaN-safe (a stray NaN distance can't panic the sort). - pairs.sort_by(|a, b| a.1.total_cmp(&b.1)); - pairs.truncate(effective_k); - - let mapped: Vec<(usize, f64)> = pairs - .into_iter() - .map(|(pool_idx, dist)| (selector.map_or(pool_idx, |s| s[pool_idx] as usize), dist)) - .collect(); - out.push(mapped); - } - - Ok(out) - } - - /// Persist vectors + args to `/vectors.bin` (flat little-endian f32) and - /// `/args.json`. - pub fn save(&self, dir: &Path) -> std::io::Result<()> { - std::fs::create_dir_all(dir)?; - let mut bytes = Vec::with_capacity(self.vectors.len() * self.dim * 4); - for row in &self.vectors { - for &x in row { - bytes.extend_from_slice(&x.to_le_bytes()); - } - } - std::fs::write(dir.join("vectors.bin"), &bytes)?; - - let meta = BackendMeta { - rows: self.vectors.len(), - dim: self.dim, - arguments: self.arguments.clone(), - }; - let json = serde_json::to_string(&meta) - .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; - std::fs::write(dir.join("args.json"), json) - } - - /// Inverse of [`save`](Self::save). - pub fn load(dir: &Path) -> Result { - let meta_raw = std::fs::read_to_string(dir.join("args.json")).map_err(|e| e.to_string())?; - let meta: BackendMeta = serde_json::from_str(&meta_raw).map_err(|e| e.to_string())?; - - let bytes = std::fs::read(dir.join("vectors.bin")).map_err(|e| e.to_string())?; - let expected = meta.rows * meta.dim * 4; - if bytes.len() != expected { - return Err(format!( - "Vector file size mismatch: expected {expected} bytes, got {}", - bytes.len() - )); - } - - let mut vectors = Vec::with_capacity(meta.rows); - for r in 0..meta.rows { - let mut row = Vec::with_capacity(meta.dim); - for c in 0..meta.dim { - let off = (r * meta.dim + c) * 4; - let arr: [u8; 4] = bytes[off..off + 4].try_into().expect("4-byte chunk"); - row.push(f32::from_le_bytes(arr)); - } - vectors.push(row); - } - Self::new(vectors, meta.arguments) - } -} - -#[derive(Serialize, Deserialize)] -struct BackendMeta { - rows: usize, - dim: usize, - arguments: BasicArgs, -} +pub use backend::{BasicArgs, SelectableBasicBackend}; #[cfg(test)] -mod tests { - use super::*; - use tempfile::tempdir; - - fn chunk(content: &str) -> Chunk { - Chunk { - content: content.to_string(), - file_path: "f.ts".to_string(), - start_line: 1, - end_line: 1, - language: None, - } - } - - // --- stub parity (golden vectors captured from the TS implementation) --- - - #[test] - fn fnv1a_matches_ts() { - assert_eq!(fnv1a("hello"), 1_335_831_723); - } - - #[test] - fn stub_embed_matches_ts_golden() { - // Golden values captured from the TS `stubEmbed` (Float32Array entries - // widened to f64); `as f32` reproduces the exact stored f32. - let expected_hello: [f64; 8] = [ - 0.085_591_696_202_754_97, - -0.438_301_533_460_617_07, - -0.693_752_408_027_648_9, - 0.431_218_117_475_509_64, - -0.016_508_268_192_410_47, - -0.213_292_211_294_174_2, - 0.267_603_516_578_674_3, - 0.126_279_816_031_456, - ]; - let hello = stub_embed("hello", 8); - for (got, want) in hello.iter().zip(&expected_hello) { - assert_eq!(*got, *want as f32); - } - - let expected_foo: [f64; 4] = [ - 0.054_837_439_209_222_794, - -0.873_466_372_489_929_2, - -0.401_930_719_614_028_93, - -0.269_260_287_284_851_1, - ]; - let foo = stub_embed("foo", 4); - for (got, want) in foo.iter().zip(&expected_foo) { - assert_eq!(*got, *want as f32); - } - } - - #[test] - fn stub_embed_is_unit_length() { - let v = stub_embed("anything", 256); - let norm: f64 = v - .iter() - .map(|&x| (x as f64) * (x as f64)) - .sum::() - .sqrt(); - assert!((norm - 1.0).abs() < 1e-5); - } - - // --- load_model / embed_chunks --- - - #[test] - fn load_model_defaults_path_via_seam() { - // Offline: inject a loader so no network/model download happens. - let (model, path) = load_model_with(None, |_| Ok(make_stub_model(7))); - assert_eq!(path, DEFAULT_MODEL_NAME); - assert!(model.dim() > 0); - } - - #[test] - fn load_model_resolves_distinct_paths_and_caches() { - // Distinct paths each load once; a repeat path is served from cache. - let (_, a) = load_model_with(Some("seam/path-X"), |_| Ok(make_stub_model(4))); - let (_, b) = load_model_with(Some("seam/path-Y"), |_| Ok(make_stub_model(4))); - // The loader must NOT fire for an already-cached path — panic proves it. - let (_, a2) = load_model_with(Some("seam/path-X"), |_| { - panic!("cached path must not reload") - }); - assert_eq!(a, "seam/path-X"); - assert_eq!(b, "seam/path-Y"); - assert_eq!(a2, "seam/path-X"); - } - - #[test] - fn load_model_falls_back_to_stub_on_error() { - let (model, path) = load_model_with(Some("seam/will-fail"), |_| Err("boom".to_string())); - assert_eq!(path, "seam/will-fail"); - assert_eq!(model.dim(), DEFAULT_STUB_DIM); // stub fallback - } - - /// Real Model2Vec load — downloads `minishlab/potion-code-16M` from HF on - /// first run, so it's network-gated and not part of the default suite. - /// Run with: `cargo test -p csp -- --ignored real_model2vec`. - #[test] - #[ignore = "network: downloads potion-code-16M from Hugging Face"] - fn real_model2vec_loads_and_embeds() { - let model = load_static(DEFAULT_MODEL_NAME).expect("load real model"); - assert!(model.dim() > 0); - let vecs = model.encode(&["fn main() {}".to_string(), "def main(): pass".to_string()]); - assert_eq!(vecs.len(), 2); - assert_eq!(vecs[0].len(), model.dim()); - assert_ne!(vecs[0], vecs[1]); - } - - #[test] - fn embed_empty_is_empty() { - let model = make_stub_model(8); - assert!(embed_chunks(&model, &[]).is_empty()); - } - - #[test] - fn embed_one_per_chunk() { - let model = make_stub_model(8); - let vectors = embed_chunks(&model, &[chunk("a"), chunk("b")]); - assert_eq!(vectors.len(), 2); - for v in &vectors { - assert_eq!(v.len(), 8); - } - } - - #[test] - fn embed_is_deterministic() { - let model = make_stub_model(16); - let v1 = embed_chunks(&model, &[chunk("same")]); - let v2 = embed_chunks(&model, &[chunk("same")]); - assert_eq!(v1, v2); - } - - #[test] - fn embed_differs_by_content() { - let model = make_stub_model(16); - let v1 = embed_chunks(&model, &[chunk("alpha")]); - let v2 = embed_chunks(&model, &[chunk("beta")]); - assert_ne!(v1, v2); - } - - // --- SelectableBasicBackend::query --- - - fn backend(n: usize, dim: usize) -> SelectableBasicBackend { - let model = make_stub_model(dim); - let vectors: Vec> = (0..n) - .map(|i| stub_embed(&format!("doc{i}"), dim)) - .collect(); - let _ = model; - SelectableBasicBackend::from_vectors(vectors).unwrap() - } - - #[test] - fn query_rejects_k_below_one() { - let b = backend(3, 8); - assert!(b.query(&[b.vectors[0].clone()], 0, None).is_err()); - } - - #[test] - fn new_rejects_inconsistent_dims() { - let v0 = stub_embed("x", 8); - let truncated = v0[..4].to_vec(); - let err = SelectableBasicBackend::from_vectors(vec![v0, truncated]).unwrap_err(); - assert!(err.contains("Inconsistent vector dimensions")); - } - - #[test] - fn query_rejects_dim_mismatch() { - let b = backend(3, 8); - let bad = vec![0f32; 4]; - let err = b.query(&[bad], 1, None).unwrap_err(); - assert!(err.contains("Query vector dimension mismatch")); - } - - #[test] - fn query_rejects_selector_out_of_bounds() { - let b = backend(3, 8); - let err = b.query(&[b.vectors[0].clone()], 1, Some(&[5])).unwrap_err(); - assert!(err.contains("Selector index out of bounds")); - } - - #[test] - fn query_returns_sorted_topk_with_self_nearest() { - let b = backend(3, 8); - let results = b.query(&[b.vectors[0].clone()], 3, None).unwrap(); - assert_eq!(results.len(), 1); - let hits = &results[0]; - assert_eq!(hits.len(), 3); - assert_eq!(hits[0].0, 0); - assert!(hits[0].1.abs() < 1e-5); - for i in 1..hits.len() { - assert!(hits[i].1 >= hits[i - 1].1); - } - } - - #[test] - fn query_respects_selector_pool() { - let b = backend(4, 8); - let results = b.query(&[b.vectors[0].clone()], 2, Some(&[1, 2])).unwrap(); - let hits = &results[0]; - assert_eq!(hits.len(), 2); - for (idx, _) in hits { - assert!(*idx == 1 || *idx == 2); - } - } - - #[test] - fn query_handles_multiple_queries() { - let b = backend(3, 8); - let results = b - .query(&[b.vectors[0].clone(), b.vectors[1].clone()], 1, None) - .unwrap(); - assert_eq!(results.len(), 2); - assert_eq!(results[0][0].0, 0); - assert_eq!(results[1][0].0, 1); - } - - #[test] - fn query_caps_k_at_num_vectors() { - let b = backend(2, 8); - let results = b.query(&[b.vectors[0].clone()], 5, None).unwrap(); - assert_eq!(results[0].len(), 2); - } - - // --- save / load --- - - #[test] - fn save_load_round_trips() { - let original = backend(3, 8); - let dir = tempdir().unwrap(); - original.save(dir.path()).unwrap(); - - let loaded = SelectableBasicBackend::load(dir.path()).unwrap(); - assert_eq!(loaded.vectors.len(), original.vectors.len()); - assert_eq!(loaded.dim, original.dim); - for (a, b) in loaded.vectors.iter().zip(&original.vectors) { - assert_eq!(a, b); - } - - let q = vec![original.vectors[0].clone()]; - let orig_hits: Vec = original.query(&q, 3, None).unwrap()[0] - .iter() - .map(|h| h.0) - .collect(); - let loaded_hits: Vec = loaded.query(&q, 3, None).unwrap()[0] - .iter() - .map(|h| h.0) - .collect(); - assert_eq!(orig_hits, loaded_hits); - } - - #[test] - fn load_rejects_truncated_vectors() { - let original = backend(3, 8); - let dir = tempdir().unwrap(); - original.save(dir.path()).unwrap(); - // Truncate vectors.bin to half its size. - let path = dir.path().join("vectors.bin"); - let bytes = std::fs::read(&path).unwrap(); - std::fs::write(&path, &bytes[..bytes.len() / 2]).unwrap(); - assert!(SelectableBasicBackend::load(dir.path()).is_err()); - } -} +mod tests; diff --git a/crates/csp/src/indexing/dense/backend.rs b/crates/csp/src/indexing/dense/backend.rs new file mode 100644 index 0000000..4b88694 --- /dev/null +++ b/crates/csp/src/indexing/dense/backend.rs @@ -0,0 +1,227 @@ +//! In-memory cosine vector backend with selector filtering and persistence. + +use std::path::Path; + +use serde::{Deserialize, Serialize}; + +/// Backend arguments. For parity only cosine is supported. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct BasicArgs { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metric: Option, +} + +impl Default for BasicArgs { + fn default() -> Self { + Self { + metric: Some("cosine".to_string()), + } + } +} + +/// L2-normalise a vector in place (f64 accumulation, f32 storage — matching TS). +/// Zero vectors stay zero. +fn normalize_in_place(v: &mut [f32]) { + let mut n: f64 = 0.0; + for &x in v.iter() { + n += (x as f64) * (x as f64); + } + n = n.sqrt(); + if n == 0.0 { + return; + } + for x in v.iter_mut() { + *x = ((*x as f64) / n) as f32; + } +} + +fn dot(a: &[f32], b: &[f32]) -> f64 { + let mut s = 0.0; + for i in 0..a.len() { + s += (a[i] as f64) * (b[i] as f64); + } + s +} + +/// In-memory cosine vector backend with optional candidate-selector filtering — +/// port of `SelectableBasicBackend(CosineBasicBackend)`. +#[derive(Debug)] +pub struct SelectableBasicBackend { + /// Pre-normalised row vectors. + pub vectors: Vec>, + pub arguments: BasicArgs, + pub dim: usize, +} + +impl SelectableBasicBackend { + /// Build from raw vectors (defensively copied and L2-normalised so cosine + /// distance reduces to `1 - dot`). Errors on inconsistent dimensions. + pub fn new(vectors: Vec>, arguments: BasicArgs) -> Result { + let dim = vectors.first().map(Vec::len).unwrap_or(0); + if !vectors.is_empty() && dim == 0 { + return Err( + "Vector dimension must be greater than 0 for a non-empty index".to_string(), + ); + } + let mut normalized = Vec::with_capacity(vectors.len()); + for v in vectors { + if v.len() != dim { + return Err(format!( + "Inconsistent vector dimensions: expected {dim}, got {}", + v.len() + )); + } + let mut copy = v; + normalize_in_place(&mut copy); + normalized.push(copy); + } + Ok(Self { + vectors: normalized, + arguments, + dim, + }) + } + + /// Convenience constructor with default (cosine) arguments. + pub fn from_vectors(vectors: Vec>) -> Result { + Self::new(vectors, BasicArgs::default()) + } + + /// Batched k-NN query. Returns, per query, `[(chunk_index, cosine_distance)]` + /// sorted by ascending distance. `selector` constrains results to a pool. + pub fn query( + &self, + query_vectors: &[Vec], + k: usize, + selector: Option<&[u32]>, + ) -> Result>, String> { + if k < 1 { + return Err(format!("k should be >= 1, is now {k}")); + } + + let num_vectors = self.vectors.len(); + let mut effective_k = k.min(num_vectors); + if let Some(sel) = selector { + for &idx in sel { + if idx as usize >= num_vectors { + return Err(format!( + "Selector index out of bounds: {idx} (total vectors: {num_vectors})" + )); + } + } + effective_k = effective_k.min(sel.len()); + } + + let mut out: Vec> = Vec::with_capacity(query_vectors.len()); + if effective_k == 0 { + out.resize(query_vectors.len(), Vec::new()); + return Ok(out); + } + + for raw in query_vectors { + if raw.len() != self.dim { + return Err(format!( + "Query vector dimension mismatch: expected {}, got {}", + self.dim, + raw.len() + )); + } + let mut q = raw.clone(); + normalize_in_place(&mut q); + + let pool_size = selector.map(<[u32]>::len).unwrap_or(num_vectors); + // (pool_idx, distance) pairs, stably sorted by ascending distance. + let mut pairs: Vec<(usize, f64)> = (0..pool_size) + .map(|i| { + let vec_idx = selector.map_or(i, |s| s[i] as usize); + (i, 1.0 - dot(&q, &self.vectors[vec_idx])) + }) + .collect(); + // total_cmp is NaN-safe (a stray NaN distance can't panic the sort). + pairs.sort_by(|a, b| a.1.total_cmp(&b.1)); + pairs.truncate(effective_k); + + let mapped: Vec<(usize, f64)> = pairs + .into_iter() + .map(|(pool_idx, dist)| (selector.map_or(pool_idx, |s| s[pool_idx] as usize), dist)) + .collect(); + out.push(mapped); + } + + Ok(out) + } + + /// Persist vectors + args to `/vectors.bin` (flat little-endian f32) and + /// `/args.json`. + pub fn save(&self, dir: &Path) -> std::io::Result<()> { + std::fs::create_dir_all(dir)?; + let mut bytes = Vec::with_capacity(self.vectors.len() * self.dim * 4); + for row in &self.vectors { + for &x in row { + bytes.extend_from_slice(&x.to_le_bytes()); + } + } + std::fs::write(dir.join("vectors.bin"), &bytes)?; + + let meta = BackendMeta { + rows: self.vectors.len(), + dim: self.dim, + arguments: self.arguments.clone(), + }; + let json = serde_json::to_string(&meta) + .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?; + std::fs::write(dir.join("args.json"), json) + } + + /// Inverse of [`save`](Self::save). + pub fn load(dir: &Path) -> Result { + let meta_raw = std::fs::read_to_string(dir.join("args.json")).map_err(|e| e.to_string())?; + let meta: BackendMeta = serde_json::from_str(&meta_raw).map_err(|e| e.to_string())?; + if meta.rows > 0 && meta.dim == 0 { + return Err( + "Invalid vector dimension: dim must be greater than 0 for a non-empty index" + .to_string(), + ); + } + + let bytes = std::fs::read(dir.join("vectors.bin")).map_err(|e| e.to_string())?; + let expected = meta + .rows + .checked_mul(meta.dim) + .and_then(|elements| elements.checked_mul(std::mem::size_of::())) + .ok_or_else(|| "Vector dimensions are too large (overflow)".to_string())?; + if bytes.len() != expected { + return Err(format!( + "Vector file size mismatch: expected {expected} bytes, got {}", + bytes.len() + )); + } + + let mut byte_chunks = bytes.chunks_exact(4); + let mut vectors = Vec::with_capacity(meta.rows); + for _ in 0..meta.rows { + let mut row = Vec::with_capacity(meta.dim); + for _ in 0..meta.dim { + let arr: [u8; 4] = byte_chunks + .next() + .expect("validated vector byte count") + .try_into() + .expect("4-byte chunk"); + row.push(f32::from_le_bytes(arr)); + } + vectors.push(row); + } + let mut backend = Self::new(vectors, meta.arguments)?; + if meta.rows == 0 { + backend.dim = meta.dim; + } + Ok(backend) + } +} + +#[derive(Serialize, Deserialize)] +struct BackendMeta { + rows: usize, + dim: usize, + arguments: BasicArgs, +} diff --git a/crates/csp/src/indexing/dense/tests.rs b/crates/csp/src/indexing/dense/tests.rs new file mode 100644 index 0000000..28255e1 --- /dev/null +++ b/crates/csp/src/indexing/dense/tests.rs @@ -0,0 +1,312 @@ +use super::*; +use tempfile::tempdir; + +fn chunk(content: &str) -> Chunk { + Chunk { + content: content.to_string(), + file_path: "f.ts".to_string(), + start_line: 1, + end_line: 1, + language: None, + } +} + +// --- stub parity (golden vectors captured from the TS implementation) --- + +#[test] +fn fnv1a_matches_ts() { + assert_eq!(fnv1a("hello"), 1_335_831_723); +} + +#[test] +fn stub_embed_matches_ts_golden() { + // Golden values captured from the TS `stubEmbed` (Float32Array entries + // widened to f64); `as f32` reproduces the exact stored f32. + let expected_hello: [f64; 8] = [ + 0.085_591_696_202_754_97, + -0.438_301_533_460_617_07, + -0.693_752_408_027_648_9, + 0.431_218_117_475_509_64, + -0.016_508_268_192_410_47, + -0.213_292_211_294_174_2, + 0.267_603_516_578_674_3, + 0.126_279_816_031_456, + ]; + let hello = stub_embed("hello", 8); + for (got, want) in hello.iter().zip(&expected_hello) { + assert_eq!(*got, *want as f32); + } + + let expected_foo: [f64; 4] = [ + 0.054_837_439_209_222_794, + -0.873_466_372_489_929_2, + -0.401_930_719_614_028_93, + -0.269_260_287_284_851_1, + ]; + let foo = stub_embed("foo", 4); + for (got, want) in foo.iter().zip(&expected_foo) { + assert_eq!(*got, *want as f32); + } +} + +#[test] +fn stub_embed_is_unit_length() { + let v = stub_embed("anything", 256); + let norm: f64 = v + .iter() + .map(|&x| (x as f64) * (x as f64)) + .sum::() + .sqrt(); + assert!((norm - 1.0).abs() < 1e-5); +} + +// --- load_model / embed_chunks --- + +#[test] +fn load_model_defaults_path_via_seam() { + // Offline: inject a loader so no network/model download happens. + let (model, path) = load_model_with(None, |_| Ok(make_stub_model(7))); + assert_eq!(path, DEFAULT_MODEL_NAME); + assert!(model.dim() > 0); +} + +#[test] +fn load_model_resolves_distinct_paths_and_caches() { + // Distinct paths each load once; a repeat path is served from cache. + let (_, a) = load_model_with(Some("seam/path-X"), |_| Ok(make_stub_model(4))); + let (_, b) = load_model_with(Some("seam/path-Y"), |_| Ok(make_stub_model(4))); + // The loader must NOT fire for an already-cached path — panic proves it. + let (_, a2) = load_model_with(Some("seam/path-X"), |_| { + panic!("cached path must not reload") + }); + assert_eq!(a, "seam/path-X"); + assert_eq!(b, "seam/path-Y"); + assert_eq!(a2, "seam/path-X"); +} + +#[test] +fn load_model_falls_back_to_stub_on_error() { + let (model, path) = load_model_with(Some("seam/will-fail"), |_| Err("boom".to_string())); + assert_eq!(path, "seam/will-fail"); + assert_eq!(model.dim(), DEFAULT_STUB_DIM); // stub fallback +} + +/// Real Model2Vec load — downloads `minishlab/potion-code-16M-v2` from HF on +/// first run, so it's network-gated and not part of the default suite. +/// Run with: `cargo test -p csp -- --ignored real_model2vec`. +#[test] +#[ignore = "network: downloads potion-code-16M-v2 from Hugging Face"] +fn real_model2vec_loads_and_embeds() { + let model = load_static(DEFAULT_MODEL_NAME).expect("load real model"); + assert!(model.dim() > 0); + let vecs = model.encode(&["fn main() {}".to_string(), "def main(): pass".to_string()]); + assert_eq!(vecs.len(), 2); + assert_eq!(vecs[0].len(), model.dim()); + assert_ne!(vecs[0], vecs[1]); +} + +#[test] +fn embed_empty_is_empty() { + let model = make_stub_model(8); + assert!(embed_chunks(&model, &[]).is_empty()); +} + +#[test] +fn embed_one_per_chunk() { + let model = make_stub_model(8); + let vectors = embed_chunks(&model, &[chunk("a"), chunk("b")]); + assert_eq!(vectors.len(), 2); + for v in &vectors { + assert_eq!(v.len(), 8); + } +} + +#[test] +fn embed_is_deterministic() { + let model = make_stub_model(16); + let v1 = embed_chunks(&model, &[chunk("same")]); + let v2 = embed_chunks(&model, &[chunk("same")]); + assert_eq!(v1, v2); +} + +#[test] +fn embed_differs_by_content() { + let model = make_stub_model(16); + let v1 = embed_chunks(&model, &[chunk("alpha")]); + let v2 = embed_chunks(&model, &[chunk("beta")]); + assert_ne!(v1, v2); +} + +// --- SelectableBasicBackend::query --- + +fn backend(n: usize, dim: usize) -> SelectableBasicBackend { + let model = make_stub_model(dim); + let vectors: Vec> = (0..n) + .map(|i| stub_embed(&format!("doc{i}"), dim)) + .collect(); + let _ = model; + SelectableBasicBackend::from_vectors(vectors).unwrap() +} + +#[test] +fn query_rejects_k_below_one() { + let b = backend(3, 8); + assert!(b.query(&[b.vectors[0].clone()], 0, None).is_err()); +} + +#[test] +fn new_rejects_non_empty_zero_dimension_vectors() { + let err = SelectableBasicBackend::from_vectors(vec![vec![]]).unwrap_err(); + assert!(err.contains("dimension must be greater than 0")); +} + +#[test] +fn new_rejects_inconsistent_dims() { + let v0 = stub_embed("x", 8); + let truncated = v0[..4].to_vec(); + let err = SelectableBasicBackend::from_vectors(vec![v0, truncated]).unwrap_err(); + assert!(err.contains("Inconsistent vector dimensions")); +} + +#[test] +fn query_rejects_dim_mismatch() { + let b = backend(3, 8); + let bad = vec![0f32; 4]; + let err = b.query(&[bad], 1, None).unwrap_err(); + assert!(err.contains("Query vector dimension mismatch")); +} + +#[test] +fn query_rejects_selector_out_of_bounds() { + let b = backend(3, 8); + let err = b.query(&[b.vectors[0].clone()], 1, Some(&[5])).unwrap_err(); + assert!(err.contains("Selector index out of bounds")); +} + +#[test] +fn query_returns_sorted_topk_with_self_nearest() { + let b = backend(3, 8); + let results = b.query(&[b.vectors[0].clone()], 3, None).unwrap(); + assert_eq!(results.len(), 1); + let hits = &results[0]; + assert_eq!(hits.len(), 3); + assert_eq!(hits[0].0, 0); + assert!(hits[0].1.abs() < 1e-5); + for i in 1..hits.len() { + assert!(hits[i].1 >= hits[i - 1].1); + } +} + +#[test] +fn query_respects_selector_pool() { + let b = backend(4, 8); + let results = b.query(&[b.vectors[0].clone()], 2, Some(&[1, 2])).unwrap(); + let hits = &results[0]; + assert_eq!(hits.len(), 2); + for (idx, _) in hits { + assert!(*idx == 1 || *idx == 2); + } +} + +#[test] +fn query_handles_multiple_queries() { + let b = backend(3, 8); + let results = b + .query(&[b.vectors[0].clone(), b.vectors[1].clone()], 1, None) + .unwrap(); + assert_eq!(results.len(), 2); + assert_eq!(results[0][0].0, 0); + assert_eq!(results[1][0].0, 1); +} + +#[test] +fn query_caps_k_at_num_vectors() { + let b = backend(2, 8); + let results = b.query(&[b.vectors[0].clone()], 5, None).unwrap(); + assert_eq!(results[0].len(), 2); +} + +// --- save / load --- + +#[test] +fn save_load_round_trips() { + let original = backend(3, 8); + let dir = tempdir().unwrap(); + original.save(dir.path()).unwrap(); + + let loaded = SelectableBasicBackend::load(dir.path()).unwrap(); + assert_eq!(loaded.vectors.len(), original.vectors.len()); + assert_eq!(loaded.dim, original.dim); + for (a, b) in loaded.vectors.iter().zip(&original.vectors) { + assert_eq!(a, b); + } + + let q = vec![original.vectors[0].clone()]; + let orig_hits: Vec = original.query(&q, 3, None).unwrap()[0] + .iter() + .map(|h| h.0) + .collect(); + let loaded_hits: Vec = loaded.query(&q, 3, None).unwrap()[0] + .iter() + .map(|h| h.0) + .collect(); + assert_eq!(orig_hits, loaded_hits); +} + +#[test] +fn load_preserves_dimension_for_empty_index() { + let dir = tempdir().unwrap(); + std::fs::write( + dir.path().join("args.json"), + r#"{"rows":0,"dim":8,"arguments":{"metric":"cosine"}}"#, + ) + .unwrap(); + std::fs::write(dir.path().join("vectors.bin"), []).unwrap(); + + let loaded = SelectableBasicBackend::load(dir.path()).unwrap(); + assert!(loaded.vectors.is_empty()); + assert_eq!(loaded.dim, 8); +} + +#[test] +fn load_rejects_zero_dimension_for_non_empty_index() { + let dir = tempdir().unwrap(); + std::fs::write( + dir.path().join("args.json"), + r#"{"rows":1,"dim":0,"arguments":{"metric":"cosine"}}"#, + ) + .unwrap(); + std::fs::write(dir.path().join("vectors.bin"), []).unwrap(); + + let err = SelectableBasicBackend::load(dir.path()).unwrap_err(); + assert!(err.contains("dim must be greater than 0")); +} + +#[test] +fn load_rejects_overflowing_dimensions() { + let dir = tempdir().unwrap(); + std::fs::write( + dir.path().join("args.json"), + format!( + r#"{{"rows":{},"dim":2,"arguments":{{"metric":"cosine"}}}}"#, + usize::MAX + ), + ) + .unwrap(); + std::fs::write(dir.path().join("vectors.bin"), []).unwrap(); + + let err = SelectableBasicBackend::load(dir.path()).unwrap_err(); + assert!(err.contains("overflow")); +} + +#[test] +fn load_rejects_truncated_vectors() { + let original = backend(3, 8); + let dir = tempdir().unwrap(); + original.save(dir.path()).unwrap(); + // Truncate vectors.bin to half its size. + let path = dir.path().join("vectors.bin"); + let bytes = std::fs::read(&path).unwrap(); + std::fs::write(&path, &bytes[..bytes.len() / 2]).unwrap(); + assert!(SelectableBasicBackend::load(dir.path()).is_err()); +} diff --git a/crates/csp/src/indexing/index.rs b/crates/csp/src/indexing/index.rs index 17f1185..1de41c1 100644 --- a/crates/csp/src/indexing/index.rs +++ b/crates/csp/src/indexing/index.rs @@ -1,6 +1,5 @@ //! `CspIndex` — the hybrid (dense + BM25) search orchestrator. Port of -//! `src/indexing/index.ts` (← semble `index/index.py`), plus the -//! `load_or_build_index` cache orchestration from `src/indexing/cache.ts`. +//! `src/indexing/index.ts` (← semble `index/index.py`). use std::collections::{BTreeMap, HashSet}; use std::fmt::Write as _; @@ -11,17 +10,11 @@ use serde::{Deserialize, Serialize}; use sha2::{Digest, Sha256}; use crate::chunking::source::DESIRED_CHUNK_LENGTH_CHARS; -use crate::indexing::cache::{ - compute_content_hash, ensure_cache_dir, resolve_cache_dir, CacheFile, CacheLocation, -}; -use crate::indexing::create::{create_index_from_path, CreateIndexOptions, MAX_FILE_BYTES}; +use crate::indexing::create::{create_index_from_path, CreateIndexOptions}; use crate::indexing::dense::{load_model, make_stub_model, Model, SelectableBasicBackend}; -use crate::indexing::file_walker::walk_files; -use crate::indexing::files::get_extensions; use crate::indexing::sparse::Bm25Index; use crate::search::{search as run_search, SearchOptions as RunSearchOptions, SearchResult}; use crate::types::{chunk_from_dict, chunk_to_dict, Chunk, ChunkDict, ContentType, IndexStats}; -use crate::utils::is_git_url; /// On-disk index schema version. pub const INDEX_SCHEMA_VERSION: u32 = 1; @@ -41,6 +34,9 @@ pub struct IndexManifest { pub source_id: Option, pub content: Vec, pub model_id: String, + /// Runtime model implementation used to build vectors (`static` or `stub`). + /// Absent in legacy manifests, which are conservatively treated as stale. + pub model_kind: Option, /// Target chunk length the index was built with. Changing it alters every /// chunk boundary, so a cache built with a different value must be rebuilt /// (mirrors semble `_metadata_matches`). `None` = built before this field @@ -86,7 +82,7 @@ pub struct CspIndex { pub content: Vec, } -fn normalize_content(content: Option>) -> Vec { +pub(crate) fn normalize_content(content: Option>) -> Vec { content.unwrap_or_else(|| DEFAULT_CONTENT.to_vec()) } @@ -295,6 +291,7 @@ impl CspIndex { source_id: self.root.clone(), content: self.content.clone(), model_id: self.model_path.clone(), + model_kind: Some(self.model.kind().to_string()), chunk_size: Some(DESIRED_CHUNK_LENGTH_CHARS as u32), }; let manifest_json = serde_json::to_string(&manifest).map_err(|e| e.to_string())?; @@ -431,6 +428,15 @@ pub fn parse_manifest(raw: &serde_json::Value) -> Result .and_then(serde_json::Value::as_str) .ok_or("Invalid manifest: modelId must be a string")? .to_string(); + let model_kind = match obj.get("modelKind") { + None | Some(serde_json::Value::Null) => None, + Some(serde_json::Value::String(kind)) if matches!(kind.as_str(), "static" | "stub") => { + Some(kind.clone()) + } + Some(_) => { + return Err("Invalid manifest: modelKind must be 'static', 'stub', or null".to_string()) + } + }; // Absent/null = built before the field existed → None (treated as a cache // mismatch by `try_reuse`). A present-but-non-numeric value is malformed. let chunk_size = obj @@ -460,481 +466,12 @@ pub fn parse_manifest(raw: &serde_json::Value) -> Result source_id, content, model_id, + model_kind, chunk_size, }) } -// --- load_or_build_index (cache.ts orchestration) --------------------------- - -/// Options for [`load_or_build_index`]. -#[derive(Debug, Clone, Default)] -pub struct LoadOrBuildOptions { - pub base_dir: Option, - pub git_ref: Option, - pub content: Option>, - pub model_path: Option, -} - -/// Collect the source files `from_path` would index, as [`CacheFile`] entries. -fn collect_source_files(root: &Path, content: &[ContentType]) -> Vec { - let resolved = get_extensions(content, None); - let ext_refs: Vec<&str> = resolved.iter().map(String::as_str).collect(); - let mut files = Vec::new(); - for file_path in walk_files(root, &ext_refs, &[]) { - let Ok(meta) = std::fs::metadata(&file_path) else { - continue; - }; - if meta.len() > MAX_FILE_BYTES { - continue; - } - let Ok(raw) = std::fs::read(&file_path) else { - continue; - }; - let rel = file_path.strip_prefix(root).unwrap_or(&file_path); - files.push(CacheFile { - path: rel.to_string_lossy().into_owned(), - content: raw, - }); - } - files -} - -/// Load a cached index for `source` if fresh, else build, persist, and return. -pub fn load_or_build_index(source: &str, options: &LoadOrBuildOptions) -> Result { - let content = normalize_content(options.content.clone()); - let is_git = is_git_url(source); - - let location = CacheLocation { - base_dir: options.base_dir.clone(), - git_ref: options.git_ref.clone(), - }; - let cache_dir = resolve_cache_dir(source, &content, &location); - let base_only = CacheLocation { - base_dir: options.base_dir.clone(), - git_ref: None, - }; - ensure_cache_dir(&cache_dir, &base_only)?; - - // Local sources: the source-file hash is the cache-validity oracle. Git - // sources are URL+ref keyed (no cheap live hash). - let source_hash = if is_git { - None - } else { - Some(compute_content_hash(&collect_source_files( - Path::new(source), - &content, - ))) - }; - - if let Some(cached) = try_reuse(&cache_dir, is_git, source_hash.as_deref()) { - return Ok(cached); - } - - let load_options = LoadOptions { - model_path: options.model_path.clone(), - content: Some(content), - }; - let index = if is_git { - CspIndex::from_git(source, &load_options, options.git_ref.as_deref())? - } else { - CspIndex::from_path(Path::new(source), &load_options)? - }; - index.save(&cache_dir, source_hash.as_deref())?; - Ok(index) -} - -/// Reuse a cached index when present and valid, else `None`. -fn try_reuse(cache_dir: &Path, is_git: bool, source_hash: Option<&str>) -> Option { - let manifest_path = cache_dir.join("manifest.json"); - if !manifest_path.exists() { - return None; - } - let raw = std::fs::read_to_string(&manifest_path).ok()?; - let value: serde_json::Value = serde_json::from_str(&raw).ok()?; - let manifest = parse_manifest(&value).ok()?; - // A chunk_size change re-chunks every file, so a cache built with a different - // target length is stale even if the source files are byte-identical. - if manifest.chunk_size != Some(DESIRED_CHUNK_LENGTH_CHARS as u32) { - return None; - } - // Local sources additionally validate the live source-file hash; git sources - // are URL+ref keyed (no cheap live hash). - if !is_git && Some(manifest.content_hash.as_str()) != source_hash { - return None; - } - CspIndex::load_from_disk(cache_dir).ok() -} +pub use crate::indexing::cache_orchestrator::{load_or_build_index, LoadOrBuildOptions}; #[cfg(test)] -mod tests { - use super::*; - use crate::indexing::dense::make_stub_model; - use tempfile::tempdir; - - fn make_chunk( - file_path: &str, - start: u32, - end: u32, - language: Option<&str>, - content: &str, - ) -> Chunk { - Chunk { - content: content.to_string(), - file_path: file_path.to_string(), - start_line: start, - end_line: end, - language: language.map(str::to_string), - } - } - - fn build_index(chunks: Vec) -> CspIndex { - let model = make_stub_model(4); - let vectors: Vec> = (0..chunks.len()) - .map(|i| { - let mut v = vec![0f32; 4]; - v[0] = (i + 1) as f32; - v - }) - .collect(); - CspIndex::new(CspIndexState { - model, - bm25_index: Bm25Index::build(&vec![vec!["x".to_string()]; chunks.len()]), - semantic_index: SelectableBasicBackend::from_vectors(vectors).unwrap(), - chunks, - model_path: "test-model".to_string(), - root: None, - content: DEFAULT_CONTENT.to_vec(), - }) - } - - #[test] - fn stats_zero_for_empty() { - let idx = build_index(vec![]); - let stats = idx.stats(); - assert_eq!(stats.indexed_files, 0); - assert_eq!(stats.total_chunks, 0); - assert!(stats.languages.is_empty()); - } - - #[test] - fn stats_reflect_distribution() { - let chunks = vec![ - make_chunk("a.ts", 1, 10, Some("typescript"), "x"), - make_chunk("a.ts", 11, 20, Some("typescript"), "y"), - make_chunk("b.py", 1, 5, Some("python"), "z"), - make_chunk("c.bin", 1, 1, None, "w"), - ]; - let stats = build_index(chunks).stats(); - assert_eq!(stats.indexed_files, 3); - assert_eq!(stats.total_chunks, 4); - assert_eq!(stats.languages.get("typescript"), Some(&2)); - assert_eq!(stats.languages.get("python"), Some(&1)); - assert_eq!(stats.languages.len(), 2); - } - - #[test] - fn search_empty_query_and_index() { - let idx = build_index(vec![make_chunk("a.ts", 1, 1, Some("typescript"), "x")]); - assert!(idx.search("", &QueryOptions::default()).is_empty()); - assert!(idx.search(" ", &QueryOptions::default()).is_empty()); - let empty = build_index(vec![]); - assert!(empty - .search("anything", &QueryOptions::default()) - .is_empty()); - } - - #[test] - fn search_top_k_zero() { - let idx = build_index(vec![make_chunk("a.ts", 1, 1, Some("typescript"), "x")]); - let opts = QueryOptions { - top_k: Some(0), - ..Default::default() - }; - assert!(idx.search("anything", &opts).is_empty()); - } - - #[test] - fn search_filters_matching_nothing() { - let chunks = vec![ - make_chunk("a.ts", 1, 10, Some("typescript"), "alpha"), - make_chunk("b.py", 1, 10, Some("python"), "beta"), - ]; - let idx = build_index(chunks); - let lang_opts = QueryOptions { - filter_languages: Some(vec!["nonexistent".to_string()]), - ..Default::default() - }; - assert!(idx.search("anything", &lang_opts).is_empty()); - let path_opts = QueryOptions { - filter_paths: Some(vec!["nope.ts".to_string()]), - ..Default::default() - }; - assert!(idx.search("anything", &path_opts).is_empty()); - } - - #[test] - fn find_related_excludes_seed() { - let chunks = vec![ - make_chunk("a.ts", 1, 10, Some("typescript"), "seed chunk"), - make_chunk("a.ts", 11, 20, Some("typescript"), "companion 1"), - make_chunk("b.ts", 1, 5, Some("typescript"), "companion 2"), - ]; - let idx = build_index(chunks.clone()); - let opts = QueryOptions { - top_k: Some(5), - ..Default::default() - }; - let results = idx.find_related(&chunks[0], &opts); - assert!(!results.iter().any(|r| r.chunk == chunks[0])); - assert!(results.len() <= 5); - } - - #[test] - fn save_load_roundtrip() { - let chunks = vec![ - make_chunk("a.ts", 1, 10, Some("typescript"), "A"), - make_chunk("b.ts", 1, 5, Some("python"), "B"), - ]; - let idx = build_index(chunks); - let dir = tempdir().unwrap(); - idx.save(dir.path(), None).unwrap(); - let loaded = CspIndex::load_from_disk(dir.path()).unwrap(); - assert_eq!(loaded.chunks.len(), 2); - let paths: Vec<&str> = loaded.chunks.iter().map(|c| c.file_path.as_str()).collect(); - assert_eq!(paths, ["a.ts", "b.ts"]); - let stats = loaded.stats(); - assert_eq!(stats.total_chunks, 2); - assert_eq!(stats.languages.get("typescript"), Some(&1)); - assert_eq!(stats.languages.get("python"), Some(&1)); - } - - #[test] - fn load_missing_directory() { - let dir = tempdir().unwrap(); - let err = CspIndex::load_from_disk(&dir.path().join("nope")).unwrap_err(); - assert!(err.contains("Index not found")); - } - - #[test] - fn load_missing_artifact() { - let dir = tempdir().unwrap(); - let err = CspIndex::load_from_disk(dir.path()).unwrap_err(); - assert!(err.contains("Missing:")); - } - - #[test] - fn load_schema_version_mismatch() { - let idx = build_index(vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]); - let dir = tempdir().unwrap(); - idx.save(dir.path(), None).unwrap(); - let manifest_path = dir.path().join("manifest.json"); - let raw = std::fs::read_to_string(&manifest_path).unwrap(); - let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); - value["schemaVersion"] = serde_json::json!(999); - std::fs::write(&manifest_path, value.to_string()).unwrap(); - let err = CspIndex::load_from_disk(dir.path()).unwrap_err(); - assert!(err.to_lowercase().contains("schema version")); - } - - #[test] - fn load_rejects_invalid_content() { - let idx = build_index(vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]); - let dir = tempdir().unwrap(); - idx.save(dir.path(), None).unwrap(); - let manifest_path = dir.path().join("manifest.json"); - let raw = std::fs::read_to_string(&manifest_path).unwrap(); - let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); - value["content"] = serde_json::json!(["bogus"]); - std::fs::write(&manifest_path, value.to_string()).unwrap(); - assert!(CspIndex::load_from_disk(dir.path()).is_err()); - } - - #[test] - fn save_writes_manifest_fields() { - let chunks = vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]; - let idx = build_index(chunks); - let dir = tempdir().unwrap(); - idx.save(dir.path(), None).unwrap(); - let raw = std::fs::read_to_string(dir.path().join("manifest.json")).unwrap(); - let value: serde_json::Value = serde_json::from_str(&raw).unwrap(); - assert_eq!(value["schemaVersion"], 1); - assert_eq!(value["modelId"], "test-model"); - assert_eq!(value["content"], serde_json::json!(["code"])); - assert!(value["contentHash"].as_str().unwrap().len() == 64); - assert_eq!( - value["chunkSize"].as_u64(), - Some(u64::from(DESIRED_CHUNK_LENGTH_CHARS as u32)) - ); - } - - #[test] - fn try_reuse_rejects_stale_chunk_size() { - let chunks = vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]; - let idx = build_index(chunks); - let dir = tempdir().unwrap(); - idx.save(dir.path(), Some("deadbeef")).unwrap(); - - // Fresh cache (matching hash + current chunk_size) is reused. - assert!(try_reuse(dir.path(), false, Some("deadbeef")).is_some()); - - // Rewrite the manifest with a different chunk_size → stale → rebuild. - let manifest_path = dir.path().join("manifest.json"); - let raw = std::fs::read_to_string(&manifest_path).unwrap(); - let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); - value["chunkSize"] = serde_json::json!(9999); - std::fs::write(&manifest_path, value.to_string()).unwrap(); - assert!(try_reuse(dir.path(), false, Some("deadbeef")).is_none()); - - // A manifest predating the field (absent chunkSize) is also stale. - let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); - value.as_object_mut().unwrap().remove("chunkSize"); - std::fs::write(&manifest_path, value.to_string()).unwrap(); - assert!(try_reuse(dir.path(), false, Some("deadbeef")).is_none()); - } - - #[test] - fn save_deterministic_content_hash() { - let chunks = vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]; - let dir_a = tempdir().unwrap(); - let dir_b = tempdir().unwrap(); - build_index(chunks.clone()) - .save(dir_a.path(), None) - .unwrap(); - build_index(chunks).save(dir_b.path(), None).unwrap(); - let ha: serde_json::Value = serde_json::from_str( - &std::fs::read_to_string(dir_a.path().join("manifest.json")).unwrap(), - ) - .unwrap(); - let hb: serde_json::Value = serde_json::from_str( - &std::fs::read_to_string(dir_b.path().join("manifest.json")).unwrap(), - ) - .unwrap(); - assert_eq!(ha["contentHash"], hb["contentHash"]); - } - - #[test] - fn from_path_errors_on_missing() { - let dir = tempdir().unwrap(); - let err = - CspIndex::from_path(&dir.path().join("nope"), &LoadOptions::default()).unwrap_err(); - assert!(err.contains("Path does not exist")); - } - - #[test] - fn from_path_errors_on_file() { - let dir = tempdir().unwrap(); - let file = dir.path().join("f.ts"); - std::fs::write(&file, "x").unwrap(); - let err = CspIndex::from_path(&file, &LoadOptions::default()).unwrap_err(); - assert!(err.contains("Path is not a directory")); - } - - #[test] - fn from_path_builds_index() { - let dir = tempdir().unwrap(); - std::fs::write(dir.path().join("sample.ts"), "export const x = 1\n").unwrap(); - let idx = CspIndex::from_path(dir.path(), &LoadOptions::default()).unwrap(); - assert!(!idx.chunks.is_empty()); - assert_eq!(idx.content, DEFAULT_CONTENT.to_vec()); - } - - // --- from_git --- - - #[test] - fn from_git_rejects_dash_ref() { - // No clone runs — the ref guard rejects a flag-injection ref first. - let err = CspIndex::from_git( - "file:///nonexistent", - &LoadOptions::default(), - Some("--upload-pack=evil"), - ) - .unwrap_err(); - assert!(err.contains("Invalid git ref")); - } - - #[test] - fn from_git_errors_on_bad_url() { - let dir = tempdir().unwrap(); - let bogus = dir.path().join("not-a-repo"); - let err = CspIndex::from_git( - &format!("file://{}", bogus.display()), - &LoadOptions::default(), - None, - ) - .unwrap_err(); - assert!(err.contains("git clone failed")); - } - - #[test] - fn from_git_clones_and_builds() { - let repo = tempdir().unwrap(); - let run = |args: &[&str]| { - Command::new("git") - .args(args) - .current_dir(repo.path()) - .env("GIT_TERMINAL_PROMPT", "0") - .output() - .expect("git available") - }; - if !run(&["init", "-q"]).status.success() { - return; // git unavailable — skip rather than fail. - } - run(&["config", "user.email", "test@example.com"]); - run(&["config", "user.name", "Test"]); - run(&["config", "commit.gpgsign", "false"]); - std::fs::write(repo.path().join("a.ts"), "export const x = 1\n").unwrap(); - run(&["add", "."]); - run(&["commit", "-q", "-m", "initial"]); - - let url = format!("file://{}", repo.path().display()); - let idx = CspIndex::from_git(&url, &LoadOptions::default(), None).unwrap(); - assert!(!idx.chunks.is_empty()); - assert_eq!(idx.root.as_deref(), Some(url.as_str())); - } - - // --- load_or_build_index (cache.ts loadOrBuildIndex parity) --- - - #[test] - fn load_or_build_miss_then_hit_then_invalidate() { - let home = tempdir().unwrap(); - let src = tempdir().unwrap(); - let base = home.path().join(".csp"); - std::fs::write( - src.path().join("a.ts"), - "export function alpha() { return 1 }\n", - ) - .unwrap(); - let src_str = src.path().to_string_lossy().into_owned(); - let opts = LoadOrBuildOptions { - base_dir: Some(base.clone()), - ..Default::default() - }; - - // Miss: builds and writes a manifest. - let first = load_or_build_index(&src_str, &opts).unwrap(); - assert!(!first.chunks.is_empty()); - let cache_dir = resolve_cache_dir( - &src_str, - DEFAULT_CONTENT, - &CacheLocation { - base_dir: Some(base.clone()), - git_ref: None, - }, - ); - assert!(cache_dir.join("manifest.json").exists()); - - // Hit: a second call reuses the cache (same chunk count). - let second = load_or_build_index(&src_str, &opts).unwrap(); - assert_eq!(second.chunks.len(), first.chunks.len()); - - // Invalidation: add a file → content hash changes → rebuild reflects it. - std::fs::write( - src.path().join("b.ts"), - "export function beta() { return 2 }\n", - ) - .unwrap(); - let third = load_or_build_index(&src_str, &opts).unwrap(); - assert!(third.chunks.iter().any(|c| c.file_path == "b.ts")); - assert!(third.chunks.len() >= first.chunks.len()); - } -} +mod tests; diff --git a/crates/csp/src/indexing/index/tests.rs b/crates/csp/src/indexing/index/tests.rs new file mode 100644 index 0000000..de9577a --- /dev/null +++ b/crates/csp/src/indexing/index/tests.rs @@ -0,0 +1,299 @@ +use super::*; +use crate::indexing::dense::make_stub_model; +use tempfile::tempdir; + +fn make_chunk( + file_path: &str, + start: u32, + end: u32, + language: Option<&str>, + content: &str, +) -> Chunk { + Chunk { + content: content.to_string(), + file_path: file_path.to_string(), + start_line: start, + end_line: end, + language: language.map(str::to_string), + } +} + +fn build_index(chunks: Vec) -> CspIndex { + let model = make_stub_model(4); + let vectors: Vec> = (0..chunks.len()) + .map(|i| { + let mut v = vec![0f32; 4]; + v[0] = (i + 1) as f32; + v + }) + .collect(); + CspIndex::new(CspIndexState { + model, + bm25_index: Bm25Index::build(&vec![vec!["x".to_string()]; chunks.len()]), + semantic_index: SelectableBasicBackend::from_vectors(vectors).unwrap(), + chunks, + model_path: "test-model".to_string(), + root: None, + content: DEFAULT_CONTENT.to_vec(), + }) +} + +#[test] +fn stats_zero_for_empty() { + let idx = build_index(vec![]); + let stats = idx.stats(); + assert_eq!(stats.indexed_files, 0); + assert_eq!(stats.total_chunks, 0); + assert!(stats.languages.is_empty()); +} + +#[test] +fn stats_reflect_distribution() { + let chunks = vec![ + make_chunk("a.ts", 1, 10, Some("typescript"), "x"), + make_chunk("a.ts", 11, 20, Some("typescript"), "y"), + make_chunk("b.py", 1, 5, Some("python"), "z"), + make_chunk("c.bin", 1, 1, None, "w"), + ]; + let stats = build_index(chunks).stats(); + assert_eq!(stats.indexed_files, 3); + assert_eq!(stats.total_chunks, 4); + assert_eq!(stats.languages.get("typescript"), Some(&2)); + assert_eq!(stats.languages.get("python"), Some(&1)); + assert_eq!(stats.languages.len(), 2); +} + +#[test] +fn search_empty_query_and_index() { + let idx = build_index(vec![make_chunk("a.ts", 1, 1, Some("typescript"), "x")]); + assert!(idx.search("", &QueryOptions::default()).is_empty()); + assert!(idx.search(" ", &QueryOptions::default()).is_empty()); + let empty = build_index(vec![]); + assert!(empty + .search("anything", &QueryOptions::default()) + .is_empty()); +} + +#[test] +fn search_top_k_zero() { + let idx = build_index(vec![make_chunk("a.ts", 1, 1, Some("typescript"), "x")]); + let opts = QueryOptions { + top_k: Some(0), + ..Default::default() + }; + assert!(idx.search("anything", &opts).is_empty()); +} + +#[test] +fn search_filters_matching_nothing() { + let chunks = vec![ + make_chunk("a.ts", 1, 10, Some("typescript"), "alpha"), + make_chunk("b.py", 1, 10, Some("python"), "beta"), + ]; + let idx = build_index(chunks); + let lang_opts = QueryOptions { + filter_languages: Some(vec!["nonexistent".to_string()]), + ..Default::default() + }; + assert!(idx.search("anything", &lang_opts).is_empty()); + let path_opts = QueryOptions { + filter_paths: Some(vec!["nope.ts".to_string()]), + ..Default::default() + }; + assert!(idx.search("anything", &path_opts).is_empty()); +} + +#[test] +fn find_related_excludes_seed() { + let chunks = vec![ + make_chunk("a.ts", 1, 10, Some("typescript"), "seed chunk"), + make_chunk("a.ts", 11, 20, Some("typescript"), "companion 1"), + make_chunk("b.ts", 1, 5, Some("typescript"), "companion 2"), + ]; + let idx = build_index(chunks.clone()); + let opts = QueryOptions { + top_k: Some(5), + ..Default::default() + }; + let results = idx.find_related(&chunks[0], &opts); + assert!(!results.iter().any(|r| r.chunk == chunks[0])); + assert!(results.len() <= 5); +} + +#[test] +fn save_load_roundtrip() { + let chunks = vec![ + make_chunk("a.ts", 1, 10, Some("typescript"), "A"), + make_chunk("b.ts", 1, 5, Some("python"), "B"), + ]; + let idx = build_index(chunks); + let dir = tempdir().unwrap(); + idx.save(dir.path(), None).unwrap(); + let loaded = CspIndex::load_from_disk(dir.path()).unwrap(); + assert_eq!(loaded.chunks.len(), 2); + let paths: Vec<&str> = loaded.chunks.iter().map(|c| c.file_path.as_str()).collect(); + assert_eq!(paths, ["a.ts", "b.ts"]); + let stats = loaded.stats(); + assert_eq!(stats.total_chunks, 2); + assert_eq!(stats.languages.get("typescript"), Some(&1)); + assert_eq!(stats.languages.get("python"), Some(&1)); +} + +#[test] +fn load_missing_directory() { + let dir = tempdir().unwrap(); + let err = CspIndex::load_from_disk(&dir.path().join("nope")).unwrap_err(); + assert!(err.contains("Index not found")); +} + +#[test] +fn load_missing_artifact() { + let dir = tempdir().unwrap(); + let err = CspIndex::load_from_disk(dir.path()).unwrap_err(); + assert!(err.contains("Missing:")); +} + +#[test] +fn load_schema_version_mismatch() { + let idx = build_index(vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]); + let dir = tempdir().unwrap(); + idx.save(dir.path(), None).unwrap(); + let manifest_path = dir.path().join("manifest.json"); + let raw = std::fs::read_to_string(&manifest_path).unwrap(); + let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); + value["schemaVersion"] = serde_json::json!(999); + std::fs::write(&manifest_path, value.to_string()).unwrap(); + let err = CspIndex::load_from_disk(dir.path()).unwrap_err(); + assert!(err.to_lowercase().contains("schema version")); +} + +#[test] +fn load_rejects_invalid_content() { + let idx = build_index(vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]); + let dir = tempdir().unwrap(); + idx.save(dir.path(), None).unwrap(); + let manifest_path = dir.path().join("manifest.json"); + let raw = std::fs::read_to_string(&manifest_path).unwrap(); + let mut value: serde_json::Value = serde_json::from_str(&raw).unwrap(); + value["content"] = serde_json::json!(["bogus"]); + std::fs::write(&manifest_path, value.to_string()).unwrap(); + assert!(CspIndex::load_from_disk(dir.path()).is_err()); +} + +#[test] +fn save_writes_manifest_fields() { + let chunks = vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]; + let idx = build_index(chunks); + let dir = tempdir().unwrap(); + idx.save(dir.path(), None).unwrap(); + let raw = std::fs::read_to_string(dir.path().join("manifest.json")).unwrap(); + let value: serde_json::Value = serde_json::from_str(&raw).unwrap(); + assert_eq!(value["schemaVersion"], 1); + assert_eq!(value["modelId"], "test-model"); + assert_eq!(value["content"], serde_json::json!(["code"])); + assert!(value["contentHash"].as_str().unwrap().len() == 64); + assert_eq!( + value["chunkSize"].as_u64(), + Some(u64::from(DESIRED_CHUNK_LENGTH_CHARS as u32)) + ); +} + +#[test] +fn save_deterministic_content_hash() { + let chunks = vec![make_chunk("a.ts", 1, 10, Some("typescript"), "A")]; + let dir_a = tempdir().unwrap(); + let dir_b = tempdir().unwrap(); + build_index(chunks.clone()) + .save(dir_a.path(), None) + .unwrap(); + build_index(chunks).save(dir_b.path(), None).unwrap(); + let ha: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(dir_a.path().join("manifest.json")).unwrap()) + .unwrap(); + let hb: serde_json::Value = + serde_json::from_str(&std::fs::read_to_string(dir_b.path().join("manifest.json")).unwrap()) + .unwrap(); + assert_eq!(ha["contentHash"], hb["contentHash"]); +} + +#[test] +fn from_path_errors_on_missing() { + let dir = tempdir().unwrap(); + let err = CspIndex::from_path(&dir.path().join("nope"), &LoadOptions::default()).unwrap_err(); + assert!(err.contains("Path does not exist")); +} + +#[test] +fn from_path_errors_on_file() { + let dir = tempdir().unwrap(); + let file = dir.path().join("f.ts"); + std::fs::write(&file, "x").unwrap(); + let err = CspIndex::from_path(&file, &LoadOptions::default()).unwrap_err(); + assert!(err.contains("Path is not a directory")); +} + +#[test] +fn from_path_builds_index() { + let dir = tempdir().unwrap(); + std::fs::write(dir.path().join("sample.ts"), "export const x = 1\n").unwrap(); + let idx = CspIndex::from_path(dir.path(), &LoadOptions::default()).unwrap(); + assert!(!idx.chunks.is_empty()); + assert_eq!(idx.content, DEFAULT_CONTENT.to_vec()); +} + +// --- from_git --- + +#[test] +fn from_git_rejects_dash_ref() { + // No clone runs — the ref guard rejects a flag-injection ref first. + let err = CspIndex::from_git( + "file:///nonexistent", + &LoadOptions::default(), + Some("--upload-pack=evil"), + ) + .unwrap_err(); + assert!(err.contains("Invalid git ref")); +} + +#[test] +fn from_git_errors_on_bad_url() { + let dir = tempdir().unwrap(); + let bogus = dir.path().join("not-a-repo"); + let err = CspIndex::from_git( + &format!("file://{}", bogus.display()), + &LoadOptions::default(), + None, + ) + .unwrap_err(); + assert!(err.contains("git clone failed")); +} + +#[test] +fn from_git_clones_and_builds() { + let repo = tempdir().unwrap(); + let run = |args: &[&str]| { + Command::new("git") + .args(args) + .current_dir(repo.path()) + .env("GIT_TERMINAL_PROMPT", "0") + .output() + .ok() + }; + let Some(output) = run(&["init", "-q"]) else { + return; // git unavailable — skip rather than fail. + }; + if !output.status.success() { + return; + } + run(&["config", "user.email", "test@example.com"]); + run(&["config", "user.name", "Test"]); + run(&["config", "commit.gpgsign", "false"]); + std::fs::write(repo.path().join("a.ts"), "export const x = 1\n").unwrap(); + run(&["add", "."]); + run(&["commit", "-q", "-m", "initial"]); + + let url = format!("file://{}", repo.path().display()); + let idx = CspIndex::from_git(&url, &LoadOptions::default(), None).unwrap(); + assert!(!idx.chunks.is_empty()); + assert_eq!(idx.root.as_deref(), Some(url.as_str())); +} diff --git a/crates/csp/src/indexing/mod.rs b/crates/csp/src/indexing/mod.rs index 868f6f8..771e6a4 100644 --- a/crates/csp/src/indexing/mod.rs +++ b/crates/csp/src/indexing/mod.rs @@ -5,6 +5,7 @@ //! Phase 3. pub mod cache; +mod cache_orchestrator; pub mod create; pub mod dense; pub mod file_walker; diff --git a/crates/csp/src/search.rs b/crates/csp/src/search.rs index 70a824b..f2b90e2 100644 --- a/crates/csp/src/search.rs +++ b/crates/csp/src/search.rs @@ -215,6 +215,9 @@ pub fn search( let b = normalized_bm25.get(&idx).copied().unwrap_or(0.0); combined.insert(idx, alpha_weight * s + (1.0 - alpha_weight) * b); } + // Drop chunks whose fused score is exactly 0.0 before ranking (parity with + // semble#219's `combined_scores = {... if score}`). + combined.retain(|_, &mut score| score != 0.0); let ranked: Vec<(usize, f64)> = if rerank { boost_multi_chunk_files(&mut combined, chunks); @@ -238,315 +241,4 @@ pub fn search( } #[cfg(test)] -mod tests { - use super::*; - use std::cell::RefCell; - - fn make_chunk(content: &str, file_path: &str, start_line: u32, end_line: u32) -> Chunk { - Chunk { - content: content.to_string(), - file_path: file_path.to_string(), - start_line, - end_line, - language: Some("ts".to_string()), - } - } - - fn make_chunks() -> Vec { - vec![ - make_chunk("class Alpha {}", "src/alpha.ts", 10, 20), - make_chunk("function beta() {}", "src/alpha.ts", 30, 40), - make_chunk("export const gamma = 1", "src/gamma.ts", 1, 5), - make_chunk("function delta() {}", "src/delta.ts", 5, 15), - make_chunk("class Epsilon {}", "src/epsilon.ts", 50, 60), - ] - } - - struct MockModel; - impl EmbeddingModel for MockModel { - fn encode(&self, texts: &[String]) -> Vec> { - texts.iter().map(|_| vec![0.1, 0.2, 0.3]).collect() - } - } - - #[derive(Default)] - struct QueryCall { - k: usize, - selector: Option>, - } - - struct MockSemantic { - results: Vec<(usize, f64)>, - calls: RefCell>, - } - impl MockSemantic { - fn new(results: Vec<(usize, f64)>) -> Self { - Self { - results, - calls: RefCell::new(Vec::new()), - } - } - } - impl VectorBackend for MockSemantic { - fn query( - &self, - _vectors: &[Vec], - k: usize, - selector: Option<&[u32]>, - ) -> Vec> { - self.calls.borrow_mut().push(QueryCall { - k, - selector: selector.map(<[u32]>::to_vec), - }); - vec![self.results.clone()] - } - } - - struct Bm25Call { - mask: Option>, - } - struct MockBm25 { - scores: Vec, - calls: RefCell>, - } - impl MockBm25 { - fn new(scores: Vec) -> Self { - Self { - scores, - calls: RefCell::new(Vec::new()), - } - } - } - impl SparseBackend for MockBm25 { - fn get_scores(&self, _tokens: &[String], weight_mask: Option<&[u8]>) -> Vec { - self.calls.borrow_mut().push(Bm25Call { - mask: weight_mask.map(<[u8]>::to_vec), - }); - self.scores.clone() - } - } - - fn opts(alpha: Option, rerank: Option) -> SearchOptions { - SearchOptions { - alpha, - selector: None, - rerank, - } - } - - // --- sort_top_k --- - - #[test] - fn sort_top_k_descending() { - let out = sort_top_k(&[0.1, 0.9, 0.5, 0.3, 0.7], 3); - assert_eq!(out, [1, 4, 2]); - } - - #[test] - fn sort_top_k_clamps() { - let out = sort_top_k(&[1.0, 2.0, 3.0], 10); - assert_eq!(out, [2, 1, 0]); - } - - #[test] - fn sort_top_k_empty() { - assert!(sort_top_k(&[], 5).is_empty()); - } - - // --- rrf_scores --- - - #[test] - fn rrf_assigns_by_rank() { - let mut raw = Scores::new(); - raw.insert(0, 0.1); - raw.insert(1, 0.9); - raw.insert(2, 0.5); - let rrf = rrf_scores(&raw); - assert!((rrf[&1] - 1.0 / (RRF_K as f64 + 1.0)).abs() < 1e-12); - assert!((rrf[&2] - 1.0 / (RRF_K as f64 + 2.0)).abs() < 1e-12); - assert!((rrf[&0] - 1.0 / (RRF_K as f64 + 3.0)).abs() < 1e-12); - } - - #[test] - fn rrf_empty() { - assert!(rrf_scores(&Scores::new()).is_empty()); - } - - #[test] - fn rrf_first_rank_is_one_over_61() { - let mut raw = Scores::new(); - raw.insert(0, 5.0); - let rrf = rrf_scores(&raw); - assert!((rrf[&0] - 1.0 / 61.0).abs() < 1e-12); - } - - // --- search_semantic / search_bm25 --- - - #[test] - fn semantic_distance_to_similarity() { - let chunks = make_chunks(); - let idx = MockSemantic::new(vec![(0, 0.2), (2, 0.7)]); - let results = search_semantic("q", &MockModel, &idx, &chunks, 5, None); - assert_eq!(results.len(), 2); - assert_eq!(results[0].0, 0); - assert!((results[0].1 - 0.8).abs() < 1e-10); - assert_eq!(results[1].0, 2); - assert!((results[1].1 - 0.3).abs() < 1e-10); - } - - #[test] - fn semantic_passes_selector_and_k() { - let chunks = make_chunks(); - let idx = MockSemantic::new(vec![(0, 0.5)]); - let selector = vec![0u32, 2]; - search_semantic("q", &MockModel, &idx, &chunks, 5, Some(&selector)); - let calls = idx.calls.borrow(); - assert_eq!(calls.len(), 1); - assert_eq!(calls[0].selector.as_deref(), Some([0u32, 2].as_slice())); - assert_eq!(calls[0].k, 5); - } - - #[test] - fn bm25_excludes_zero_and_sorts() { - let chunks = make_chunks(); - let bm = MockBm25::new(vec![0.5, 0.0, 0.9, 0.2, 0.0]); - let results = search_bm25("alpha beta", &bm, &chunks, 5, None); - let idxs: Vec = results.iter().map(|r| r.0).collect(); - assert_eq!(idxs, [2, 0, 3]); - assert!((results[0].1 - 0.9).abs() < 1e-5); - } - - #[test] - fn bm25_empty_tokens() { - let chunks = make_chunks(); - let bm = MockBm25::new(vec![1.0; 5]); - assert!(search_bm25(" ", &bm, &chunks, 5, None).is_empty()); - } - - #[test] - fn bm25_builds_mask_from_selector() { - let chunks = make_chunks(); - let bm = MockBm25::new(vec![1.0; 5]); - search_bm25("alpha", &bm, &chunks, 5, Some(&[1, 3])); - let calls = bm.calls.borrow(); - assert_eq!(calls.len(), 1); - assert_eq!(calls[0].mask.as_deref(), Some([0u8, 1, 0, 1, 0].as_slice())); - } - - // --- search --- - - #[test] - fn search_alpha_one_is_semantic() { - let chunks = make_chunks(); - let idx = MockSemantic::new(vec![(2, 0.05), (0, 0.10)]); - let bm = MockBm25::new(vec![0.0, 0.0, 0.0, 0.0, 9.0]); - let results = search( - "alpha", - &MockModel, - &idx, - &bm, - &chunks, - 3, - &opts(Some(1.0), Some(false)), - ); - assert_eq!(results[0].chunk, chunks[2]); - assert_eq!(results[1].chunk, chunks[0]); - assert!(results[0].score > 0.0); - assert!(results[1].score > 0.0); - if let Some(r) = results.iter().find(|r| r.chunk == chunks[4]) { - assert_eq!(r.score, 0.0); - } - } - - #[test] - fn search_alpha_zero_is_bm25() { - let chunks = make_chunks(); - let idx = MockSemantic::new(vec![(0, 0.05)]); - let bm = MockBm25::new(vec![0.5, 0.0, 0.9, 0.2, 0.0]); - let results = search( - "alpha", - &MockModel, - &idx, - &bm, - &chunks, - 3, - &opts(Some(0.0), Some(false)), - ); - let got: Vec<&Chunk> = results.iter().map(|r| &r.chunk).collect(); - assert_eq!(got, vec![&chunks[2], &chunks[0], &chunks[3]]); - } - - #[test] - fn search_rrf_first_rank_score() { - let chunks = make_chunks(); - let idx = MockSemantic::new(vec![(0, 0.0)]); - let bm = MockBm25::new(vec![0.0; 5]); - let results = search( - "q", - &MockModel, - &idx, - &bm, - &chunks, - 5, - &opts(Some(1.0), Some(false)), - ); - assert_eq!(results.len(), 1); - assert!((results[0].score - 1.0 / 61.0).abs() < 1e-10); - } - - #[test] - fn search_sorts_ties_by_start_line() { - let chunks = vec![ - make_chunk("foo", "src/late.ts", 100, 100), - make_chunk("bar", "src/early.ts", 1, 1), - ]; - let idx = MockSemantic::new(vec![(0, 0.5)]); - let bm = MockBm25::new(vec![0.0, 1.0]); - let results = search( - "q", - &MockModel, - &idx, - &bm, - &chunks, - 5, - &opts(Some(0.5), Some(false)), - ); - assert_eq!(results.len(), 2); - assert_eq!(results[0].chunk.start_line, 1); - assert_eq!(results[1].chunk.start_line, 100); - } - - #[test] - fn search_empty_inputs() { - let chunks = make_chunks(); - let idx = MockSemantic::new(vec![]); - let bm = MockBm25::new(vec![0.0; 5]); - let results = search( - "q", - &MockModel, - &idx, - &bm, - &chunks, - 5, - &SearchOptions::default(), - ); - assert!(results.is_empty()); - } - - #[test] - fn search_rerank_applies_multi_chunk_boost() { - let chunks = make_chunks(); - let idx = MockSemantic::new(vec![(0, 0.10), (1, 0.20), (2, 0.30)]); - let bm = MockBm25::new(vec![0.0; 5]); - let ranked = search( - "q", - &MockModel, - &idx, - &bm, - &chunks, - 3, - &opts(Some(1.0), Some(true)), - ); - assert_eq!(ranked[0].chunk.file_path, "src/alpha.ts"); - } -} +mod tests; diff --git a/crates/csp/src/search/tests.rs b/crates/csp/src/search/tests.rs new file mode 100644 index 0000000..661350f --- /dev/null +++ b/crates/csp/src/search/tests.rs @@ -0,0 +1,321 @@ +use super::*; +use std::cell::RefCell; + +fn make_chunk(content: &str, file_path: &str, start_line: u32, end_line: u32) -> Chunk { + Chunk { + content: content.to_string(), + file_path: file_path.to_string(), + start_line, + end_line, + language: Some("ts".to_string()), + } +} + +fn make_chunks() -> Vec { + vec![ + make_chunk("class Alpha {}", "src/alpha.ts", 10, 20), + make_chunk("function beta() {}", "src/alpha.ts", 30, 40), + make_chunk("export const gamma = 1", "src/gamma.ts", 1, 5), + make_chunk("function delta() {}", "src/delta.ts", 5, 15), + make_chunk("class Epsilon {}", "src/epsilon.ts", 50, 60), + ] +} + +struct MockModel; +impl EmbeddingModel for MockModel { + fn encode(&self, texts: &[String]) -> Vec> { + texts.iter().map(|_| vec![0.1, 0.2, 0.3]).collect() + } +} + +#[derive(Default)] +struct QueryCall { + k: usize, + selector: Option>, +} + +struct MockSemantic { + results: Vec<(usize, f64)>, + calls: RefCell>, +} +impl MockSemantic { + fn new(results: Vec<(usize, f64)>) -> Self { + Self { + results, + calls: RefCell::new(Vec::new()), + } + } +} +impl VectorBackend for MockSemantic { + fn query( + &self, + _vectors: &[Vec], + k: usize, + selector: Option<&[u32]>, + ) -> Vec> { + self.calls.borrow_mut().push(QueryCall { + k, + selector: selector.map(<[u32]>::to_vec), + }); + vec![self.results.clone()] + } +} + +struct Bm25Call { + mask: Option>, +} +struct MockBm25 { + scores: Vec, + calls: RefCell>, +} +impl MockBm25 { + fn new(scores: Vec) -> Self { + Self { + scores, + calls: RefCell::new(Vec::new()), + } + } +} +impl SparseBackend for MockBm25 { + fn get_scores(&self, _tokens: &[String], weight_mask: Option<&[u8]>) -> Vec { + self.calls.borrow_mut().push(Bm25Call { + mask: weight_mask.map(<[u8]>::to_vec), + }); + self.scores.clone() + } +} + +fn opts(alpha: Option, rerank: Option) -> SearchOptions { + SearchOptions { + alpha, + selector: None, + rerank, + } +} + +// --- sort_top_k --- + +#[test] +fn sort_top_k_descending() { + let out = sort_top_k(&[0.1, 0.9, 0.5, 0.3, 0.7], 3); + assert_eq!(out, [1, 4, 2]); +} + +#[test] +fn sort_top_k_clamps() { + let out = sort_top_k(&[1.0, 2.0, 3.0], 10); + assert_eq!(out, [2, 1, 0]); +} + +#[test] +fn sort_top_k_empty() { + assert!(sort_top_k(&[], 5).is_empty()); +} + +// --- rrf_scores --- + +#[test] +fn rrf_assigns_by_rank() { + let mut raw = Scores::new(); + raw.insert(0, 0.1); + raw.insert(1, 0.9); + raw.insert(2, 0.5); + let rrf = rrf_scores(&raw); + assert!((rrf[&1] - 1.0 / (RRF_K as f64 + 1.0)).abs() < 1e-12); + assert!((rrf[&2] - 1.0 / (RRF_K as f64 + 2.0)).abs() < 1e-12); + assert!((rrf[&0] - 1.0 / (RRF_K as f64 + 3.0)).abs() < 1e-12); +} + +#[test] +fn rrf_empty() { + assert!(rrf_scores(&Scores::new()).is_empty()); +} + +#[test] +fn rrf_first_rank_is_one_over_61() { + let mut raw = Scores::new(); + raw.insert(0, 5.0); + let rrf = rrf_scores(&raw); + assert!((rrf[&0] - 1.0 / 61.0).abs() < 1e-12); +} + +// --- search_semantic / search_bm25 --- + +#[test] +fn semantic_distance_to_similarity() { + let chunks = make_chunks(); + let idx = MockSemantic::new(vec![(0, 0.2), (2, 0.7)]); + let results = search_semantic("q", &MockModel, &idx, &chunks, 5, None); + assert_eq!(results.len(), 2); + assert_eq!(results[0].0, 0); + assert!((results[0].1 - 0.8).abs() < 1e-10); + assert_eq!(results[1].0, 2); + assert!((results[1].1 - 0.3).abs() < 1e-10); +} + +#[test] +fn semantic_passes_selector_and_k() { + let chunks = make_chunks(); + let idx = MockSemantic::new(vec![(0, 0.5)]); + let selector = vec![0u32, 2]; + search_semantic("q", &MockModel, &idx, &chunks, 5, Some(&selector)); + let calls = idx.calls.borrow(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].selector.as_deref(), Some([0u32, 2].as_slice())); + assert_eq!(calls[0].k, 5); +} + +#[test] +fn bm25_excludes_zero_and_sorts() { + let chunks = make_chunks(); + let bm = MockBm25::new(vec![0.5, 0.0, 0.9, 0.2, 0.0]); + let results = search_bm25("alpha beta", &bm, &chunks, 5, None); + let idxs: Vec = results.iter().map(|r| r.0).collect(); + assert_eq!(idxs, [2, 0, 3]); + assert!((results[0].1 - 0.9).abs() < 1e-5); +} + +#[test] +fn bm25_empty_tokens() { + let chunks = make_chunks(); + let bm = MockBm25::new(vec![1.0; 5]); + assert!(search_bm25(" ", &bm, &chunks, 5, None).is_empty()); +} + +#[test] +fn bm25_builds_mask_from_selector() { + let chunks = make_chunks(); + let bm = MockBm25::new(vec![1.0; 5]); + search_bm25("alpha", &bm, &chunks, 5, Some(&[1, 3])); + let calls = bm.calls.borrow(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].mask.as_deref(), Some([0u8, 1, 0, 1, 0].as_slice())); +} + +// --- search --- + +#[test] +fn search_alpha_one_is_semantic() { + let chunks = make_chunks(); + let idx = MockSemantic::new(vec![(2, 0.05), (0, 0.10)]); + let bm = MockBm25::new(vec![0.0, 0.0, 0.0, 0.0, 9.0]); + let results = search( + "alpha", + &MockModel, + &idx, + &bm, + &chunks, + 3, + &opts(Some(1.0), Some(false)), + ); + assert_eq!(results[0].chunk, chunks[2]); + assert_eq!(results[1].chunk, chunks[0]); + assert!(results[0].score > 0.0); + assert!(results[1].score > 0.0); + // The BM25-only chunk has a fused score of 0.0 at alpha=1.0 and is dropped + // (semble#219 filters `combined_scores` before ranking). + assert!(results.iter().all(|r| r.chunk != chunks[4])); +} + +#[test] +fn search_alpha_zero_is_bm25() { + let chunks = make_chunks(); + let idx = MockSemantic::new(vec![(0, 0.05)]); + let bm = MockBm25::new(vec![0.5, 0.0, 0.9, 0.2, 0.0]); + let results = search( + "alpha", + &MockModel, + &idx, + &bm, + &chunks, + 3, + &opts(Some(0.0), Some(false)), + ); + let got: Vec<&Chunk> = results.iter().map(|r| &r.chunk).collect(); + assert_eq!(got, vec![&chunks[2], &chunks[0], &chunks[3]]); +} + +#[test] +fn search_rrf_first_rank_score() { + let chunks = make_chunks(); + let idx = MockSemantic::new(vec![(0, 0.0)]); + let bm = MockBm25::new(vec![0.0; 5]); + let results = search( + "q", + &MockModel, + &idx, + &bm, + &chunks, + 5, + &opts(Some(1.0), Some(false)), + ); + assert_eq!(results.len(), 1); + assert!((results[0].score - 1.0 / 61.0).abs() < 1e-10); +} + +#[test] +fn search_sorts_ties_by_start_line() { + let chunks = vec![ + make_chunk("foo", "src/late.ts", 100, 100), + make_chunk("bar", "src/early.ts", 1, 1), + ]; + let idx = MockSemantic::new(vec![(0, 0.5)]); + let bm = MockBm25::new(vec![0.0, 1.0]); + let results = search( + "q", + &MockModel, + &idx, + &bm, + &chunks, + 5, + &opts(Some(0.5), Some(false)), + ); + assert_eq!(results.len(), 2); + assert_eq!(results[0].chunk.start_line, 1); + assert_eq!(results[1].chunk.start_line, 100); +} + +#[test] +fn search_empty_inputs() { + let chunks = make_chunks(); + let idx = MockSemantic::new(vec![]); + let bm = MockBm25::new(vec![0.0; 5]); + let results = search( + "q", + &MockModel, + &idx, + &bm, + &chunks, + 5, + &SearchOptions::default(), + ); + assert!(results.is_empty()); +} + +#[test] +fn search_rerank_applies_multi_chunk_boost() { + let chunks = make_chunks(); + let idx = MockSemantic::new(vec![(2, 0.05), (0, 0.10), (1, 0.20)]); + let bm = MockBm25::new(vec![0.0; 5]); + let unranked = search( + "q", + &MockModel, + &idx, + &bm, + &chunks, + 3, + &opts(Some(1.0), Some(false)), + ); + assert_eq!(unranked[0].chunk.file_path, "src/gamma.ts"); + + let ranked = search( + "q", + &MockModel, + &idx, + &bm, + &chunks, + 3, + &opts(Some(1.0), Some(true)), + ); + assert_eq!(ranked[0].chunk.file_path, "src/alpha.ts"); +}