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;
#[derive(Debug, Serialize, Deserialize)]
pub struct TrigramIndex {
trigrams: HashMap<String, HashSet<u32>>,
total_trigrams: usize,
}
impl TrigramIndex {
pub fn new() -> Self {
Self {
trigrams: HashMap::new(),
total_trigrams: 0,
}
}
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;
}
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(())
}
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)
}
pub fn extract_regex_trigrams(&self, regex_pattern: &str) -> Result<Vec<String>> {
let _regex = Regex::new(regex_pattern)?;
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())
}
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() {
let all_handles: HashSet<u32> = self.trigrams
.values()
.flat_map(|handles| handles.iter())
.cloned()
.collect();
return Ok(all_handles);
}
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 {
return Ok(HashSet::new());
}
}
Ok(candidates.unwrap_or_default())
}
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,
}
fn extract_trigrams(text: &str) -> HashSet<String> {
let mut trigrams = HashSet::new();
let normalized = text.to_lowercase();
let chars: Vec<char> = normalized.chars().collect();
for window in chars.windows(3) {
let trigram: String = window.iter().collect();
if trigram.chars().all(|c| c.is_alphanumeric() || c == '_') {
trigrams.insert(trigram);
}
}
trigrams
}
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] {
'.' | '*' | '+' | '?' | '^' | '$' | '|' | '(' | ')' | '[' | ']' | '{' | '}' => {
if !current_literal.is_empty() {
literals.push(current_literal.clone());
current_literal.clear();
}
i += 1;
}
'\\' => {
if i + 1 < chars.len() {
current_literal.push(chars[i + 1]);
i += 2;
} else {
i += 1;
}
}
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"]), ("(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");
let candidates = index.trigrams.get("wor").unwrap();
assert!(candidates.contains(&1));
assert!(candidates.contains(&2));
let candidates = index.trigrams.get("hel").unwrap();
assert!(candidates.contains(&1));
assert!(!candidates.contains(&2));
}
}