use std::path::{Path, PathBuf};
use std::sync::Arc;
use once_cell::sync::Lazy;
use rayon::prelude::*;
use regex::Regex;
use rustc_hash::FxHashSet;
use crate::config::fragmentation::FRAGMENTATION;
use crate::config::limits::LIMITS;
use crate::config::tokenization::TOKENIZATION;
use crate::git::{self, CatFileBatch};
use crate::parsers::fragment_file;
use crate::tokenizer::count_tokens;
use crate::types::{Fragment, FragmentId, FragmentKind, extract_identifiers};
static BINARY_CTRL_RE: Lazy<Regex> = Lazy::new(|| Regex::new(r"\x00").unwrap());
static GENERATED_FILENAME_PATTERNS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
[
".pb.go",
"_pb2.py",
"_pb2_grpc.py",
".pb.h",
".pb.cc",
".pb.swift",
".min.js",
".min.css",
".designer.cs",
".api",
]
.into_iter()
.collect()
});
const GENERATED_FILENAME_SUFFIXES: &[&str] = &["_generated.", "OuterClass.java"];
static GENERATED_PATH_SEGMENTS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
[
"generated",
"gen-java",
"gen-go",
"gen-py",
"gen-cpp",
"gen-swift",
"__generated__",
"autogen",
"codegen",
]
.into_iter()
.collect()
});
const GENERATED_CONTENT_MARKERS: &[&str] = &[
"@generated",
"do not edit",
"code generated",
"auto-generated",
"this file is generated",
"generated by",
"automatically generated",
"auto generated",
];
static KNOWN_BINARY_EXTENSIONS: Lazy<FxHashSet<&'static str>> = Lazy::new(|| {
[
".png", ".jpg", ".jpeg", ".gif", ".bmp", ".ico", ".svg", ".webp", ".mp3", ".mp4", ".wav",
".ogg", ".flac", ".avi", ".mkv", ".mov", ".zip", ".gz", ".tar", ".bz2", ".xz", ".7z",
".rar", ".jar", ".war", ".ear", ".class", ".pyc", ".pyo", ".o", ".a", ".so", ".dylib",
".dll", ".exe", ".bin", ".dat", ".db", ".sqlite", ".pdf", ".doc", ".docx", ".xls", ".xlsx",
".ppt", ".pptx", ".woff", ".woff2", ".ttf", ".otf", ".eot",
]
.into_iter()
.collect()
});
fn looks_binary(content: &str) -> bool {
let mut check_len = content
.len()
.min(FRAGMENTATION.binary_detection_buffer_size);
while check_len > 0 && !content.is_char_boundary(check_len) {
check_len -= 1;
}
BINARY_CTRL_RE.is_match(&content[..check_len])
}
fn has_generated_filename(name: &str) -> bool {
GENERATED_FILENAME_PATTERNS
.iter()
.any(|p| name.ends_with(p))
|| GENERATED_FILENAME_SUFFIXES
.iter()
.any(|s| name.ends_with(s))
}
fn has_generated_path_segment(path: &Path) -> bool {
path.components().any(|c| {
let s = c.as_os_str().to_string_lossy().to_lowercase();
GENERATED_PATH_SEGMENTS.contains(s.as_str())
})
}
fn has_generated_content_marker(content: &str) -> bool {
let header: String = content
.lines()
.take(FRAGMENTATION.generated_marker_header_lines)
.collect::<Vec<_>>()
.join("\n")
.to_lowercase();
for marker in GENERATED_CONTENT_MARKERS {
if !header.contains(marker) {
continue;
}
if *marker != "@generated" {
return true;
}
if header.contains("@generated") {
let after_idx = header.find("@generated").unwrap() + "@generated".len();
let next_char = header[after_idx..].chars().next();
if next_char.is_none() || !next_char.unwrap().is_ascii_lowercase() {
return true;
}
}
}
false
}
fn is_generated_file(path: &Path, content: &str) -> bool {
let name = path
.file_name()
.map(|n| n.to_string_lossy().to_string())
.unwrap_or_default();
has_generated_filename(&name)
|| has_generated_path_segment(path)
|| has_generated_content_marker(content)
}
fn truncate_generated_fragments(file_frags: Vec<Fragment>) -> Vec<Fragment> {
let max_lines = LIMITS.max_generated_lines as u32;
file_frags
.into_iter()
.map(|frag| {
if frag.line_count() <= max_lines {
return frag;
}
let lines: Vec<&str> = frag.content.lines().collect();
let remaining = lines.len() - max_lines as usize;
let truncated_lines = &lines[..max_lines as usize];
let truncated_content = format!(
"{}\n# ... [{} more lines]",
truncated_lines.join("\n"),
remaining
);
let new_end = frag.start_line() + max_lines - 1;
let identifiers = extract_identifiers(
&truncated_content,
TOKENIZATION.fragment_min_identifier_length,
);
Fragment {
id: FragmentId::new(frag.id.path.clone(), frag.start_line(), new_end),
kind: frag.kind,
content: Arc::from(truncated_content),
identifiers,
token_count: 0,
symbol_name: frag.symbol_name,
}
})
.collect()
}
fn dedup_fragments(raw_frags: Vec<Fragment>, seen: &mut FxHashSet<FragmentId>) -> Vec<Fragment> {
let mut result = Vec::new();
for f in raw_frags {
if !seen.contains(&f.id) {
seen.insert(f.id.clone());
result.push(f);
}
}
result
}
fn normalize_path(path: &Path, root_dir: &Path) -> PathBuf {
if path.is_absolute() {
path.canonicalize().unwrap_or_else(|_| path.to_path_buf())
} else {
let joined = root_dir.join(path);
joined.canonicalize().unwrap_or_else(|_| joined)
}
}
fn read_file_content(
file_path: &Path,
root_dir: &Path,
preferred_revs: &[String],
mut batch_reader: Option<&mut CatFileBatch>,
is_changed: bool,
) -> Option<String> {
let ext = file_path
.extension()
.map(|e| format!(".{}", e.to_string_lossy().to_lowercase()))
.unwrap_or_default();
if KNOWN_BINARY_EXTENSIONS.contains(ext.as_str()) {
return None;
}
let abs_path = normalize_path(file_path, root_dir);
let resolved_root = root_dir
.canonicalize()
.unwrap_or_else(|_| root_dir.to_path_buf());
let rel = abs_path.strip_prefix(&resolved_root).ok()?;
let max_size = if is_changed {
LIMITS.max_changed_file_size
} else {
LIMITS.max_file_size
};
for rev in preferred_revs {
if let Some(reader) = batch_reader.as_deref_mut() {
match reader.get(rev, rel) {
Ok(content) if content.len() <= max_size && !looks_binary(&content) => {
return Some(content);
}
_ => continue,
}
} else {
match git::show_file_at_revision(root_dir, rev, rel) {
Ok(content) if content.len() <= max_size && !looks_binary(&content) => {
return Some(content);
}
_ => continue,
}
}
}
if abs_path.exists() && abs_path.is_file() {
if let Ok(meta) = std::fs::metadata(&abs_path) {
if meta.len() as usize > max_size {
return None;
}
}
if let Ok(content) = std::fs::read_to_string(&abs_path) {
if !looks_binary(&content) {
return Some(content);
}
}
}
None
}
pub fn process_files_for_fragments(
files: &[PathBuf],
root_dir: &Path,
preferred_revs: &[String],
seen_frag_ids: &mut FxHashSet<FragmentId>,
mut batch_reader: Option<&mut CatFileBatch>,
is_changed: bool,
) -> Vec<Fragment> {
let max_frags = LIMITS.max_fragments;
let max_generated = LIMITS.max_generated_fragments;
let chunk_size = rayon::current_num_threads().max(1);
let mut parsed: Vec<Vec<Fragment>> = Vec::with_capacity(files.len());
for chunk in files.chunks(chunk_size) {
let chunk_contents: Vec<(PathBuf, String)> = chunk
.iter()
.filter_map(|file_path| {
let content = read_file_content(
file_path,
root_dir,
preferred_revs,
batch_reader.as_deref_mut(),
is_changed,
)?;
Some((file_path.clone(), content))
})
.collect();
parsed.extend(
chunk_contents
.par_iter()
.map(|(file_path, content)| {
let path_arc: Arc<str> = Arc::from(file_path.to_string_lossy().as_ref());
let mut raw_frags = fragment_file(path_arc, content);
let generated = !is_changed && is_generated_file(file_path, content);
let cap = if generated {
max_generated
} else if is_changed {
max_frags.saturating_mul(10)
} else {
max_frags
};
if raw_frags.len() > cap {
raw_frags.sort_by(|a, b| b.line_count().cmp(&a.line_count()));
raw_frags.truncate(cap);
}
if generated {
raw_frags = truncate_generated_fragments(raw_frags);
}
raw_frags
})
.collect::<Vec<_>>(),
);
}
let mut fragments: Vec<Fragment> = Vec::new();
for file_frags in parsed {
for frag in dedup_fragments(file_frags, seen_frag_ids) {
seen_frag_ids.insert(frag.id.clone());
fragments.push(frag);
}
}
fragments
}
pub fn create_whole_file_fragment(
path: &Path,
root_dir: &Path,
preferred_revs: &[String],
batch_reader: Option<&mut CatFileBatch>,
) -> Option<Fragment> {
let content = read_file_content(path, root_dir, preferred_revs, batch_reader, true)?;
let trimmed = content.trim();
if trimmed.is_empty() {
return None;
}
let content = if is_generated_file(path, &content) {
let lines: Vec<&str> = content.lines().collect();
let max_lines = LIMITS.max_generated_lines;
if lines.len() > max_lines {
let remaining = lines.len() - max_lines;
format!(
"{}\n# ... [{} more lines]",
lines[..max_lines].join("\n"),
remaining
)
} else {
content
}
} else {
content
};
let lines: Vec<&str> = content.lines().collect();
let line_count = lines.len() as u32;
let path_arc: Arc<str> = Arc::from(path.to_string_lossy().as_ref());
let token_count = count_tokens(&content) + LIMITS.overhead_per_fragment;
let identifiers = extract_identifiers(&content, TOKENIZATION.fragment_min_identifier_length);
Some(Fragment {
id: FragmentId::new(path_arc, 1, line_count),
kind: FragmentKind::Chunk,
content: Arc::from(content),
identifiers,
token_count,
symbol_name: None,
})
}