_diffctx/

filtering.rs

use std::path::Path;
use std::sync::Arc;

use rayon::prelude::*;
use rustc_hash::{FxHashMap, FxHashSet};

use crate::config::extensions::CODE_EXTENSIONS;
use crate::config::filtering::FILTERING;
use crate::graph::{EdgeCategory, Graph};
use crate::types::{DiffHunk, Fragment, FragmentId, FragmentKind};

fn fragment_hunk_gap(frag_start: u32, frag_end: u32, hunk_start: u32, hunk_end: u32) -> u32 {
    if frag_end < hunk_start {
        hunk_start - frag_end
    } else if frag_start > hunk_end {
        frag_start - hunk_end
    } else {
        0
    }
}

fn proximity_score(frag: &Fragment, file_hunks: &[(u32, u32)]) -> f64 {
    let min_gap = file_hunks
        .iter()
        .map(|&(h_start, h_end)| {
            fragment_hunk_gap(frag.start_line(), frag.end_line(), h_start, h_end)
        })
        .min()
        .unwrap_or(u32::MAX);
    let half_decay = if frag.kind == FragmentKind::Definition {
        FILTERING.definition_proximity_half_decay
    } else {
        FILTERING.proximity_half_decay
    };
    FILTERING.proximity_floor_max / (1.0 + min_gap as f64 / half_decay)
}

fn effective_relevance_threshold(token_count: u32) -> f64 {
    let size_factor = (token_count as f64 / FILTERING.size_penalty_base_tokens)
        .max(1.0)
        .powf(FILTERING.size_penalty_exponent);
    FILTERING.low_relevance_threshold * size_factor
}

pub fn apply_hunk_proximity_bonus(
    rel: &mut FxHashMap<FragmentId, f64>,
    core_ids: &FxHashSet<FragmentId>,
    fragments: &[Fragment],
    hunks: &[DiffHunk],
) {
    let mut hunks_by_path: FxHashMap<&str, Vec<(u32, u32)>> = FxHashMap::default();
    for h in hunks {
        let (h_start, h_end) = h.core_selection_range();
        hunks_by_path
            .entry(h.path.as_ref())
            .or_default()
            .push((h_start, h_end));
    }

    let bonuses: Vec<(FragmentId, f64)> = fragments
        .par_iter()
        .filter(|frag| !core_ids.contains(&frag.id))
        .filter_map(|frag| {
            let file_hunks = hunks_by_path.get(frag.path())?;
            let bonus = proximity_score(frag, file_hunks);
            Some((frag.id.clone(), bonus))
        })
        .collect();

    for (id, bonus) in bonuses {
        let current = rel.get(&id).copied().unwrap_or(0.0);
        if current < bonus {
            rel.insert(id, bonus);
        }
    }
}

fn classify_semantic_edges(
    graph: &Graph,
    changed_paths: &FxHashSet<Arc<str>>,
) -> (
    FxHashMap<Arc<str>, FxHashSet<Arc<str>>>,
    FxHashSet<Arc<str>>,
) {
    let mut reverse_deps: FxHashMap<Arc<str>, FxHashSet<Arc<str>>> = FxHashMap::default();
    let mut direct_edge_paths: FxHashSet<Arc<str>> = FxHashSet::default();

    graph.for_each_categorized_edge(|src, dst, category| {
        if category != EdgeCategory::Semantic {
            return;
        }
        let src_changed = changed_paths.contains(&src.path);
        let dst_changed = changed_paths.contains(&dst.path);
        if !(src_changed ^ dst_changed) {
            return;
        }

        let (changed_frag, other_frag) = if src_changed { (src, dst) } else { (dst, src) };

        // `graph.edge_categories` is capped in lockstep with the CSR
        // (see `graph::assemble_graph`), so a categorized edge here is
        // guaranteed to exist in the CSR too -- `fwd_w == rev_w == 0.0`
        // can only mean a genuinely near-zero weight, never a
        // capped-away phantom silently suppressing hub-noise filtering.
        let fwd_w = graph
            .forward_edge_weight(changed_frag, other_frag)
            .unwrap_or(0.0);
        let rev_w = graph
            .forward_edge_weight(other_frag, changed_frag)
            .unwrap_or(0.0);

        if rev_w > fwd_w {
            reverse_deps
                .entry(changed_frag.path.clone())
                .or_default()
                .insert(other_frag.path.clone());
        } else {
            direct_edge_paths.insert(other_frag.path.clone());
        }
    });

    (reverse_deps, direct_edge_paths)
}

fn find_hub_noise_paths(graph: &Graph, changed_paths: &FxHashSet<Arc<str>>) -> FxHashSet<Arc<str>> {
    let (reverse_deps, direct_edge_paths) = classify_semantic_edges(graph, changed_paths);

    let changed_dirs: FxHashSet<String> = changed_paths
        .iter()
        .filter_map(|p| {
            Path::new(p.as_ref())
                .parent()
                .map(|d| d.to_string_lossy().into_owned())
        })
        .collect();

    let mut noise_counts: FxHashMap<Arc<str>, usize> = FxHashMap::default();
    for (hub_path, deps) in &reverse_deps {
        if changed_paths.contains(hub_path) {
            continue;
        }
        if deps.len() >= FILTERING.hub_reverse_threshold {
            for dep in deps {
                *noise_counts.entry(dep.clone()).or_insert(0) += 1;
            }
        }
    }

    noise_counts
        .into_iter()
        .filter(|(p, _count)| {
            !direct_edge_paths.contains(p)
                && !changed_dirs.contains(
                    &Path::new(p.as_ref())
                        .parent()
                        .map(|d| d.to_string_lossy().into_owned())
                        .unwrap_or_default(),
                )
        })
        .map(|(p, _)| p)
        .collect()
}

fn find_config_generic_code_files(
    graph: &Graph,
    changed_paths: &FxHashSet<Arc<str>>,
) -> FxHashSet<Arc<str>> {
    let mut has_real_edge: FxHashSet<Arc<str>> = FxHashSet::default();
    let mut has_generic_config: FxHashSet<Arc<str>> = FxHashSet::default();
    let mut generic_edge_count: FxHashMap<Arc<str>, usize> = FxHashMap::default();
    let config_stems: FxHashSet<String> = changed_paths
        .iter()
        .filter_map(|p| {
            Path::new(p.as_ref())
                .file_stem()
                .map(|s| s.to_string_lossy().to_lowercase())
        })
        .collect();

    graph.for_each_categorized_edge(|src, dst, category| {
        let src_changed = changed_paths.contains(&src.path);
        let dst_changed = changed_paths.contains(&dst.path);
        if !(src_changed ^ dst_changed) {
            return;
        }
        let other_path = if src_changed { &dst.path } else { &src.path };
        match category {
            EdgeCategory::ConfigGeneric => {
                has_generic_config.insert(other_path.clone());
                *generic_edge_count.entry(other_path.clone()).or_insert(0) += 1;
            }
            EdgeCategory::Semantic | EdgeCategory::Config => {
                has_real_edge.insert(other_path.clone());
            }
            _ => {}
        }
    });

    let generic_only: FxHashSet<Arc<str>> = has_generic_config
        .difference(&has_real_edge)
        .cloned()
        .collect();

    generic_only
        .into_iter()
        .filter(|p| {
            let path = Path::new(p.as_ref());
            let ext = path
                .extension()
                .map(|e| format!(".{}", e.to_string_lossy().to_lowercase()))
                .unwrap_or_default();
            let stem = path
                .file_stem()
                .map(|s| s.to_string_lossy().to_lowercase())
                .unwrap_or_default();
            CODE_EXTENSIONS.contains(ext.as_str())
                && generic_edge_count.get(p).copied().unwrap_or(0) <= 1
                && !config_stems.contains(&stem)
        })
        .collect()
}

pub fn filter_unrelated_fragments(
    fragments: &[Fragment],
    core_ids: &FxHashSet<FragmentId>,
    graph: &Graph,
) -> Vec<Fragment> {
    let changed_paths: FxHashSet<Arc<str>> = core_ids.iter().map(|fid| fid.path.clone()).collect();

    let mut paths_to_remove = find_hub_noise_paths(graph, &changed_paths);
    let config_generic = find_config_generic_code_files(graph, &changed_paths);
    for p in config_generic {
        paths_to_remove.insert(p);
    }
    for p in &changed_paths {
        paths_to_remove.remove(p);
    }

    fragments
        .iter()
        .filter(|f| !paths_to_remove.contains(&f.id.path))
        .cloned()
        .collect()
}

pub fn filter_low_relevance(
    fragments: Vec<Fragment>,
    core_ids: &FxHashSet<FragmentId>,
    rel: &FxHashMap<FragmentId, f64>,
) -> Vec<Fragment> {
    fragments
        .into_iter()
        .filter(|f| {
            core_ids.contains(&f.id)
                || rel.get(&f.id).copied().unwrap_or(0.0)
                    >= effective_relevance_threshold(f.token_count)
        })
        .collect()
}

pub fn filter_positive_relevance(
    fragments: Vec<Fragment>,
    core_ids: &FxHashSet<FragmentId>,
    rel: &FxHashMap<FragmentId, f64>,
) -> Vec<Fragment> {
    fragments
        .into_iter()
        .filter(|f| core_ids.contains(&f.id) || rel.get(&f.id).copied().unwrap_or(0.0) > 0.0)
        .collect()
}

pub fn cap_context_fragments(
    fragments: Vec<Fragment>,
    core_ids: &FxHashSet<FragmentId>,
    rel: &FxHashMap<FragmentId, f64>,
) -> Vec<Fragment> {
    let changed_paths: FxHashSet<Arc<str>> = core_ids.iter().map(|fid| fid.path.clone()).collect();

    let mut ctx_by_path: FxHashMap<Arc<str>, Vec<Fragment>> = FxHashMap::default();
    let mut result: Vec<Fragment> = Vec::new();

    for f in fragments {
        if changed_paths.contains(&f.id.path) {
            result.push(f);
        } else {
            ctx_by_path.entry(f.id.path.clone()).or_default().push(f);
        }
    }

    for (_path, mut file_frags) in ctx_by_path {
        if file_frags.len() <= FILTERING.max_context_fragments_per_file {
            result.extend(file_frags);
        } else {
            file_frags.sort_by(|a, b| {
                let sa = rel.get(&a.id).copied().unwrap_or(0.0);
                let sb = rel.get(&b.id).copied().unwrap_or(0.0);
                sb.total_cmp(&sa).then_with(|| a.id.cmp(&b.id))
            });
            result.extend(
                file_frags
                    .into_iter()
                    .take(FILTERING.max_context_fragments_per_file),
            );
        }
    }

    result.sort_by(|a, b| a.id.cmp(&b.id));
    result
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::types::FragmentKind;
    use std::sync::Arc;

    fn frag(path: &str, start: u32, end: u32) -> Fragment {
        Fragment {
            id: FragmentId::new(Arc::from(path), start, end),
            kind: FragmentKind::Function,
            content: Arc::from(""),
            identifiers: FxHashSet::default(),
            token_count: 10,
            symbol_name: None,
        }
    }

    #[test]
    fn cap_context_fragments_output_is_sorted_and_shuffle_invariant() {
        let changed_path = "changed.rs";
        let mut core_ids: FxHashSet<FragmentId> = FxHashSet::default();
        let mut fragments: Vec<Fragment> = Vec::new();
        for i in 0..5u32 {
            let f = frag(changed_path, i * 10, i * 10 + 5);
            core_ids.insert(f.id.clone());
            fragments.push(f);
        }

        let mut rel: FxHashMap<FragmentId, f64> = FxHashMap::default();
        let over_cap_path = "hub.rs";
        let n_context = FILTERING.max_context_fragments_per_file + 5;
        for i in 0..n_context {
            let f = frag(over_cap_path, (i as u32) * 10, (i as u32) * 10 + 5);
            // Distinct, strictly descending scores: no ties, so the
            // top-K selection itself is unambiguous and any remaining
            // non-determinism can only come from the final id sort.
            rel.insert(f.id.clone(), (n_context - i) as f64);
            fragments.push(f);
        }

        let baseline = cap_context_fragments(fragments.clone(), &core_ids, &rel);

        assert_eq!(
            baseline.len(),
            5 + FILTERING.max_context_fragments_per_file,
            "core fragments bypass the per-file cap; context fragments truncate to it"
        );

        let baseline_ids: Vec<FragmentId> = baseline.iter().map(|f| f.id.clone()).collect();
        let mut sorted_ids = baseline_ids.clone();
        sorted_ids.sort();
        assert_eq!(
            baseline_ids, sorted_ids,
            "cap_context_fragments output must be sorted by fragment id"
        );

        for shuffled in [
            {
                let mut v = fragments.clone();
                v.reverse();
                v
            },
            {
                let mut v = fragments.clone();
                v.rotate_left(7);
                v
            },
            {
                let mut v = fragments.clone();
                v.sort_by(|a, b| {
                    rel.get(&a.id)
                        .copied()
                        .unwrap_or(0.0)
                        .total_cmp(&rel.get(&b.id).copied().unwrap_or(0.0))
                });
                v
            },
        ] {
            let result = cap_context_fragments(shuffled, &core_ids, &rel);
            let ids: Vec<FragmentId> = result.iter().map(|f| f.id.clone()).collect();
            assert_eq!(
                ids, baseline_ids,
                "cap_context_fragments must be invariant under input ordering"
            );
        }
    }
}