use fxhash::FxHashSet;
use memchr::memmem;
use once_cell::sync::Lazy;
use scribe_core::FileInfo;
use std::collections::HashSet;
use std::path::{Path, PathBuf};
static COLD_EXTENSIONS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
[
"md", "txt", "rst", "adoc", "wiki", "png", "jpg", "jpeg", "gif", "bmp", "ico", "svg", "webp", "tiff", "mp3", "mp4", "avi",
"mkv", "mov", "wmv", "flv", "webm", "m4v", "wav", "flac", "ogg", "aac", "wma",
"zip", "tar", "gz", "bz2", "xz", "7z", "rar", "jar", "war", "ear",
"exe", "dll", "so", "dylib", "a", "lib", "bin", "out", "pdf", "doc", "docx", "xls", "xlsx", "ppt", "pptx", "odt", "ods", "odp",
"ttf", "otf", "woff", "woff2", "eot", "tmp", "temp", "cache", "log", "bak", "swp", "swo", "min.js", "min.css",
]
.into_iter()
.collect()
});
static HOT_EXTENSIONS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
[
"rs",
"py",
"js",
"ts",
"jsx",
"tsx",
"go",
"java",
"c",
"cpp",
"h",
"hpp",
"cs",
"php",
"rb",
"swift",
"kt",
"scala",
"clj",
"hs",
"elm",
"ml",
"ocaml",
"json",
"yaml",
"yml",
"toml",
"xml",
"html",
"css",
"scss",
"less",
"sass",
"sh",
"bash",
"zsh",
"fish",
"ps1",
"cmd",
"bat",
"dockerfile",
"makefile",
"sql",
"graphql",
"prisma",
]
.into_iter()
.collect()
});
static COLD_DIRS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
[
"node_modules",
"__pycache__",
".pytest_cache",
".mypy_cache",
"target",
"build",
"dist",
".git",
".hg",
".svn",
"vendor",
"third_party",
"external",
"deps",
".idea",
".vscode",
".vs",
".gradle",
".maven",
"coverage",
".coverage",
".nyc_output",
"logs",
"tmp",
"temp",
".tmp",
".temp",
]
.into_iter()
.collect()
});
static BINARY_MARKERS: Lazy<Vec<&'static [u8]>> = Lazy::new(|| {
vec![
b"\x7fELF", b"MZ", b"\xca\xfe\xba\xbe", b"\xfe\xed\xfa\xce", b"\x89PNG", b"\xff\xd8\xff", b"GIF8", b"RIFF", b"%PDF", b"PK\x03\x04", ]
});
const MAX_CONTENT_SIZE: u64 = 8 * 1024 * 1024;
const BINARY_SAMPLE_SIZE: usize = 512;
#[derive(Debug)]
pub struct FileFilter {
allow_extensions: Option<FxHashSet<String>>,
deny_extensions: FxHashSet<String>,
max_file_size: u64,
include_hidden: bool,
binary_detection: bool,
stats: FilterStats,
}
#[derive(Debug, Default, Clone)]
pub struct FilterStats {
pub files_walked: u64,
pub dirs_skipped: u64,
pub extension_filtered: u64,
pub size_filtered: u64,
pub binary_filtered: u64,
pub passed_filter: u64,
pub bytes_read_for_detection: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FilterResult {
Include,
Exclude(FilterReason),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FilterReason {
ColdExtension,
ColdDirectory,
TooLarge(u64),
Hidden,
Binary,
CustomExtensionFilter,
}
impl FileFilter {
pub fn new() -> Self {
Self {
allow_extensions: None,
deny_extensions: FxHashSet::default(),
max_file_size: MAX_CONTENT_SIZE,
include_hidden: false,
binary_detection: true,
stats: FilterStats::default(),
}
}
pub fn with_allow_extensions(mut self, extensions: Vec<String>) -> Self {
self.allow_extensions = Some(extensions.into_iter().map(|e| e.to_lowercase()).collect());
self
}
pub fn with_deny_extensions(mut self, extensions: Vec<String>) -> Self {
self.deny_extensions = extensions.into_iter().map(|e| e.to_lowercase()).collect();
self
}
pub fn with_max_file_size(mut self, size: u64) -> Self {
self.max_file_size = size;
self
}
pub fn with_include_hidden(mut self, include: bool) -> Self {
self.include_hidden = include;
self
}
pub fn with_binary_detection(mut self, detect: bool) -> Self {
self.binary_detection = detect;
self
}
pub fn pre_filter_path(&mut self, path: &Path) -> FilterResult {
self.stats.files_walked += 1;
if !self.include_hidden {
if let Some(name) = path.file_name() {
if name.to_string_lossy().starts_with('.') {
return FilterResult::Exclude(FilterReason::Hidden);
}
}
}
for component in path.components() {
if let std::path::Component::Normal(name) = component {
if COLD_DIRS.contains(name.to_str().unwrap_or("")) {
self.stats.dirs_skipped += 1;
return FilterResult::Exclude(FilterReason::ColdDirectory);
}
}
}
let extension = path
.extension()
.and_then(|ext| ext.to_str())
.unwrap_or("")
.to_lowercase();
if let Some(ref allow_list) = self.allow_extensions {
if !allow_list.contains(&extension) {
self.stats.extension_filtered += 1;
return FilterResult::Exclude(FilterReason::CustomExtensionFilter);
}
}
if self.deny_extensions.contains(&extension) {
self.stats.extension_filtered += 1;
return FilterResult::Exclude(FilterReason::CustomExtensionFilter);
}
if COLD_EXTENSIONS.contains(extension.as_str()) {
self.stats.extension_filtered += 1;
return FilterResult::Exclude(FilterReason::ColdExtension);
}
FilterResult::Include
}
pub async fn filter_file(&mut self, path: &Path) -> FilterResult {
match self.pre_filter_path(path) {
FilterResult::Exclude(reason) => return FilterResult::Exclude(reason),
FilterResult::Include => {}
}
if let Ok(metadata) = tokio::fs::metadata(path).await {
if metadata.len() > self.max_file_size {
self.stats.size_filtered += 1;
return FilterResult::Exclude(FilterReason::TooLarge(metadata.len()));
}
if self.binary_detection && self.should_check_binary(path) {
if self.is_binary_file(path).await {
self.stats.binary_filtered += 1;
return FilterResult::Exclude(FilterReason::Binary);
}
}
}
self.stats.passed_filter += 1;
FilterResult::Include
}
fn should_check_binary(&self, path: &Path) -> bool {
let extension = path
.extension()
.and_then(|ext| ext.to_str())
.unwrap_or("")
.to_lowercase();
if HOT_EXTENSIONS.contains(extension.as_str()) {
return false;
}
if extension.is_empty() {
return false;
}
true
}
pub async fn is_binary_file(&mut self, path: &Path) -> bool {
match tokio::fs::File::open(path).await {
Ok(mut file) => {
use tokio::io::AsyncReadExt;
let mut buffer = vec![0u8; BINARY_SAMPLE_SIZE];
match file.read(&mut buffer).await {
Ok(bytes_read) => {
self.stats.bytes_read_for_detection += bytes_read as u64;
buffer.truncate(bytes_read);
let extension = path.extension().and_then(|ext| ext.to_str());
if FileInfo::detect_binary_from_bytes(&buffer, extension) {
return true;
}
self.detect_binary_content(&buffer)
}
Err(_) => false, }
}
Err(_) => false, }
}
fn detect_binary_content(&self, content: &[u8]) -> bool {
for marker in BINARY_MARKERS.iter() {
if content.starts_with(marker) {
return true;
}
}
if memchr::memchr(0, content).is_some() {
return true;
}
let non_printable = content
.iter()
.filter(|&&b| b < 32 && b != b'\t' && b != b'\n' && b != b'\r')
.count();
let ratio = non_printable as f64 / content.len() as f64;
ratio > 0.05 }
pub fn stats(&self) -> &FilterStats {
&self.stats
}
pub fn reset_stats(&mut self) {
self.stats = FilterStats::default();
}
}
impl Default for FileFilter {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug)]
pub struct DirectoryFilter {
cold_dirs: FxHashSet<String>,
stats: DirectoryFilterStats,
}
#[derive(Debug, Default)]
pub struct DirectoryFilterStats {
pub dirs_walked: u64,
pub dirs_skipped: u64,
}
impl DirectoryFilter {
pub fn new() -> Self {
Self {
cold_dirs: COLD_DIRS.iter().map(|s| s.to_string()).collect(),
stats: DirectoryFilterStats::default(),
}
}
pub fn with_additional_cold_dirs(mut self, dirs: Vec<String>) -> Self {
self.cold_dirs.extend(dirs);
self
}
pub fn should_skip_directory(&mut self, path: &Path) -> bool {
self.stats.dirs_walked += 1;
if let Some(name) = path.file_name() {
if let Some(name_str) = name.to_str() {
if self.cold_dirs.contains(name_str) {
self.stats.dirs_skipped += 1;
return true;
}
}
}
false
}
pub fn stats(&self) -> &DirectoryFilterStats {
&self.stats
}
}
impl Default for DirectoryFilter {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::TempDir;
use tokio::fs;
#[tokio::test]
async fn test_cold_extension_filtering() {
let mut filter = FileFilter::new();
assert_eq!(
filter.pre_filter_path(Path::new("test.png")),
FilterResult::Exclude(FilterReason::ColdExtension)
);
assert_eq!(
filter.pre_filter_path(Path::new("code.rs")),
FilterResult::Include
);
}
#[tokio::test]
async fn test_cold_directory_filtering() {
let mut filter = FileFilter::new();
assert_eq!(
filter.pre_filter_path(Path::new("node_modules/package/index.js")),
FilterResult::Exclude(FilterReason::ColdDirectory)
);
assert_eq!(
filter.pre_filter_path(Path::new("src/main.rs")),
FilterResult::Include
);
}
#[tokio::test]
async fn test_custom_extension_filtering() {
let mut filter =
FileFilter::new().with_allow_extensions(vec!["rs".to_string(), "py".to_string()]);
assert_eq!(
filter.pre_filter_path(Path::new("test.js")),
FilterResult::Exclude(FilterReason::CustomExtensionFilter)
);
assert_eq!(
filter.pre_filter_path(Path::new("test.rs")),
FilterResult::Include
);
}
#[tokio::test]
async fn test_file_size_filtering() {
let large_file = Path::new("test_large_file.rs");
let content = "x".repeat(2000);
fs::write(&large_file, &content).await.unwrap();
let mut filter = FileFilter::new().with_max_file_size(1000);
let result = filter.filter_file(&large_file).await;
let _ = fs::remove_file(&large_file).await;
match result {
FilterResult::Exclude(FilterReason::TooLarge(_)) => {}
other => panic!("Expected TooLarge, got {:?}", other),
}
}
#[tokio::test]
async fn test_binary_detection() {
let temp_dir = TempDir::new().unwrap();
let test_dir = temp_dir.path().join("project");
fs::create_dir_all(&test_dir).await.unwrap();
let binary_file = test_dir.join("binary.dat");
fs::write(&binary_file, &[0u8, 1u8, 2u8, 0u8])
.await
.unwrap();
let text_file = test_dir.join("text.txt");
fs::write(&text_file, "Hello, world!").await.unwrap();
let mut filter = FileFilter::new();
assert!(filter.is_binary_file(&binary_file).await);
assert!(!filter.is_binary_file(&text_file).await);
}
#[tokio::test]
async fn test_hidden_file_filtering() {
let mut filter = FileFilter::new().with_include_hidden(false);
assert_eq!(
filter.pre_filter_path(Path::new(".hidden")),
FilterResult::Exclude(FilterReason::Hidden)
);
let mut filter = FileFilter::new().with_include_hidden(true);
assert_eq!(
filter.pre_filter_path(Path::new(".hidden")),
FilterResult::Include
);
}
#[test]
fn test_binary_content_detection() {
let filter = FileFilter::new();
assert!(filter.detect_binary_content(b"\x7fELF\x01\x01\x01"));
assert!(filter.detect_binary_content(b"%PDF-1.4\n"));
assert!(filter.detect_binary_content(b"text\x00more text"));
assert!(!filter.detect_binary_content(b"Hello, world!\n"));
assert!(!filter.detect_binary_content(b"fn main() {\n\tprintln!(\"Hello\");\n}"));
}
#[test]
fn test_directory_filtering() {
let mut dir_filter = DirectoryFilter::new();
assert!(dir_filter.should_skip_directory(Path::new("node_modules")));
assert!(dir_filter.should_skip_directory(Path::new("target")));
assert!(!dir_filter.should_skip_directory(Path::new("src")));
assert_eq!(dir_filter.stats().dirs_walked, 3);
assert_eq!(dir_filter.stats().dirs_skipped, 2);
}
#[test]
fn test_filter_statistics() {
let mut filter = FileFilter::new();
filter.pre_filter_path(Path::new("test.rs")); filter.pre_filter_path(Path::new("test.png")); filter.pre_filter_path(Path::new("node_modules/pkg/index.js")); filter.pre_filter_path(Path::new(".hidden"));
let stats = filter.stats();
assert_eq!(stats.files_walked, 4);
assert_eq!(stats.extension_filtered, 1);
assert_eq!(stats.dirs_skipped, 1);
assert_eq!(stats.passed_filter, 0); }
}