_diffctx/

fragmentation.rs

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};

// Content reaching here is already-decoded UTF-8 text, so the only reliable
// binary signal is an embedded NUL (matching git's own heuristic). The old
// range flagged ESC/BS/etc., wrongly dropping changed text fixtures that embed
// ANSI escape codes (snapshot/terminal-recording files).
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;

    // Process files in chunks: sequential read (CatFileBatch is &mut, !Send) then
    // parallel parse within each chunk. Peak raw-content memory = chunk_size × max_file_size
    // instead of N_files × avg_file_size — eliminates OOM on large repos like astropy.
    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);
                    // Changed files are the subject of the diff: never apply the
                    // aggressive generated-file reduction (cap=5 + 30-line content
                    // truncation), which can drop the small fragment covering the
                    // edited hunk before core identification runs.
                    let generated = !is_changed && is_generated_file(file_path, content);
                    // Changed files also get 10x cap headroom: the biggest-N
                    // truncation below keeps the LONGEST fragments, and the
                    // edited hunk is typically a small leaf that would be
                    // dropped first (hunks are only known later, in core
                    // identification). max_changed_file_size still bounds cost.
                    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<_>>(),
        );
        // chunk_contents dropped here — raw file text freed before next chunk
    }

    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,
    })
}