use std::path::{Path, PathBuf};
use crate::chunking::source::chunk_source;
use crate::indexing::cache::sha256_hex;
use crate::indexing::dense::{embed_chunk_refs, Model, SelectableBasicBackend};
use crate::indexing::file_walker::walk_files;
use crate::indexing::files::{
detect_language, get_extensions, get_max_file_bytes, DEFAULT_MAX_FILE_BYTES, MAX_FILE_BYTES_ENV,
};
use crate::indexing::sparse::{enrich_for_bm25, Bm25Index};
use crate::indexing::types::{make_chunk_id, FileManifest, FileManifestEntry, PreviousIndex};
use crate::tokens::tokenize;
use crate::types::{Chunk, ContentType};
#[deprecated(
since = "0.1.10",
note = "use `indexing::files::DEFAULT_MAX_FILE_BYTES` or `get_max_file_bytes()` \
(the limit is now overridable via `CSP_MAX_FILE_BYTES`)"
)]
pub const MAX_FILE_BYTES: u64 = DEFAULT_MAX_FILE_BYTES;
pub struct CreateIndexOptions<'a> {
pub model: &'a Model,
pub extensions: Option<Vec<String>>,
pub content: Option<Vec<ContentType>>,
pub display_root: Option<PathBuf>,
pub max_file_bytes: Option<u64>,
}
impl<'a> CreateIndexOptions<'a> {
pub fn new(model: &'a Model) -> Self {
Self {
model,
extensions: None,
content: None,
display_root: None,
max_file_bytes: None,
}
}
}
#[derive(Debug)]
pub struct CreateIndexResult {
pub bm25_index: Bm25Index,
pub semantic_index: SelectableBasicBackend,
pub chunks: Vec<Chunk>,
pub files: FileManifest,
}
fn escape_control(path: &str) -> String {
let mut out = String::with_capacity(path.len());
for c in path.chars() {
if c.is_control() {
out.extend(c.escape_default());
} else {
out.push(c);
}
}
out
}
pub(crate) fn skipped_large_warning(
skipped: &[String],
max_file_bytes: u64,
from_env: bool,
) -> Option<String> {
if skipped.is_empty() {
return None;
}
let shown: Vec<String> = skipped.iter().take(5).map(|p| escape_control(p)).collect();
let knob = if from_env {
MAX_FILE_BYTES_ENV
} else {
"max_file_bytes"
};
Some(format!(
"Skipped {} file(s) exceeding the maximum file size of {} bytes \
(raise {} to include them): {}{}",
skipped.len(),
max_file_bytes,
knob,
shown.join(", "),
if skipped.len() > 5 { " ..." } else { "" },
))
}
fn reindex_file(
bm25_index: &mut Bm25Index,
indexed_path: &str,
file_chunks: &[Chunk],
previous_entry: Option<&FileManifestEntry>,
) -> Result<(), String> {
if let Some(entry) = previous_entry {
for slot in 0..entry.count {
bm25_index.remove_document(&make_chunk_id(indexed_path, slot));
}
}
for (slot, chunk) in file_chunks.iter().enumerate() {
bm25_index.add_document(
&make_chunk_id(indexed_path, slot),
&tokenize(&enrich_for_bm25(chunk)),
)?;
}
Ok(())
}
type PreviousParts = (
Bm25Index,
FileManifest,
Vec<Option<Chunk>>,
Vec<Option<Vec<f32>>>,
);
fn open_previous(previous: Option<PreviousIndex>) -> PreviousParts {
match previous {
Some(prev) => (
prev.bm25_index,
prev.files,
prev.chunks.into_iter().map(Some).collect(),
prev.vectors.into_iter().map(Some).collect(),
),
None => (
Bm25Index::new(),
FileManifest::new(),
Vec::new(),
Vec::new(),
),
}
}
fn display_path(file_path: &Path, display_root: Option<&Path>) -> String {
match display_root {
Some(root) => file_path
.strip_prefix(root)
.unwrap_or(file_path)
.to_string_lossy()
.into_owned(),
None => file_path.to_string_lossy().into_owned(),
}
}
fn take_previous_rows(
entry: Option<&FileManifestEntry>,
hash: &str,
previous_chunks: &mut [Option<Chunk>],
previous_vectors: &mut [Option<Vec<f32>>],
) -> Option<(Vec<Chunk>, Vec<Vec<f32>>)> {
let entry = entry?;
if entry.hash != hash
|| entry.end() > previous_chunks.len()
|| entry.end() > previous_vectors.len()
{
return None;
}
let rows: Option<Vec<Chunk>> = previous_chunks[entry.start..entry.end()]
.iter_mut()
.map(Option::take)
.collect();
let vecs: Option<Vec<Vec<f32>>> = previous_vectors[entry.start..entry.end()]
.iter_mut()
.map(Option::take)
.collect();
rows.zip(vecs)
}
fn embed_fresh_rows(
model: &Model,
chunks: &[Chunk],
fresh_rows: &[usize],
mut vectors: Vec<Option<Vec<f32>>>,
) -> Result<Vec<Vec<f32>>, String> {
let fresh_chunks: Vec<&Chunk> = fresh_rows.iter().map(|&i| &chunks[i]).collect();
let fresh_vectors =
SelectableBasicBackend::from_vectors(embed_chunk_refs(model, &fresh_chunks))?.vectors;
if fresh_vectors.len() != fresh_rows.len() {
return Err("Embedder returned the wrong number of rows".to_string());
}
for (&row, vector) in fresh_rows.iter().zip(fresh_vectors) {
vectors[row] = Some(vector);
}
let vectors: Option<Vec<Vec<f32>>> = vectors.into_iter().collect();
vectors.ok_or_else(|| "Internal error: an embedding row was left unfilled".to_string())
}
pub fn create_index_from_path(
path: &Path,
options: &CreateIndexOptions,
previous: Option<PreviousIndex>,
) -> Result<CreateIndexResult, String> {
let content = options
.content
.clone()
.unwrap_or_else(|| vec![ContentType::Code]);
let resolved = get_extensions(&content, options.extensions.as_deref());
let ext_refs: Vec<&str> = resolved.iter().map(String::as_str).collect();
let max_file_bytes = options.max_file_bytes.unwrap_or_else(get_max_file_bytes);
let (mut bm25_index, previous_files, mut previous_chunks, mut previous_vectors) =
open_previous(previous);
let mut chunks: Vec<Chunk> = Vec::new();
let mut chunk_ids: Vec<String> = Vec::new();
let mut vectors: Vec<Option<Vec<f32>>> = Vec::new();
let mut fresh_rows: Vec<usize> = Vec::new();
let mut files = FileManifest::new();
let mut skipped_large: Vec<String> = Vec::new();
for file_path in walk_files(path, &ext_refs, &[]) {
let language = detect_language(&file_path.to_string_lossy());
let size = match std::fs::metadata(&file_path) {
Ok(meta) => meta.len(),
Err(_) => continue,
};
if size > max_file_bytes {
skipped_large.push(display_path(&file_path, options.display_root.as_deref()));
continue;
}
let Ok(bytes) = std::fs::read(&file_path) else {
continue;
};
let hash = sha256_hex(&bytes);
let indexed_path = display_path(&file_path, options.display_root.as_deref());
if files.contains_key(&indexed_path) {
eprintln!(
"csp: skipping {}: its display path collides with an already indexed file \
(non-UTF-8 file name)",
escape_control(&file_path.display().to_string())
);
continue;
}
let previous_entry = previous_files.get(&indexed_path);
let reused = take_previous_rows(
previous_entry,
&hash,
&mut previous_chunks,
&mut previous_vectors,
);
let start = chunks.len();
let file_chunks = match reused {
Some((file_chunks, file_vectors)) => {
vectors.extend(file_vectors.into_iter().map(Some));
file_chunks
}
None => {
let source = String::from_utf8_lossy(&bytes).into_owned();
let file_chunks = chunk_source(&source, &indexed_path, language);
reindex_file(&mut bm25_index, &indexed_path, &file_chunks, previous_entry)?;
fresh_rows.extend(start..start + file_chunks.len());
vectors.extend(std::iter::repeat_n(None, file_chunks.len()));
file_chunks
}
};
let count = file_chunks.len();
chunk_ids.extend((0..count).map(|slot| make_chunk_id(&indexed_path, slot)));
chunks.extend(file_chunks);
files.insert(indexed_path, FileManifestEntry { hash, start, count });
}
for (indexed_path, entry) in &previous_files {
if !files.contains_key(indexed_path) {
reindex_file(&mut bm25_index, indexed_path, &[], Some(entry))?;
}
}
if let Some(warning) = skipped_large_warning(
&skipped_large,
max_file_bytes,
options.max_file_bytes.is_none(),
) {
eprintln!("csp: {warning}");
}
if chunks.is_empty() {
return Err(format!(
"No supported files found under {}.",
path.display()
));
}
let vectors = embed_fresh_rows(options.model, &chunks, &fresh_rows, vectors)?;
bm25_index.set_doc_order(chunk_ids);
let semantic_index = SelectableBasicBackend::from_normalized(vectors)?;
Ok(CreateIndexResult {
bm25_index,
semantic_index,
chunks,
files,
})
}
#[cfg(test)]
mod tests;