use std::collections::{HashMap, HashSet};
use std::fs;
use std::path::PathBuf;
use std::time::Instant;
use crate::db::queries::QueryBuilder;
use crate::extraction::grammars::{
detect_language, init_grammars, is_language_supported, language_key,
load_grammars_for_languages,
};
use crate::extraction::tree_sitter::extract_from_source;
use crate::types::{ExtractionError, ExtractionResult, Language};
use super::constants::{
FILE_IO_BATCH_SIZE, MAX_FILE_SIZE, PARSE_TIMEOUT_MS, WORKER_RECYCLE_INTERVAL,
};
use super::discovery::{get_git_changed_files, scan_directory};
use super::helpers::{extraction_error, hash_content, tracked_files_placeholder};
use super::results::{GitChanges, IndexPhase, IndexProgress, IndexResult, SyncResult};
pub struct ExtractionOrchestrator<'db> {
root_dir: PathBuf,
queries: QueryBuilder<'db>,
detected_framework_names: Option<Vec<String>>,
}
impl<'db> ExtractionOrchestrator<'db> {
pub fn new(root_dir: impl Into<PathBuf>, queries: QueryBuilder<'db>) -> Self {
Self {
root_dir: root_dir.into(),
queries,
detected_framework_names: None,
}
}
fn ensure_detected_frameworks(&mut self, files: Option<&[String]>) -> Vec<String> {
if let Some(names) = &self.detected_framework_names {
return names.clone();
}
let _files = files
.map(|files| files.to_vec())
.unwrap_or_else(|| scan_directory(&self.root_dir, None));
self.detected_framework_names = Some(Vec::new());
Vec::new()
}
pub async fn index_all<F>(&mut self, mut on_progress: Option<F>, verbose: bool) -> IndexResult
where
F: FnMut(IndexProgress),
{
let _ = init_grammars().await;
let start = Instant::now();
let mut errors = Vec::new();
let mut files_indexed = 0usize;
let mut files_skipped = 0usize;
let mut files_errored = 0usize;
let mut total_nodes = 0usize;
let mut total_edges = 0usize;
emit_progress(&mut on_progress, IndexPhase::Scanning, 0, 0, None);
let mut scan_cb = |current: usize, file: String| {
emit_progress(
&mut on_progress,
IndexPhase::Scanning,
current,
0,
Some(file),
);
};
let files = scan_directory(&self.root_dir, Some(&mut scan_cb));
self.detected_framework_names = None;
let framework_names = self.ensure_detected_frameworks(Some(&files));
emit_progress(&mut on_progress, IndexPhase::Parsing, 0, files.len(), None);
let mut languages = files
.iter()
.map(|file| detect_language(file, None))
.collect::<Vec<_>>();
if languages
.iter()
.any(|language| language_key(language) == "c")
{
languages.push(Language::Cpp);
}
let _ = load_grammars_for_languages(&languages).await;
for batch in files.chunks(FILE_IO_BATCH_SIZE) {
for file_path in batch {
match self.index_file_inner(file_path, &framework_names).await {
Ok(result) => {
if result.errors.is_empty() {
files_indexed += 1;
} else {
files_errored += 1;
errors.extend(result.errors.clone());
}
total_nodes += result.nodes.len();
total_edges += result.edges.len();
}
Err(error) => {
files_skipped += 1;
errors.push(error);
}
}
emit_progress(
&mut on_progress,
IndexPhase::Parsing,
files_indexed + files_skipped + files_errored,
files.len(),
Some(file_path.clone()),
);
}
}
emit_progress(
&mut on_progress,
IndexPhase::Resolving,
files.len(),
files.len(),
None,
);
if verbose {
let _ = (PARSE_TIMEOUT_MS, WORKER_RECYCLE_INTERVAL);
}
IndexResult {
success: errors.is_empty(),
files_indexed,
files_skipped,
files_errored,
nodes_created: total_nodes,
edges_created: total_edges,
errors,
duration_ms: start.elapsed().as_millis() as u64,
}
}
pub async fn index_files(&mut self, file_paths: &[String]) -> IndexResult {
let start = Instant::now();
let framework_names = self.ensure_detected_frameworks(Some(file_paths));
let mut out = IndexResult {
success: true,
files_indexed: 0,
files_skipped: 0,
files_errored: 0,
nodes_created: 0,
edges_created: 0,
errors: Vec::new(),
duration_ms: 0,
};
for file_path in file_paths {
match self.index_file_inner(file_path, &framework_names).await {
Ok(result) => {
out.files_indexed += 1;
out.nodes_created += result.nodes.len();
out.edges_created += result.edges.len();
out.errors.extend(result.errors);
}
Err(error) => {
out.files_errored += 1;
out.errors.push(error);
}
}
}
out.success = out.errors.is_empty();
out.duration_ms = start.elapsed().as_millis() as u64;
out
}
pub async fn index_file(&mut self, relative_path: &str) -> ExtractionResult {
let framework_names = self.ensure_detected_frameworks(None);
self.index_file_inner(relative_path, &framework_names)
.await
.unwrap_or_else(|error| ExtractionResult {
nodes: Vec::new(),
edges: Vec::new(),
unresolved_references: Vec::new(),
errors: vec![error],
duration_ms: 0,
})
}
pub async fn index_file_with_content(
&mut self,
relative_path: &str,
content: &str,
) -> ExtractionResult {
let framework_names = self.ensure_detected_frameworks(None);
let language = detect_language(relative_path, Some(content));
if !is_language_supported(language) {
return ExtractionResult {
nodes: Vec::new(),
edges: Vec::new(),
unresolved_references: Vec::new(),
errors: vec![extraction_error(
format!("Unsupported language: {}", language_key(&language)),
Some(relative_path.to_owned()),
"unsupported_language",
)],
duration_ms: 0,
};
}
let result = extract_from_source(
relative_path,
content,
Some(language),
Some(&framework_names),
);
self.store_extraction_result(relative_path, content, &result);
result
}
async fn index_file_inner(
&mut self,
relative_path: &str,
framework_names: &[String],
) -> Result<ExtractionResult, ExtractionError> {
let full_path = self.root_dir.join(relative_path);
let metadata = fs::metadata(&full_path).map_err(|err| {
extraction_error(
format!("Unable to stat file: {err}"),
Some(relative_path.to_owned()),
"file_error",
)
})?;
if metadata.len() > MAX_FILE_SIZE {
return Err(extraction_error(
"File exceeds maximum indexable size".to_owned(),
Some(relative_path.to_owned()),
"file_too_large",
));
}
let content = fs::read_to_string(&full_path).map_err(|err| {
extraction_error(
format!("Unable to read file: {err}"),
Some(relative_path.to_owned()),
"file_error",
)
})?;
let language = detect_language(relative_path, Some(&content));
let result = extract_from_source(
relative_path,
&content,
Some(language),
Some(framework_names),
);
self.store_extraction_result(relative_path, &content, &result);
Ok(result)
}
fn store_extraction_result(
&mut self,
relative_path: &str,
content: &str,
result: &ExtractionResult,
) {
let _content_hash = hash_content(content);
let _ = (relative_path, result, &mut self.queries);
}
pub async fn sync<F>(&mut self, mut on_progress: Option<F>) -> SyncResult
where
F: FnMut(IndexProgress),
{
let start = Instant::now();
let changes = self.get_changed_files();
let mut changed_file_paths = Vec::new();
changed_file_paths.extend(changes.added.iter().cloned());
changed_file_paths.extend(changes.modified.iter().cloned());
changed_file_paths.extend(changes.deleted.iter().cloned());
emit_progress(
&mut on_progress,
IndexPhase::Parsing,
0,
changes.added.len() + changes.modified.len(),
None,
);
let mut nodes_updated = 0usize;
for file in changes.added.iter().chain(changes.modified.iter()) {
let result = self.index_file(file).await;
nodes_updated += result.nodes.len();
}
for removed in &changes.deleted {
let _ = (removed, &mut self.queries);
}
SyncResult {
files_checked: changed_file_paths.len(),
files_added: changes.added.len(),
files_modified: changes.modified.len(),
files_removed: changes.deleted.len(),
nodes_updated,
duration_ms: start.elapsed().as_millis() as u64,
changed_file_paths: Some(changed_file_paths),
}
}
pub fn get_changed_files(&self) -> GitChanges {
if let Some(git_changes) = get_git_changed_files(&self.root_dir) {
return git_changes;
}
let current_files = scan_directory(&self.root_dir, None)
.into_iter()
.collect::<HashSet<_>>();
let tracked_files = tracked_files_placeholder(&self.queries);
let tracked_map = tracked_files
.iter()
.map(|file| (file.path.clone(), file.clone()))
.collect::<HashMap<_, _>>();
let mut changes = GitChanges::default();
for tracked in &tracked_files {
if !current_files.contains(&tracked.path) {
changes.deleted.push(tracked.path.clone());
}
}
for file_path in current_files {
let full_path = self.root_dir.join(&file_path);
let Ok(content) = fs::read_to_string(full_path) else {
continue;
};
let content_hash = hash_content(&content);
match tracked_map.get(&file_path) {
None => changes.added.push(file_path),
Some(tracked) if tracked.content_hash != content_hash => {
changes.modified.push(file_path)
}
_ => {}
}
}
changes
}
}
fn emit_progress<F>(
callback: &mut Option<F>,
phase: IndexPhase,
current: usize,
total: usize,
current_file: Option<String>,
) where
F: FnMut(IndexProgress),
{
if let Some(callback) = callback.as_mut() {
callback(IndexProgress {
phase,
current,
total,
current_file,
});
}
}