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;
#[derive(Debug)]
pub struct PathFSTIndex {
path_map: Option<Map<Mmap>>,
handle_to_path: HashMap<u32, PathBuf>,
}
impl PathFSTIndex {
pub fn new() -> Self {
Self {
path_map: None,
handle_to_path: HashMap::new(),
}
}
pub fn build_from_paths<P: AsRef<Path>>(
paths: HashMap<PathBuf, u32>,
output_path: P,
) -> Result<Self> {
let mut sorted_paths: Vec<_> = paths.iter().collect();
sorted_paths.sort_by(|a, b| a.0.cmp(b.0));
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()?;
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,
})
}
pub fn load_from_file<P: AsRef<Path>>(
fst_path: P,
json_path: P,
) -> Result<Self> {
let file = File::open(&fst_path)?;
let mmap = unsafe { memmap2::Mmap::map(&file)? };
let path_map = Map::new(mmap)?;
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,
})
}
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(())
}
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)
}
pub fn get_path(&self, handle: u32) -> Option<&PathBuf> {
self.handle_to_path.get(&handle)
}
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; }
if let Ok(path_str) = std::str::from_utf8(path_bytes) {
let path = PathBuf::from(path_str);
results.push((path, handle as u32));
}
}
}
results
}
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 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 {
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)
}
pub fn len(&self) -> usize {
self.handle_to_path.len()
}
pub fn is_empty(&self) -> bool {
self.handle_to_path.is_empty()
}
}
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, _ => prefix.push(ch),
}
}
if prefix.contains('/') && !prefix.is_empty() {
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)?;
assert_eq!(index.get_handle(&PathBuf::from("src/main.rs")), Some(1));
assert_eq!(index.get_handle(&PathBuf::from("nonexistent")), None);
let src_files = index.find_by_prefix("src/");
assert_eq!(src_files.len(), 2);
Ok(())
}
}