siftdb_core/

inverted_index.rs

use std::path::{Path, PathBuf};
use std::collections::{HashMap, HashSet};
use std::fs::File;
use std::io::{BufWriter, BufReader};
use fst::{IntoStreamer, Streamer, Map, MapBuilder};
use anyhow::{Result, Context};
use memmap2::Mmap;
use serde::{Serialize, Deserialize};

/// Inverted index mapping terms to file handles for O(1) search
#[derive(Debug)]
pub struct InvertedIndex {
    /// FST mapping term -> encoded file handle list
    term_map: Option<Map<Mmap>>,
    /// Decoded posting lists: term -> set of file handles
    posting_lists: HashMap<String, HashSet<u32>>,
}

#[derive(Serialize, Deserialize)]
struct PostingListData {
    posting_lists: HashMap<String, HashSet<u32>>,
}

impl InvertedIndex {
    /// Create a new empty inverted index
    pub fn new() -> Self {
        Self {
            term_map: None,
            posting_lists: HashMap::new(),
        }
    }

    /// Build inverted index from file contents
    pub fn build_from_content<P: AsRef<Path>>(
        file_contents: HashMap<u32, String>, // file_handle -> content
        output_fst_path: P,
        output_json_path: P,
    ) -> Result<Self> {
        let mut posting_lists: HashMap<String, HashSet<u32>> = HashMap::new();

        // Tokenize and build posting lists
        for (file_handle, content) in file_contents {
            let tokens = tokenize(&content);
            for token in tokens {
                posting_lists
                    .entry(token.to_lowercase())
                    .or_insert_with(HashSet::new)
                    .insert(file_handle);
            }
        }

        // Build FST from terms
        let mut sorted_terms: Vec<_> = posting_lists.keys().collect();
        sorted_terms.sort();

        let file = File::create(&output_fst_path)?;
        let mut builder = MapBuilder::new(BufWriter::new(file))?;

        for (i, term) in sorted_terms.iter().enumerate() {
            let term_bytes = term.as_bytes();
            builder.insert(term_bytes, i as u64)?;
        }

        builder.finish()?;

        // Save posting lists to JSON
        let data = PostingListData {
            posting_lists: posting_lists.clone(),
        };
        let file = File::create(&output_json_path)?;
        let writer = BufWriter::new(file);
        serde_json::to_writer(writer, &data)?;

        // Load the built FST
        let file = File::open(&output_fst_path)?;
        let mmap = unsafe { memmap2::Mmap::map(&file)? };
        let term_map = Map::new(mmap)?;

        Ok(Self {
            term_map: Some(term_map),
            posting_lists,
        })
    }

    /// Load existing inverted index from files
    pub fn load_from_files<P: AsRef<Path>>(
        fst_path: P,
        json_path: P,
    ) -> Result<Self> {
        // Load FST
        let file = File::open(&fst_path)?;
        let mmap = unsafe { memmap2::Mmap::map(&file)? };
        let term_map = Map::new(mmap)?;

        // Load posting lists from JSON
        let file = File::open(&json_path)?;
        let reader = BufReader::new(file);
        let data: PostingListData = serde_json::from_reader(reader)
            .context("Failed to parse posting lists")?;

        Ok(Self {
            term_map: Some(term_map),
            posting_lists: data.posting_lists,
        })
    }

    /// Find all file handles containing a term (O(1) lookup!)
    pub fn find_files_with_term(&self, term: &str) -> HashSet<u32> {
        let term_lower = term.to_lowercase();
        self.posting_lists
            .get(&term_lower)
            .cloned()
            .unwrap_or_default()
    }

    /// Find files containing all terms (AND query)
    pub fn find_files_with_all_terms(&self, terms: &[&str]) -> HashSet<u32> {
        if terms.is_empty() {
            return HashSet::new();
        }

        let mut result = self.find_files_with_term(terms[0]);
        
        for &term in &terms[1..] {
            let term_files = self.find_files_with_term(term);
            result = result.intersection(&term_files).cloned().collect();
            
            if result.is_empty() {
                break; // Early termination
            }
        }

        result
    }

    /// Find files containing any terms (OR query) 
    pub fn find_files_with_any_terms(&self, terms: &[&str]) -> HashSet<u32> {
        let mut result = HashSet::new();
        
        for &term in terms {
            let term_files = self.find_files_with_term(term);
            result = result.union(&term_files).cloned().collect();
        }

        result
    }

    /// Check if a term exists in the index
    pub fn contains_term(&self, term: &str) -> bool {
        let term_lower = term.to_lowercase();
        self.posting_lists.contains_key(&term_lower)
    }

    /// Get number of files containing a term
    pub fn term_frequency(&self, term: &str) -> usize {
        let term_lower = term.to_lowercase();
        self.posting_lists
            .get(&term_lower)
            .map(|files| files.len())
            .unwrap_or(0)
    }

    /// Get total number of unique terms
    pub fn term_count(&self) -> usize {
        self.posting_lists.len()
    }
}

/// Simple tokenizer - splits on whitespace and punctuation
fn tokenize(content: &str) -> Vec<String> {
    let mut tokens = Vec::new();
    let mut current_token = String::new();

    for ch in content.chars() {
        if ch.is_alphanumeric() || ch == '_' {
            current_token.push(ch);
        } else {
            if !current_token.is_empty() {
                tokens.push(current_token.clone());
                current_token.clear();
            }
        }
    }

    if !current_token.is_empty() {
        tokens.push(current_token);
    }

    // Filter out very short tokens and common noise
    tokens
        .into_iter()
        .filter(|t| t.len() >= 2 && !is_stop_word(t))
        .collect()
}

/// Check if a word is a common stop word
fn is_stop_word(word: &str) -> bool {
    matches!(word.to_lowercase().as_str(), 
        "the" | "a" | "an" | "and" | "or" | "but" | "in" | "on" | "at" | "to" | "for" | "of" | "with" | "by"
    )
}

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

    #[test]
    fn test_tokenize() {
        let content = "fn main() { println!(\"Hello, world!\"); }";
        let tokens = tokenize(content);
        assert!(tokens.contains(&"fn".to_string()));
        assert!(tokens.contains(&"main".to_string()));
        assert!(tokens.contains(&"println".to_string()));
        assert!(tokens.contains(&"Hello".to_string()));
        assert!(tokens.contains(&"world".to_string()));
    }

    #[test]
    fn test_inverted_index() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let fst_path = temp_dir.path().join("terms.fst");
        let json_path = temp_dir.path().join("posting_lists.json");

        let mut contents = HashMap::new();
        contents.insert(1, "fn main() { println!(\"Hello\"); }".to_string());
        contents.insert(2, "fn test() { assert_eq!(1, 1); }".to_string());
        contents.insert(3, "struct Point { x: i32, y: i32 }".to_string());

        let index = InvertedIndex::build_from_content(contents, &fst_path, &json_path)?;

        // Test term lookups
        let fn_files = index.find_files_with_term("fn");
        assert_eq!(fn_files.len(), 2);
        assert!(fn_files.contains(&1));
        assert!(fn_files.contains(&2));

        let struct_files = index.find_files_with_term("struct");
        assert_eq!(struct_files.len(), 1);
        assert!(struct_files.contains(&3));

        // Test AND query
        let main_fn = index.find_files_with_all_terms(&["fn", "main"]);
        assert_eq!(main_fn.len(), 1);
        assert!(main_fn.contains(&1));

        Ok(())
    }
}