use anyhow::Result;
use probe_code::search::file_list_cache;
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::time::{Duration, Instant};
use probe_code::models::{LimitedSearchResults, SearchResult};
use probe_code::path_resolver::resolve_path;
use probe_code::search::{
cache,
file_processing::{process_file_with_results, FileProcessingParams},
query::{create_query_plan, create_structured_patterns, QueryPlan},
result_ranking::rank_search_results,
search_limiter::apply_limits,
search_options::SearchOptions,
timeout,
};
pub struct SearchTimings {
pub query_preprocessing: Option<Duration>,
pub pattern_generation: Option<Duration>,
pub file_searching: Option<Duration>,
pub filename_matching: Option<Duration>,
pub early_filtering: Option<Duration>,
pub early_caching: Option<Duration>,
pub result_processing: Option<Duration>,
pub result_processing_file_io: Option<Duration>,
pub result_processing_line_collection: Option<Duration>,
pub result_processing_ast_parsing: Option<Duration>,
pub result_processing_block_extraction: Option<Duration>,
pub result_processing_result_building: Option<Duration>,
pub result_processing_ast_parsing_language_init: Option<Duration>,
pub result_processing_ast_parsing_parser_init: Option<Duration>,
pub result_processing_ast_parsing_tree_parsing: Option<Duration>,
pub result_processing_ast_parsing_line_map_building: Option<Duration>,
pub result_processing_block_extraction_code_structure: Option<Duration>,
pub result_processing_block_extraction_filtering: Option<Duration>,
pub result_processing_block_extraction_result_building: Option<Duration>,
pub result_processing_term_matching: Option<Duration>,
pub result_processing_compound_processing: Option<Duration>,
pub result_processing_line_matching: Option<Duration>,
pub result_processing_result_creation: Option<Duration>,
pub result_processing_synchronization: Option<Duration>,
pub result_processing_uncovered_lines: Option<Duration>,
pub result_ranking: Option<Duration>,
pub limit_application: Option<Duration>,
pub block_merging: Option<Duration>,
pub final_caching: Option<Duration>,
pub total_search_time: Option<Duration>,
}
pub fn format_duration(duration: Duration) -> String {
if duration.as_millis() < 1000 {
let millis = duration.as_millis();
format!("{millis}ms")
} else {
let secs = duration.as_secs_f64();
format!("{secs:.2}s")
}
}
pub fn print_timings(timings: &SearchTimings) {
let debug_mode = std::env::var("DEBUG").unwrap_or_default() == "1";
if !debug_mode {
return;
}
println!("\n=== SEARCH TIMING INFORMATION ===");
if let Some(duration) = timings.query_preprocessing {
println!("Query preprocessing: {}", format_duration(duration));
}
if let Some(duration) = timings.pattern_generation {
println!("Pattern generation: {}", format_duration(duration));
}
if let Some(duration) = timings.file_searching {
println!("File searching: {}", format_duration(duration));
}
if let Some(duration) = timings.filename_matching {
println!("Filename matching: {}", format_duration(duration));
}
if let Some(duration) = timings.early_filtering {
println!("Early AST filtering: {}", format_duration(duration));
}
if let Some(duration) = timings.early_caching {
println!("Early caching: {}", format_duration(duration));
}
if let Some(duration) = timings.result_processing {
println!("Result processing: {}", format_duration(duration));
if let Some(duration) = timings.result_processing_file_io {
println!(" - File I/O: {}", format_duration(duration));
}
if let Some(duration) = timings.result_processing_line_collection {
println!(" - Line collection: {}", format_duration(duration));
}
if let Some(duration) = timings.result_processing_ast_parsing {
println!(" - AST parsing: {}", format_duration(duration));
if let Some(d) = timings.result_processing_ast_parsing_language_init {
println!(" - Language init: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_ast_parsing_parser_init {
println!(" - Parser init: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_ast_parsing_tree_parsing {
println!(" - Tree parsing: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_ast_parsing_line_map_building {
println!(" - Line map building: {}", format_duration(d));
}
}
if let Some(duration) = timings.result_processing_block_extraction {
println!(" - Block extraction: {}", format_duration(duration));
if let Some(d) = timings.result_processing_block_extraction_code_structure {
println!(" - Code structure: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_block_extraction_filtering {
println!(" - Filtering: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_block_extraction_result_building {
println!(" - Result building: {}", format_duration(d));
}
}
if let Some(duration) = timings.result_processing_result_building {
println!(" - Result building: {}", format_duration(duration));
if let Some(d) = timings.result_processing_term_matching {
println!(" - Term matching: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_compound_processing {
println!(" - Compound processing: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_line_matching {
println!(" - Line matching: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_result_creation {
println!(" - Result creation: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_synchronization {
println!(" - Synchronization: {}", format_duration(d));
}
if let Some(d) = timings.result_processing_uncovered_lines {
println!(" - Uncovered lines: {}", format_duration(d));
}
}
}
if let Some(duration) = timings.result_ranking {
println!("Result ranking: {}", format_duration(duration));
}
if let Some(duration) = timings.limit_application {
println!("Limit application: {}", format_duration(duration));
}
if let Some(duration) = timings.block_merging {
println!("Block merging: {}", format_duration(duration));
}
if let Some(duration) = timings.final_caching {
println!("Final caching: {}", format_duration(duration));
}
if let Some(duration) = timings.total_search_time {
println!("Total search time: {}", format_duration(duration));
}
println!("===================================\n");
}
pub fn perform_probe(options: &SearchOptions) -> Result<LimitedSearchResults> {
let total_start = Instant::now();
let SearchOptions {
path,
queries,
files_only,
custom_ignores,
exclude_filenames,
reranker,
frequency_search: _,
exact,
language,
max_results,
max_bytes,
max_tokens,
allow_tests,
no_merge,
merge_threshold,
dry_run: _, session,
timeout,
} = options;
let timeout_handle = timeout::start_timeout_thread(*timeout);
let include_filenames = !exclude_filenames;
let debug_mode = std::env::var("DEBUG").unwrap_or_default() == "1";
let (effective_session, session_was_generated) = if let Some(s) = session {
if s.is_empty() || *s == "new" {
if let Ok(env_session_id) = std::env::var("PROBE_SESSION_ID") {
if !env_session_id.is_empty() {
if debug_mode {
println!("DEBUG: Using session ID from environment: {env_session_id}");
}
let static_id: &'static str = Box::leak(env_session_id.into_boxed_str());
(Some(static_id), false)
} else {
match cache::generate_session_id() {
Ok((new_id, _is_new)) => {
if debug_mode {
println!("DEBUG: Generated new session ID: {new_id}");
}
(Some(new_id), true)
}
Err(e) => {
eprintln!("Error generating session ID: {e}");
(None, false)
}
}
}
} else {
match cache::generate_session_id() {
Ok((new_id, _is_new)) => {
if debug_mode {
println!("DEBUG: Generated new session ID: {new_id}");
}
(Some(new_id), true)
}
Err(e) => {
eprintln!("Error generating session ID: {e}");
(None, false)
}
}
}
} else {
(Some(*s), false)
}
} else {
if let Ok(env_session_id) = std::env::var("PROBE_SESSION_ID") {
if !env_session_id.is_empty() {
if debug_mode {
println!("DEBUG: Using session ID from environment: {env_session_id}");
}
let static_id: &'static str = Box::leak(env_session_id.into_boxed_str());
(Some(static_id), false)
} else {
(None, false)
}
} else {
(None, false)
}
};
let mut timings = SearchTimings {
query_preprocessing: None,
pattern_generation: None,
file_searching: None,
filename_matching: None,
early_filtering: None,
early_caching: None,
result_processing: None,
result_processing_file_io: None,
result_processing_line_collection: None,
result_processing_ast_parsing: None,
result_processing_block_extraction: None,
result_processing_result_building: None,
result_processing_ast_parsing_language_init: None,
result_processing_ast_parsing_parser_init: None,
result_processing_ast_parsing_tree_parsing: None,
result_processing_ast_parsing_line_map_building: None,
result_processing_block_extraction_code_structure: None,
result_processing_block_extraction_filtering: None,
result_processing_block_extraction_result_building: None,
result_processing_term_matching: None,
result_processing_compound_processing: None,
result_processing_line_matching: None,
result_processing_result_creation: None,
result_processing_synchronization: None,
result_processing_uncovered_lines: None,
result_ranking: None,
limit_application: None,
block_merging: None,
final_caching: None,
total_search_time: None,
};
let qp_start = Instant::now();
if debug_mode {
println!("DEBUG: Starting query preprocessing...");
}
let parse_res = if queries.len() > 1 {
let combined_query = queries.join(" AND ");
create_query_plan(&combined_query, *exact)
} else {
create_query_plan(&queries[0], *exact)
};
let qp_duration = qp_start.elapsed();
timings.query_preprocessing = Some(qp_duration);
if debug_mode {
println!(
"DEBUG: Query preprocessing completed in {}",
format_duration(qp_duration)
);
}
if parse_res.is_err() {
println!("Failed to parse query as AST expression");
return Ok(LimitedSearchResults {
results: Vec::new(),
skipped_files: Vec::new(),
limits_applied: None,
cached_blocks_skipped: None,
});
}
let plan = parse_res.unwrap();
let pg_start = Instant::now();
if debug_mode {
println!("DEBUG: Starting pattern generation...");
println!("DEBUG: Using combined pattern approach for more efficient searching");
}
let structured_patterns = create_structured_patterns(&plan);
let pg_duration = pg_start.elapsed();
timings.pattern_generation = Some(pg_duration);
if debug_mode {
println!(
"DEBUG: Pattern generation completed in {}",
format_duration(pg_duration)
);
println!(
"DEBUG: Generated {patterns_len} patterns",
patterns_len = structured_patterns.len()
);
if structured_patterns.len() == 1 {
println!("DEBUG: Successfully created a single combined pattern for all terms");
}
}
let fs_start = Instant::now();
if debug_mode {
println!("DEBUG: Starting file searching...");
}
let lang_param = language.as_ref().map(|lang| normalize_language_alias(lang));
let mut file_term_map = search_with_structured_patterns(
path,
&plan,
&structured_patterns,
custom_ignores,
*allow_tests,
lang_param,
)?;
let fs_duration = fs_start.elapsed();
timings.file_searching = Some(fs_duration);
if debug_mode {
let total_matches: usize = file_term_map
.values()
.map(|term_map| term_map.values().map(|lines| lines.len()).sum::<usize>())
.sum();
let unique_files = file_term_map.keys().len();
println!(
"DEBUG: File searching completed in {} - Found {} matches in {} unique files",
format_duration(fs_duration),
total_matches,
unique_files
);
}
let mut all_files = file_term_map.keys().cloned().collect::<HashSet<_>>();
let fm_start = Instant::now();
if include_filenames && !exact {
if debug_mode {
println!("DEBUG: Starting filename matching...");
}
let resolved_path = if let Some(path_str) = path.to_str() {
match resolve_path(path_str) {
Ok(resolved_path) => {
if debug_mode {
println!(
"DEBUG: Resolved path '{}' to '{}'",
path_str,
resolved_path.display()
);
}
resolved_path
}
Err(err) => {
if debug_mode {
println!("DEBUG: Failed to resolve path '{path_str}': {err}");
}
path.to_path_buf()
}
}
} else {
path.to_path_buf()
};
let filename_matches: HashMap<PathBuf, HashSet<usize>> =
file_list_cache::find_matching_filenames(
&resolved_path,
queries,
&all_files,
custom_ignores,
*allow_tests,
&plan.term_indices,
lang_param,
)?;
if debug_mode {
println!(
"DEBUG: Found {} files matching by filename",
filename_matches.len()
);
}
for (pathbuf, matched_terms) in &filename_matches {
const MAX_FILE_SIZE: u64 = 1024 * 1024;
let resolved_path = match std::fs::canonicalize(pathbuf.as_path()) {
Ok(path) => path,
Err(e) => {
if debug_mode {
println!("DEBUG: Error resolving path for {pathbuf:?}: {e:?}");
}
continue;
}
};
let metadata = match std::fs::metadata(&resolved_path) {
Ok(meta) => meta,
Err(e) => {
if debug_mode {
println!("DEBUG: Error getting metadata for {resolved_path:?}: {e:?}");
}
continue;
}
};
if metadata.len() > MAX_FILE_SIZE {
if debug_mode {
println!(
"DEBUG: Skipping file {:?} - file too large ({} bytes > {} bytes limit)",
resolved_path,
metadata.len(),
MAX_FILE_SIZE
);
}
continue;
}
let file_content = match std::fs::read_to_string(&resolved_path) {
Ok(content) => content,
Err(e) => {
if debug_mode {
println!(
"DEBUG: Error reading file {:?}: {:?} (size: {} bytes)",
resolved_path,
e,
metadata.len()
);
}
continue;
}
};
let line_count = file_content.lines().count();
if line_count == 0 {
if debug_mode {
println!("DEBUG: File {pathbuf:?} is empty, skipping");
}
continue;
}
let all_line_numbers: HashSet<usize> = (1..=line_count).collect();
let mut term_map = if let Some(existing_map) = file_term_map.get(pathbuf) {
if debug_mode {
println!(
"DEBUG: File {pathbuf:?} already has term matches from content search, extending"
);
}
existing_map.clone()
} else {
if debug_mode {
println!("DEBUG: Creating new term map for file {pathbuf:?}");
}
HashMap::new()
};
for &term_idx in matched_terms {
term_map
.entry(term_idx)
.or_insert_with(HashSet::new)
.extend(&all_line_numbers);
if debug_mode {
println!(
"DEBUG: Added term index {term_idx} to file {pathbuf:?} with all lines"
);
}
}
file_term_map.insert(pathbuf.clone(), term_map);
all_files.insert(pathbuf.clone());
if debug_mode {
println!("DEBUG: Added file {pathbuf:?} with matching terms to file_term_map");
}
}
}
if debug_mode {
println!("DEBUG: all_files after filename matches: {all_files:?}");
}
let early_filter_start = Instant::now();
if debug_mode {
println!("DEBUG: Starting early AST filtering...");
println!("DEBUG: Before filtering: {} files", all_files.len());
}
let mut filtered_file_term_map = HashMap::new();
let mut filtered_all_files = HashSet::new();
for pathbuf in &all_files {
if let Some(term_map) = file_term_map.get(pathbuf) {
let matched_terms: HashSet<usize> = term_map.keys().copied().collect();
if plan.ast.evaluate(&matched_terms, &plan.term_indices, true) {
filtered_file_term_map.insert(pathbuf.clone(), term_map.clone());
filtered_all_files.insert(pathbuf.clone());
} else if debug_mode {
println!("DEBUG: Early filtering removed file: {pathbuf:?}");
}
} else if debug_mode {
println!("DEBUG: File {pathbuf:?} not found in file_term_map during early filtering");
}
}
file_term_map = filtered_file_term_map;
all_files = filtered_all_files;
if debug_mode {
println!(
"DEBUG: After early filtering: {} files remain",
all_files.len()
);
println!("DEBUG: all_files after early filtering: {all_files:?}");
}
let early_filter_duration = early_filter_start.elapsed();
timings.early_filtering = Some(early_filter_duration);
if debug_mode {
println!(
"DEBUG: Early AST filtering completed in {}",
format_duration(early_filter_duration)
);
}
let fm_duration = fm_start.elapsed();
timings.filename_matching = Some(fm_duration);
if debug_mode && include_filenames {
println!(
"DEBUG: Filename matching completed in {}",
format_duration(fm_duration)
);
}
if *files_only {
let mut res = Vec::new();
for f in all_files {
res.push(SearchResult {
file: f.to_string_lossy().to_string(),
lines: (1, 1),
node_type: "file".to_string(),
code: String::new(),
matched_by_filename: None,
rank: None,
score: None,
tfidf_score: None,
bm25_score: None,
tfidf_rank: None,
bm25_rank: None,
new_score: None,
hybrid2_rank: None,
combined_score_rank: None,
file_unique_terms: None,
file_total_matches: None,
file_match_rank: None,
block_unique_terms: None,
block_total_matches: None,
parent_file_id: None,
block_id: None,
matched_keywords: None,
tokenized_content: None,
});
}
let mut limited = apply_limits(res, *max_results, *max_bytes, *max_tokens);
limited.cached_blocks_skipped = None;
timings.total_search_time = Some(total_start.elapsed());
print_timings(&timings);
return Ok(limited);
}
let ec_start = Instant::now();
let mut early_skipped_count = 0;
if let Some(session_id) = effective_session {
let raw_query = if queries.len() > 1 {
queries.join(" AND ")
} else {
queries[0].clone()
};
if debug_mode {
println!(
"DEBUG: Starting early caching for session: {session_id} with query: {raw_query}"
);
if let Err(e) = cache::debug_print_cache(session_id, &raw_query) {
eprintln!("Error printing cache: {e}");
}
}
match cache::filter_matched_lines_with_cache(&mut file_term_map, session_id, &raw_query) {
Ok(skipped) => {
if debug_mode {
println!("DEBUG: Early caching skipped {skipped} matched lines");
}
early_skipped_count = skipped;
}
Err(e) => {
eprintln!("Error applying early cache: {e}");
}
}
let cached_files = file_term_map.keys().cloned().collect::<HashSet<_>>();
all_files = all_files.intersection(&cached_files).cloned().collect();
if debug_mode {
println!("DEBUG: all_files after caching: {all_files:?}");
}
}
let ec_duration = ec_start.elapsed();
timings.early_caching = Some(ec_duration);
if debug_mode && effective_session.is_some() {
println!(
"DEBUG: Early caching completed in {}",
format_duration(ec_duration)
);
}
let rp_start = Instant::now();
if debug_mode {
println!(
"DEBUG: Starting result processing for {} files after early caching...",
all_files.len()
);
}
let mut final_results = Vec::new();
let mut total_file_io_time = Duration::new(0, 0);
let mut total_line_collection_time = Duration::new(0, 0);
let mut total_ast_parsing_time = Duration::new(0, 0);
let mut total_block_extraction_time = Duration::new(0, 0);
let _total_result_building_time = Duration::new(0, 0);
let mut total_ast_parsing_language_init_time = Duration::new(0, 0);
let mut total_ast_parsing_parser_init_time = Duration::new(0, 0);
let mut total_ast_parsing_tree_parsing_time = Duration::new(0, 0);
let mut total_ast_parsing_line_map_building_time = Duration::new(0, 0);
let mut total_block_extraction_code_structure_time = Duration::new(0, 0);
let mut total_block_extraction_filtering_time = Duration::new(0, 0);
let mut total_block_extraction_result_building_time = Duration::new(0, 0);
let mut total_term_matching_time = Duration::new(0, 0);
let mut total_compound_processing_time = Duration::new(0, 0);
let mut total_line_matching_time = Duration::new(0, 0);
let mut total_result_creation_time = Duration::new(0, 0);
let mut total_synchronization_time = Duration::new(0, 0);
let mut total_uncovered_lines_time = Duration::new(0, 0);
for pathbuf in &all_files {
if debug_mode {
println!("DEBUG: Processing file: {pathbuf:?}");
}
if let Some(term_map) = file_term_map.get(pathbuf) {
if debug_mode {
println!("DEBUG: Term map for file: {term_map:?}");
}
let line_collection_start = Instant::now();
let mut all_lines = HashSet::new();
for lineset in term_map.values() {
all_lines.extend(lineset.iter());
}
let line_collection_duration = line_collection_start.elapsed();
total_line_collection_time += line_collection_duration;
if debug_mode {
println!(
"DEBUG: Found {} matched lines in file in {}",
all_lines.len(),
format_duration(line_collection_duration)
);
}
let filename_matched_queries = HashSet::new();
let term_pairs: Vec<(String, String)> = plan
.term_indices
.keys()
.map(|term| (term.clone(), term.clone()))
.collect();
let pparams = FileProcessingParams {
path: pathbuf,
line_numbers: &all_lines,
allow_tests: *allow_tests,
term_matches: term_map,
num_queries: plan.term_indices.len(),
filename_matched_queries,
queries_terms: &[term_pairs],
preprocessed_queries: None,
no_merge: *no_merge,
query_plan: &plan,
};
if debug_mode {
println!(
"DEBUG: Processing file with params: {}",
pparams.path.display()
);
}
match process_file_with_results(&pparams) {
Ok((mut file_res, file_timings)) => {
if let Some(duration) = file_timings.file_io {
total_file_io_time += duration;
}
if let Some(duration) = file_timings.ast_parsing {
total_ast_parsing_time += duration;
}
if let Some(duration) = file_timings.block_extraction {
total_block_extraction_time += duration;
}
if let Some(duration) = file_timings.ast_parsing_language_init {
total_ast_parsing_language_init_time += duration;
if debug_mode {
println!("DEBUG: - Language init: {}", format_duration(duration));
}
}
if let Some(duration) = file_timings.ast_parsing_parser_init {
total_ast_parsing_parser_init_time += duration;
if debug_mode {
println!("DEBUG: - Parser init: {}", format_duration(duration));
}
}
if let Some(duration) = file_timings.ast_parsing_tree_parsing {
total_ast_parsing_tree_parsing_time += duration;
if debug_mode {
println!("DEBUG: - Tree parsing: {}", format_duration(duration));
}
}
if let Some(duration) = file_timings.ast_parsing_line_map_building {
total_ast_parsing_line_map_building_time += duration;
if debug_mode {
println!(
"DEBUG: - Line map building: {}",
format_duration(duration)
);
}
}
if let Some(duration) = file_timings.block_extraction_code_structure {
total_block_extraction_code_structure_time += duration;
if debug_mode {
println!(
"DEBUG: - Code structure finding: {}",
format_duration(duration)
);
}
}
if let Some(duration) = file_timings.block_extraction_filtering {
total_block_extraction_filtering_time += duration;
if debug_mode {
println!("DEBUG: - Filtering: {}", format_duration(duration));
}
}
if let Some(duration) = file_timings.block_extraction_result_building {
total_block_extraction_result_building_time += duration;
if debug_mode {
println!(
"DEBUG: - Result building: {}",
format_duration(duration)
);
}
}
if let Some(duration) = file_timings.result_building_term_matching {
total_term_matching_time += duration;
if debug_mode {
println!("DEBUG: - Term matching: {}", format_duration(duration));
}
}
if let Some(duration) = file_timings.result_building_compound_processing {
total_compound_processing_time += duration;
if debug_mode {
println!(
"DEBUG: - Compound processing: {}",
format_duration(duration)
);
}
}
if let Some(duration) = file_timings.result_building_line_matching {
total_line_matching_time += duration;
if debug_mode {
println!("DEBUG: - Line matching: {}", format_duration(duration));
}
}
if let Some(duration) = file_timings.result_building_result_creation {
total_result_creation_time += duration;
if debug_mode {
println!(
"DEBUG: - Result creation: {}",
format_duration(duration)
);
}
}
if let Some(duration) = file_timings.result_building_synchronization {
total_synchronization_time += duration;
if debug_mode {
println!(
"DEBUG: - Synchronization: {}",
format_duration(duration)
);
}
}
if let Some(duration) = file_timings.result_building_uncovered_lines {
total_uncovered_lines_time += duration;
if debug_mode {
println!(
"DEBUG: - Uncovered lines: {}",
format_duration(duration)
);
}
}
if debug_mode {
println!("DEBUG: Got {} results from file processing", file_res.len());
if let Some(duration) = file_timings.file_io {
println!("DEBUG: File I/O time: {}", format_duration(duration));
}
if let Some(duration) = file_timings.ast_parsing {
println!("DEBUG: AST parsing time: {}", format_duration(duration));
}
if let Some(duration) = file_timings.block_extraction {
println!(
"DEBUG: Block extraction time: {}",
format_duration(duration)
);
}
if let Some(duration) = file_timings.block_extraction_result_building {
println!(
"DEBUG: Result building time: {}",
format_duration(duration)
);
}
}
final_results.append(&mut file_res);
}
Err(e) => {
if debug_mode {
println!("DEBUG: Error processing file: {e:?}");
}
}
}
} else {
if debug_mode {
println!("DEBUG: ERROR - File {pathbuf:?} not found in file_term_map but was in all_files");
}
}
}
let rp_duration = rp_start.elapsed();
let detailed_result_building_time = total_term_matching_time
+ total_compound_processing_time
+ total_line_matching_time
+ total_result_creation_time
+ total_synchronization_time
+ total_uncovered_lines_time;
let accounted_time = total_file_io_time
+ total_line_collection_time
+ total_ast_parsing_time
+ total_block_extraction_time;
let remaining_time = if rp_duration > accounted_time {
rp_duration - accounted_time
} else {
if detailed_result_building_time > Duration::new(0, 0) {
detailed_result_building_time
} else {
total_block_extraction_result_building_time
}
};
timings.result_processing = Some(rp_duration);
timings.result_processing_file_io = Some(total_file_io_time);
timings.result_processing_line_collection = Some(total_line_collection_time);
timings.result_processing_ast_parsing = Some(total_ast_parsing_time);
timings.result_processing_block_extraction = Some(total_block_extraction_time);
timings.result_processing_result_building = Some(remaining_time);
timings.result_processing_term_matching = Some(total_term_matching_time);
timings.result_processing_compound_processing = Some(total_compound_processing_time);
timings.result_processing_line_matching = Some(total_line_matching_time);
timings.result_processing_result_creation = Some(total_result_creation_time);
timings.result_processing_synchronization = Some(total_synchronization_time);
timings.result_processing_uncovered_lines = Some(total_uncovered_lines_time);
timings.result_processing_ast_parsing_language_init =
Some(total_ast_parsing_language_init_time);
timings.result_processing_ast_parsing_parser_init = Some(total_ast_parsing_parser_init_time);
timings.result_processing_ast_parsing_tree_parsing = Some(total_ast_parsing_tree_parsing_time);
timings.result_processing_ast_parsing_line_map_building =
Some(total_ast_parsing_line_map_building_time);
timings.result_processing_block_extraction_code_structure =
Some(total_block_extraction_code_structure_time);
timings.result_processing_block_extraction_filtering =
Some(total_block_extraction_filtering_time);
timings.result_processing_block_extraction_result_building =
Some(total_block_extraction_result_building_time);
if debug_mode {
println!(
"DEBUG: Result processing completed in {} - Generated {} results",
format_duration(rp_duration),
final_results.len()
);
println!("DEBUG: Granular result processing timings:");
println!("DEBUG: File I/O: {}", format_duration(total_file_io_time));
println!(
"DEBUG: Line collection: {}",
format_duration(total_line_collection_time)
);
println!(
"DEBUG: AST parsing: {}",
format_duration(total_ast_parsing_time)
);
println!(
"DEBUG: - Language init: {}",
format_duration(total_ast_parsing_language_init_time)
);
println!(
"DEBUG: - Parser init: {}",
format_duration(total_ast_parsing_parser_init_time)
);
println!(
"DEBUG: - Tree parsing: {}",
format_duration(total_ast_parsing_tree_parsing_time)
);
println!(
"DEBUG: - Line map building: {}",
format_duration(total_ast_parsing_line_map_building_time)
);
println!(
"DEBUG: Block extraction: {}",
format_duration(total_block_extraction_time)
);
println!(
"DEBUG: - Code structure finding: {}",
format_duration(total_block_extraction_code_structure_time)
);
println!(
"DEBUG: - Filtering: {}",
format_duration(total_block_extraction_filtering_time)
);
println!(
"DEBUG: - Result building: {}",
format_duration(total_block_extraction_result_building_time)
);
println!(
"DEBUG: Result building: {}",
format_duration(remaining_time)
);
}
let rr_start = Instant::now();
if debug_mode {
if *exact {
println!("DEBUG: Skipping result ranking due to exact flag being set");
} else {
println!("DEBUG: Starting result ranking...");
}
}
if !*exact {
rank_search_results(&mut final_results, queries, reranker);
}
let rr_duration = rr_start.elapsed();
timings.result_ranking = Some(rr_duration);
if debug_mode {
if *exact {
println!(
"DEBUG: Result ranking skipped in {}",
format_duration(rr_duration)
);
} else {
println!(
"DEBUG: Result ranking completed in {}",
format_duration(rr_duration)
);
}
}
let mut skipped_count = early_skipped_count;
let filtered_results = final_results;
let la_start = Instant::now();
if debug_mode {
println!("DEBUG: Starting limit application...");
}
let mut limited = apply_limits(filtered_results, *max_results, *max_bytes, *max_tokens);
let fc_start = Instant::now();
if let Some(session_id) = effective_session {
let raw_query = if queries.len() > 1 {
queries.join(" AND ")
} else {
queries[0].clone()
};
if debug_mode {
println!(
"DEBUG: Starting final caching for session: {session_id} with query: {raw_query}"
);
println!("DEBUG: Already skipped {early_skipped_count} lines in early caching");
if let Err(e) = cache::debug_print_cache(session_id, &raw_query) {
eprintln!("Error printing cache: {e}");
}
}
match cache::filter_results_with_cache(&limited.results, session_id, &raw_query) {
Ok((_, cached_skipped)) => {
if debug_mode {
println!("DEBUG: Final caching found {cached_skipped} cached blocks");
println!(
"DEBUG: Total skipped (early + final): {}",
early_skipped_count + cached_skipped
);
}
skipped_count += cached_skipped;
}
Err(e) => {
eprintln!("Error checking cache: {e}");
}
}
if let Err(e) = cache::add_results_to_cache(&limited.results, session_id, &raw_query) {
eprintln!("Error adding results to cache: {e}");
}
if debug_mode {
println!("DEBUG: Added limited results to cache before merging");
if let Err(e) = cache::debug_print_cache(session_id, &raw_query) {
eprintln!("Error printing updated cache: {e}");
}
}
}
limited.cached_blocks_skipped = if skipped_count > 0 {
Some(skipped_count)
} else {
None
};
let fc_duration = fc_start.elapsed();
timings.final_caching = Some(fc_duration);
if debug_mode && effective_session.is_some() {
println!(
"DEBUG: Final caching completed in {}",
format_duration(fc_duration)
);
}
let la_duration = la_start.elapsed();
timings.limit_application = Some(la_duration);
if debug_mode {
println!(
"DEBUG: Limit application completed in {} - Final result count: {}",
format_duration(la_duration),
limited.results.len()
);
}
let bm_start = Instant::now();
if debug_mode && !limited.results.is_empty() && !*no_merge {
println!("DEBUG: Starting block merging...");
}
let final_results = if !limited.results.is_empty() && !*no_merge {
use probe_code::search::block_merging::merge_ranked_blocks;
let merged = merge_ranked_blocks(limited.results.clone(), *merge_threshold);
let bm_duration = bm_start.elapsed();
timings.block_merging = Some(bm_duration);
if debug_mode {
println!(
"DEBUG: Block merging completed in {} - Merged result count: {}",
format_duration(bm_duration),
merged.len()
);
}
let merged_results = LimitedSearchResults {
results: merged.clone(),
skipped_files: limited.skipped_files,
limits_applied: limited.limits_applied,
cached_blocks_skipped: limited.cached_blocks_skipped,
};
if let Some(session_id) = effective_session {
let raw_query = if queries.len() > 1 {
queries.join(" AND ")
} else {
queries[0].clone()
};
if let Err(e) = cache::add_results_to_cache(&merged, session_id, &raw_query) {
eprintln!("Error adding merged results to cache: {e}");
}
if debug_mode {
println!("DEBUG: Added merged results to cache after merging");
if let Err(e) = cache::debug_print_cache(session_id, &raw_query) {
eprintln!("Error printing updated cache: {e}");
}
}
}
merged_results
} else {
let bm_duration = bm_start.elapsed();
timings.block_merging = Some(bm_duration);
if debug_mode && !*no_merge {
println!(
"DEBUG: Block merging skipped (no results or disabled) - {}",
format_duration(bm_duration)
);
}
limited
};
if let Some(session_id) = effective_session {
if session_was_generated {
println!("Session ID: {session_id} (generated - ALWAYS USE IT in future sessions for caching)");
} else {
println!("Session ID: {session_id}");
}
}
timings.total_search_time = Some(total_start.elapsed());
print_timings(&timings);
timeout_handle.store(true, std::sync::atomic::Ordering::SeqCst);
Ok(final_results)
}
pub fn search_with_structured_patterns(
root_path_str: &Path,
_plan: &QueryPlan,
patterns: &[(String, HashSet<usize>)],
custom_ignores: &[String],
allow_tests: bool,
language: Option<&str>,
) -> Result<HashMap<PathBuf, HashMap<usize, HashSet<usize>>>> {
let root_path = if let Some(path_str) = root_path_str.to_str() {
match resolve_path(path_str) {
Ok(resolved_path) => {
if std::env::var("DEBUG").unwrap_or_default() == "1" {
println!(
"DEBUG: Resolved path '{}' to '{}'",
path_str,
resolved_path.display()
);
}
resolved_path
}
Err(err) => {
if std::env::var("DEBUG").unwrap_or_default() == "1" {
println!("DEBUG: Failed to resolve path '{path_str}': {err}");
}
root_path_str.to_path_buf()
}
}
} else {
root_path_str.to_path_buf()
};
use rayon::prelude::*;
use regex::RegexSet;
use std::sync::{Arc, Mutex};
let debug_mode = std::env::var("DEBUG").unwrap_or_default() == "1";
let search_start = Instant::now();
if debug_mode {
println!("DEBUG: Starting parallel structured pattern search with RegexSet...");
println!("DEBUG: Creating RegexSet from {} patterns", patterns.len());
}
let pattern_strings: Vec<String> = patterns.iter().map(|(p, _)| format!("(?i){p}")).collect();
let regex_set = RegexSet::new(&pattern_strings)?;
let pattern_to_terms: Vec<HashSet<usize>> =
patterns.iter().map(|(_, terms)| terms.clone()).collect();
if debug_mode {
println!("DEBUG: RegexSet created successfully");
}
if debug_mode {
println!("DEBUG: Getting filtered file list from cache");
println!("DEBUG: Custom ignore patterns: {custom_ignores:?}");
}
let file_list = crate::search::file_list_cache::get_file_list_by_language(
&root_path,
allow_tests,
custom_ignores,
language,
)?;
if debug_mode {
println!("DEBUG: Got {} files from cache", file_list.files.len());
println!("DEBUG: Starting parallel file processing with RegexSet");
}
let regex_set = Arc::new(regex_set);
let pattern_to_terms = Arc::new(pattern_to_terms);
let file_term_maps = Arc::new(Mutex::new(HashMap::new()));
let individual_regexes: Vec<regex::Regex> = pattern_strings
.iter()
.map(|p| regex::Regex::new(p).unwrap())
.collect();
let individual_regexes = Arc::new(individual_regexes);
file_list.files.par_iter().for_each(|file_path| {
let regex_set = Arc::clone(®ex_set);
let pattern_to_terms = Arc::clone(&pattern_to_terms);
let individual_regexes = Arc::clone(&individual_regexes);
match search_file_with_regex_set(
file_path,
®ex_set,
&individual_regexes,
&pattern_to_terms,
) {
Ok(term_map) => {
if !term_map.is_empty() {
if debug_mode {
println!(
"DEBUG: File {:?} matched patterns with {} term indices",
file_path,
term_map.len()
);
}
let mut maps = file_term_maps.lock().unwrap();
maps.insert(file_path.clone(), term_map);
}
}
Err(e) => {
if debug_mode {
println!("DEBUG: Error searching file {file_path:?}: {e:?}");
}
}
}
});
let total_duration = search_start.elapsed();
let result = Arc::try_unwrap(file_term_maps)
.unwrap_or_else(|_| panic!("Failed to unwrap Arc"))
.into_inner()
.unwrap();
if debug_mode {
println!(
"DEBUG: Parallel search completed in {} - Found matches in {} files",
format_duration(total_duration),
result.len()
);
}
Ok(result)
}
fn search_file_with_regex_set(
file_path: &Path,
regex_set: ®ex::RegexSet,
individual_regexes: &[regex::Regex],
pattern_to_terms: &[HashSet<usize>],
) -> Result<HashMap<usize, HashSet<usize>>> {
let mut term_map = HashMap::new();
let debug_mode = std::env::var("DEBUG").unwrap_or_default() == "1";
const MAX_FILE_SIZE: u64 = 1024 * 1024;
let resolved_path = match std::fs::canonicalize(file_path) {
Ok(path) => path,
Err(e) => {
if debug_mode {
println!("DEBUG: Error resolving path for {file_path:?}: {e:?}");
}
return Err(anyhow::anyhow!("Failed to resolve file path: {}", e));
}
};
let metadata = match std::fs::metadata(&resolved_path) {
Ok(meta) => meta,
Err(e) => {
if debug_mode {
println!("DEBUG: Error getting metadata for {resolved_path:?}: {e:?}");
}
return Err(anyhow::anyhow!("Failed to get file metadata: {}", e));
}
};
if metadata.len() > MAX_FILE_SIZE {
if debug_mode {
println!(
"DEBUG: Skipping file {:?} - file too large ({} bytes > {} bytes limit)",
resolved_path,
metadata.len(),
MAX_FILE_SIZE
);
}
return Err(anyhow::anyhow!(
"File too large: {} bytes (limit: {} bytes)",
metadata.len(),
MAX_FILE_SIZE
));
}
let content = match std::fs::read_to_string(&resolved_path) {
Ok(content) => content,
Err(e) => {
if debug_mode {
println!(
"DEBUG: Error reading file {:?}: {:?} (size: {} bytes)",
resolved_path,
e,
metadata.len()
);
}
return Err(anyhow::anyhow!("Failed to read file: {}", e));
}
};
for (line_number, line) in content.lines().enumerate() {
if line.len() > 2000 {
if debug_mode {
println!(
"DEBUG: Skipping line {} in file {:?} - line too long ({} characters)",
line_number + 1,
file_path,
line.len()
);
}
continue;
}
let matches = regex_set.matches(line);
if matches.matched_any() {
for pattern_idx in matches.iter() {
if individual_regexes[pattern_idx].is_match(line) {
for &term_idx in &pattern_to_terms[pattern_idx] {
term_map
.entry(term_idx)
.or_insert_with(HashSet::new)
.insert(line_number + 1); }
}
}
}
}
Ok(term_map)
}
fn normalize_language_alias(lang: &str) -> &str {
match lang.to_lowercase().as_str() {
"rs" => "rust",
"js" | "jsx" => "javascript",
"ts" | "tsx" => "typescript",
"py" => "python",
"h" => "c",
"cc" | "cxx" | "hpp" | "hxx" => "cpp",
"rb" => "ruby",
"cs" => "csharp",
_ => lang, }
}