siftdb_core/

fst_index.rs

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

/// FST-based path index for fast path prefix queries and globbing
#[derive(Debug)]
pub struct PathFSTIndex {
    /// Maps path -> file_handle as FST (using Mmap for efficiency)
    path_map: Option<Map<Mmap>>,
    /// Reverse lookup: file_handle -> path
    handle_to_path: HashMap<u32, PathBuf>,
}

impl PathFSTIndex {
    /// Create a new empty FST index
    pub fn new() -> Self {
        Self {
            path_map: None,
            handle_to_path: HashMap::new(),
        }
    }

    /// Build FST index from path->handle mappings
    pub fn build_from_paths<P: AsRef<Path>>(
        paths: HashMap<PathBuf, u32>,
        output_path: P,
    ) -> Result<Self> {
        // Sort paths for FST building
        let mut sorted_paths: Vec<_> = paths.iter().collect();
        sorted_paths.sort_by(|a, b| a.0.cmp(b.0));

        // Build FST
        let file = File::create(&output_path)?;
        let mut builder = MapBuilder::new(BufWriter::new(file))?;

        let mut handle_to_path = HashMap::new();

        for (path, &handle) in sorted_paths {
            let path_str = path.to_string_lossy();
            let path_bytes = path_str.as_bytes();
            
            builder.insert(path_bytes, handle as u64)?;
            handle_to_path.insert(handle, path.clone());
        }

        builder.finish()?;

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

        Ok(Self {
            path_map: Some(path_map),
            handle_to_path,
        })
    }

    /// Load existing FST index from file
    pub fn load_from_file<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 path_map = Map::new(mmap)?;

        // Load reverse mapping from JSON
        let file = File::open(&json_path)?;
        let reader = BufReader::new(file);
        let handle_to_path: HashMap<u32, PathBuf> = serde_json::from_reader(reader)
            .context("Failed to parse handle-to-path mapping")?;

        Ok(Self {
            path_map: Some(path_map),
            handle_to_path,
        })
    }

    /// Save reverse mapping to JSON file
    pub fn save_reverse_mapping<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let file = File::create(path)?;
        let writer = BufWriter::new(file);
        serde_json::to_writer_pretty(writer, &self.handle_to_path)?;
        Ok(())
    }

    /// Get file handle for exact path
    pub fn get_handle(&self, path: &Path) -> Option<u32> {
        let path_str = path.to_string_lossy();
        self.path_map.as_ref()?
            .get(path_str.as_bytes())
            .map(|h| h as u32)
    }

    /// Get path for file handle
    pub fn get_path(&self, handle: u32) -> Option<&PathBuf> {
        self.handle_to_path.get(&handle)
    }

    /// Find all paths matching a prefix
    pub fn find_by_prefix(&self, prefix: &str) -> Vec<(PathBuf, u32)> {
        let mut results = Vec::new();
        
        if let Some(ref path_map) = self.path_map {
            let prefix_bytes = prefix.as_bytes();
            let mut stream = path_map.range().ge(prefix_bytes).into_stream();
            
            while let Some((path_bytes, handle)) = stream.next() {
                if !path_bytes.starts_with(prefix_bytes) {
                    break; // FST is sorted, so we can break early
                }
                
                if let Ok(path_str) = std::str::from_utf8(path_bytes) {
                    let path = PathBuf::from(path_str);
                    results.push((path, handle as u32));
                }
            }
        }
        
        results
    }

    /// Find all paths matching a glob pattern using FST acceleration
    pub fn find_by_glob(&self, pattern: &str) -> Result<Vec<(PathBuf, u32)>> {
        let glob = globset::Glob::new(pattern)?;
        let matcher = glob.compile_matcher();
        
        let mut results = Vec::new();
        
        // If pattern has a clear prefix, use FST range query
        if let Some(prefix) = extract_prefix(pattern) {
            let candidates = self.find_by_prefix(&prefix);
            for (path, handle) in candidates {
                if matcher.is_match(&path) {
                    results.push((path, handle));
                }
            }
        } else if let Some(ref path_map) = self.path_map {
            // Fall back to full scan (still faster than linear due to FST structure)
            let mut stream = path_map.into_stream();
            while let Some((path_bytes, handle)) = stream.next() {
                if let Ok(path_str) = std::str::from_utf8(path_bytes) {
                    let path = PathBuf::from(path_str);
                    if matcher.is_match(&path) {
                        results.push((path, handle as u32));
                    }
                }
            }
        }
        
        Ok(results)
    }

    /// Get total number of indexed paths
    pub fn len(&self) -> usize {
        self.handle_to_path.len()
    }

    /// Check if index is empty
    pub fn is_empty(&self) -> bool {
        self.handle_to_path.is_empty()
    }
}

/// Extract a prefix from a glob pattern for FST acceleration
/// Examples: "src/**/*.rs" -> Some("src/")
///          "**/*.rs" -> None (no useful prefix)
fn extract_prefix(pattern: &str) -> Option<String> {
    let mut prefix = String::new();
    let chars: Vec<char> = pattern.chars().collect();
    
    for &ch in &chars {
        match ch {
            '*' | '?' | '[' | '{' => break, // Stop at first wildcard
            _ => prefix.push(ch),
        }
    }
    
    // Only return prefix if it's useful (at least one directory level)
    if prefix.contains('/') && !prefix.is_empty() {
        // Find the last directory separator to get a clean directory prefix
        if let Some(pos) = prefix.rfind('/') {
            Some(prefix[..=pos].to_string())
        } else {
            None
        }
    } else {
        None
    }
}

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

    #[test]
    fn test_extract_prefix() {
        assert_eq!(extract_prefix("src/**/*.rs"), Some("src/".to_string()));
        assert_eq!(extract_prefix("src/main.rs"), Some("src/".to_string()));
        assert_eq!(extract_prefix("**/*.rs"), None);
        assert_eq!(extract_prefix("*.rs"), None);
        assert_eq!(extract_prefix("tests/unit/*.rs"), Some("tests/unit/".to_string()));
    }

    #[test]
    fn test_fst_index_basic() -> Result<()> {
        let temp_dir = TempDir::new()?;
        let fst_path = temp_dir.path().join("paths.fst");
        
        let mut paths = HashMap::new();
        paths.insert(PathBuf::from("src/main.rs"), 1);
        paths.insert(PathBuf::from("src/lib.rs"), 2);
        paths.insert(PathBuf::from("tests/test.rs"), 3);
        
        let index = PathFSTIndex::build_from_paths(paths, &fst_path)?;
        
        // Test exact lookup
        assert_eq!(index.get_handle(&PathBuf::from("src/main.rs")), Some(1));
        assert_eq!(index.get_handle(&PathBuf::from("nonexistent")), None);
        
        // Test prefix search
        let src_files = index.find_by_prefix("src/");
        assert_eq!(src_files.len(), 2);
        
        Ok(())
    }
}