siftdb_core/

trigram.rs

use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::fs::File;
use std::io::{BufWriter, BufReader};
use serde::{Serialize, Deserialize};
use anyhow::Result;
use regex::Regex;

/// Trigram index for fast regex preprocessing and filtering
#[derive(Debug, Serialize, Deserialize)]
pub struct TrigramIndex {
    /// Maps trigram -> set of file handles that contain it
    trigrams: HashMap<String, HashSet<u32>>,
    /// Total number of trigrams indexed
    total_trigrams: usize,
}

impl TrigramIndex {
    /// Create new empty trigram index
    pub fn new() -> Self {
        Self {
            trigrams: HashMap::new(),
            total_trigrams: 0,
        }
    }

    /// Add content from a file to the trigram index
    pub fn add_file_content(&mut self, file_handle: u32, content: &str) {
        let trigrams = extract_trigrams(content);
        
        for trigram in trigrams {
            self.trigrams
                .entry(trigram)
                .or_insert_with(HashSet::new)
                .insert(file_handle);
        }
        
        self.total_trigrams += 1;
    }

    /// Save trigram index to file
    pub fn save_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let file = File::create(path)?;
        let writer = BufWriter::new(file);
        serde_json::to_writer_pretty(writer, self)?;
        Ok(())
    }

    /// Load trigram index from file
    pub fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file = File::open(path)?;
        let reader = BufReader::new(file);
        let index: Self = serde_json::from_reader(reader)?;
        Ok(index)
    }

    /// Extract trigrams from a regex pattern for fast filtering
    pub fn extract_regex_trigrams(&self, regex_pattern: &str) -> Result<Vec<String>> {
        // Try to compile the regex to validate it
        let _regex = Regex::new(regex_pattern)?;
        
        // For now, extract literal trigrams from the pattern
        // This is a simplified approach - a full implementation would
        // analyze the regex AST to find guaranteed literal sequences
        let literals = extract_literal_parts(regex_pattern);
        
        let mut trigrams = HashSet::new();
        for literal in literals {
            if literal.len() >= 3 {
                trigrams.extend(extract_trigrams(&literal));
            }
        }
        
        Ok(trigrams.into_iter().collect())
    }

    /// Get candidate file handles that might match a regex based on trigrams
    pub fn get_regex_candidates(&self, regex_pattern: &str) -> Result<HashSet<u32>> {
        let required_trigrams = self.extract_regex_trigrams(regex_pattern)?;
        
        if required_trigrams.is_empty() {
            // If no trigrams can be extracted, we need to check all files
            // Return all file handles from all trigrams
            let all_handles: HashSet<u32> = self.trigrams
                .values()
                .flat_map(|handles| handles.iter())
                .cloned()
                .collect();
            return Ok(all_handles);
        }

        // Find intersection of files containing all required trigrams
        let mut candidates: Option<HashSet<u32>> = None;
        
        for trigram in required_trigrams {
            if let Some(handles) = self.trigrams.get(&trigram) {
                match candidates {
                    None => candidates = Some(handles.clone()),
                    Some(ref mut current) => {
                        current.retain(|h| handles.contains(h));
                    }
                }
            } else {
                // If any required trigram is missing, no matches possible
                return Ok(HashSet::new());
            }
        }
        
        Ok(candidates.unwrap_or_default())
    }

    /// Get statistics about the trigram index
    pub fn stats(&self) -> TrigramStats {
        let total_file_references: usize = self.trigrams
            .values()
            .map(|handles| handles.len())
            .sum();

        TrigramStats {
            unique_trigrams: self.trigrams.len(),
            total_file_references,
            avg_files_per_trigram: if self.trigrams.is_empty() {
                0.0
            } else {
                total_file_references as f64 / self.trigrams.len() as f64
            },
        }
    }
}

#[derive(Debug)]
pub struct TrigramStats {
    pub unique_trigrams: usize,
    pub total_file_references: usize,
    pub avg_files_per_trigram: f64,
}

/// Extract all trigrams (3-character sequences) from text
fn extract_trigrams(text: &str) -> HashSet<String> {
    let mut trigrams = HashSet::new();
    
    // Normalize to lowercase for case-insensitive matching
    let normalized = text.to_lowercase();
    let chars: Vec<char> = normalized.chars().collect();
    
    // Extract 3-character windows
    for window in chars.windows(3) {
        let trigram: String = window.iter().collect();
        // Only include alphanumeric trigrams to reduce index size
        if trigram.chars().all(|c| c.is_alphanumeric() || c == '_') {
            trigrams.insert(trigram);
        }
    }
    
    trigrams
}

/// Extract literal parts from a regex pattern
/// This is a simplified implementation - a full version would parse the regex AST
fn extract_literal_parts(pattern: &str) -> Vec<String> {
    let mut literals = Vec::new();
    let mut current_literal = String::new();
    let chars: Vec<char> = pattern.chars().collect();
    
    let mut i = 0;
    while i < chars.len() {
        match chars[i] {
            // Regex metacharacters that end a literal sequence
            '.' | '*' | '+' | '?' | '^' | '$' | '|' | '(' | ')' | '[' | ']' | '{' | '}' => {
                if !current_literal.is_empty() {
                    literals.push(current_literal.clone());
                    current_literal.clear();
                }
                i += 1;
            }
            // Escape sequences
            '\\' => {
                if i + 1 < chars.len() {
                    // Add the escaped character as literal
                    current_literal.push(chars[i + 1]);
                    i += 2;
                } else {
                    i += 1;
                }
            }
            // Regular characters
            c => {
                current_literal.push(c);
                i += 1;
            }
        }
    }
    
    if !current_literal.is_empty() {
        literals.push(current_literal);
    }
    
    literals
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_extract_trigrams() {
        let text = "hello world";
        let trigrams = extract_trigrams(text);
        
        assert!(trigrams.contains("hel"));
        assert!(trigrams.contains("ell"));
        assert!(trigrams.contains("llo"));
        assert!(trigrams.contains("wor"));
        assert!(trigrams.contains("orl"));
        assert!(trigrams.contains("rld"));
    }

    #[test]
    fn test_extract_literal_parts() {
        let patterns = vec![
            ("hello", vec!["hello"]),
            ("hello.*world", vec!["hello", "world"]),
            ("fn\\s+\\w+", vec!["fn"]),  // \\s+ and \\w+ are not literals
            ("(test|demo)", vec!["test", "demo"]),
        ];
        
        for (pattern, expected) in patterns {
            let literals = extract_literal_parts(pattern);
            assert_eq!(literals, expected, "Failed for pattern: {}", pattern);
        }
    }

    #[test]
    fn test_trigram_index() {
        let mut index = TrigramIndex::new();
        
        index.add_file_content(1, "hello world");
        index.add_file_content(2, "world peace");
        
        // Both files should contain "wor" trigram
        let candidates = index.trigrams.get("wor").unwrap();
        assert!(candidates.contains(&1));
        assert!(candidates.contains(&2));
        
        // Only file 1 should contain "hel" trigram
        let candidates = index.trigrams.get("hel").unwrap();
        assert!(candidates.contains(&1));
        assert!(!candidates.contains(&2));
    }
}