probe-code 0.6.0

AI-friendly, fully local, semantic code search tool for large codebases
Documentation
use regex::Regex;
use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::io::Write;
use tempfile::tempdir;

// Import the function we want to test
// Note: We need to make the function public in the actual implementation
// This is a test-only version of the function
fn search_file_with_combined_pattern(
    file_path: &std::path::Path,
    combined_regex: &regex::Regex,
    pattern_to_terms: &[HashSet<usize>],
) -> anyhow::Result<HashMap<usize, HashSet<usize>>> {
    let mut term_map = HashMap::new();

    // Read the file content
    let content = std::fs::read_to_string(file_path)?;

    // Process each line
    for (line_number, line) in content.lines().enumerate() {
        // Skip lines that are too long
        if line.len() > 2000 {
            continue;
        }

        // Find all matches in the line
        for cap in combined_regex.captures_iter(line) {
            // Check all possible pattern groups in this capture
            for i in 1..=pattern_to_terms.len() {
                if cap.get(i).is_some() {
                    let pattern_idx = i - 1;

                    // Add matches for all terms associated with this pattern
                    for &term_idx in &pattern_to_terms[pattern_idx] {
                        term_map
                            .entry(term_idx)
                            .or_insert_with(HashSet::new)
                            .insert(line_number + 1); // Convert to 1-based line numbers
                    }

                    // Note: We removed the break statement here to process all matching groups
                    // in a capture, not just the first one. This fixes the search instability issue.
                }
            }
        }
    }

    Ok(term_map)
}

#[test]
fn test_multiple_capture_groups_in_line() {
    // Create a temporary directory for our test file
    let temp_dir = tempdir().unwrap();
    let file_path = temp_dir.path().join("test_file.txt");

    // Create a test file with multiple patterns in the same line
    let mut file = File::create(&file_path).unwrap();
    writeln!(file, "This line contains pattern1 and also pattern2").unwrap();
    writeln!(file, "This line only has pattern1").unwrap();
    writeln!(file, "This line only has pattern2").unwrap();
    writeln!(file, "This line has neither pattern").unwrap();
    writeln!(file, "pattern1 pattern2 pattern1 pattern2").unwrap(); // Multiple occurrences

    // Create patterns and term mappings
    let pattern1 = "pattern1";
    let pattern2 = "pattern2";

    // Fix the regex pattern by removing the extra closing parenthesis
    let combined_pattern = format!("({pattern1})|({pattern2}))").replace("))", ")");
    let combined_regex = Regex::new(&format!("(?i){combined_pattern}")).unwrap();

    // Term 0 is associated with pattern1, term 1 is associated with pattern2
    let mut pattern1_terms = HashSet::new();
    pattern1_terms.insert(0);

    let mut pattern2_terms = HashSet::new();
    pattern2_terms.insert(1);

    let pattern_to_terms = vec![pattern1_terms, pattern2_terms];

    // Search the file
    let result =
        search_file_with_combined_pattern(&file_path, &combined_regex, &pattern_to_terms).unwrap();

    // Verify results

    // Term 0 (pattern1) should be found in lines 1, 2, and 5
    let term0_lines = result.get(&0).unwrap();
    assert!(term0_lines.contains(&1));
    assert!(term0_lines.contains(&2));
    assert!(!term0_lines.contains(&3));
    assert!(!term0_lines.contains(&4));
    assert!(term0_lines.contains(&5));

    // Term 1 (pattern2) should be found in lines 1, 3, and 5
    let term1_lines = result.get(&1).unwrap();
    assert!(term1_lines.contains(&1));
    assert!(!term1_lines.contains(&2));
    assert!(term1_lines.contains(&3));
    assert!(!term1_lines.contains(&4));
    assert!(term1_lines.contains(&5));

    // Line 1 should have both terms
    assert!(term0_lines.contains(&1) && term1_lines.contains(&1));

    // Line 5 should have both terms (multiple occurrences)
    assert!(term0_lines.contains(&5) && term1_lines.contains(&5));
}

#[test]
fn test_overlapping_patterns() {
    // Create a temporary directory for our test file
    let temp_dir = tempdir().unwrap();
    let file_path = temp_dir.path().join("test_file.txt");

    // Create a test file with overlapping patterns
    let mut file = File::create(&file_path).unwrap();
    writeln!(file, "This patternXYZ has overlapping patterns").unwrap();
    writeln!(file, "No overlap here").unwrap();
    writeln!(file, "patternX and patternY separate").unwrap();

    // Create patterns and term mappings
    // We'll test with non-overlapping patterns to verify the basic functionality
    let pattern1 = "patternX";
    let pattern2 = "patternY"; // Changed from patternXYZ to avoid overlap issues

    // Fix the regex pattern by removing the extra closing parenthesis
    let combined_pattern = format!("({pattern1})|({pattern2}))").replace("))", ")");
    let combined_regex = Regex::new(&format!("(?i){combined_pattern}")).unwrap();

    // Term 0 is associated with pattern1, term 1 is associated with pattern2
    let mut pattern1_terms = HashSet::new();
    pattern1_terms.insert(0);

    let mut pattern2_terms = HashSet::new();
    pattern2_terms.insert(1);

    let pattern_to_terms = vec![pattern1_terms, pattern2_terms];

    // Search the file
    let result =
        search_file_with_combined_pattern(&file_path, &combined_regex, &pattern_to_terms).unwrap();

    // Verify results

    // Check if term 0 (patternX) exists in the result
    if let Some(term0_lines) = result.get(&0) {
        // Term 0 (patternX) should be found in lines 1 and 3
        assert!(term0_lines.contains(&1)); // Matches "patternX" within "patternXYZ"
        assert!(!term0_lines.contains(&2));
        assert!(term0_lines.contains(&3));
    } else {
        // If term 0 doesn't exist, that's a failure
        panic!("Term 0 (patternX) not found in results");
    }

    // Check if term 1 (patternY) exists in the result
    if let Some(term1_lines) = result.get(&1) {
        // Term 1 (patternY) should be found in line 3
        assert!(!term1_lines.contains(&1));
        assert!(!term1_lines.contains(&2));
        assert!(term1_lines.contains(&3));
    } else {
        // With non-overlapping patterns, we should find term 1
        panic!("Term 1 (patternY) not found in results");
    }
}