ripvec_core/encoder/ripvec/

penalties.rs

//! File-path penalties + `rerank_topk` with file-saturation decay.
//!
//! Port of `~/src/semble/src/semble/ranking/penalties.py`. Two surfaces:
//!
//! 1. [`file_path_penalty`] — multiplicative penalty per file path.
//!    Combines penalties for test files, compat/legacy/examples
//!    directories, re-export barrels (`__init__.py`, `package-info.java`),
//!    and `.d.ts` declaration stubs.
//! 2. [`rerank_topk`] — greedy top-k selection that applies path
//!    penalties (when `penalise_paths == true`) then decays by 0.5 per
//!    extra chunk from the same file beyond a threshold of 1.
//!
//! ## Indexing convention
//!
//! Where Python uses `dict[Chunk, float]` keyed by hashable Chunk,
//! Rust uses `(chunk_index, score)` pairs — the same convention ripvec
//! already uses in [`crate::hybrid`]. This avoids adding `Hash`/`Eq`
//! impls to [`CodeChunk`] just to satisfy a HashMap key.

use std::collections::HashMap;
use std::path::Path;
use std::sync::OnceLock;

use regex::{Regex, RegexBuilder};

use crate::chunk::CodeChunk;

// ---------------------------------------------------------------------------
// Penalty constants (Python: `_STRONG_PENALTY` etc. in penalties.py:67).
// ---------------------------------------------------------------------------

/// Test files, compat shims, example/doc code.
pub const STRONG_PENALTY: f32 = 0.3;
/// Re-export / metadata files (`__init__.py`, `package-info.java`).
pub const MODERATE_PENALTY: f32 = 0.5;
/// `.d.ts` declaration stubs (still carry useful type info).
pub const MILD_PENALTY: f32 = 0.7;

/// Maximum chunks from the same file before saturation penalty applies.
pub const FILE_SATURATION_THRESHOLD: usize = 1;

/// Multiplicative penalty per extra chunk from the same file beyond
/// [`FILE_SATURATION_THRESHOLD`]. Excess chunks pay `decay^excess`.
pub const FILE_SATURATION_DECAY: f32 = 0.5;

// ---------------------------------------------------------------------------
// Path-prior regex bank (Python: penalties.py:8-65).
// ---------------------------------------------------------------------------

/// Regex matching test files across 14 languages plus shared helpers.
///
/// Mirrors `_TEST_FILE_RE` from `penalties.py:8`. Lazy-compiled.
fn test_file_re() -> &'static Regex {
    static RE: OnceLock<Regex> = OnceLock::new();
    RE.get_or_init(|| {
        // Match exactly the Python pattern. The non-capture group after
        // `(?:^|/)` lists per-language test-file conventions; the
        // pattern terminates with `$`.
        let pattern = concat!(
            r"(?:^|/)",
            r"(?:",
            // Python
            r"test_[^/]*\.py",
            r"|[^/]*_test\.py",
            // Go
            r"|[^/]*_test\.go",
            // Java
            r"|[^/]*Tests?\.java",
            // PHP
            r"|[^/]*Test\.php",
            // Ruby
            r"|[^/]*_spec\.rb",
            r"|[^/]*_test\.rb",
            // JavaScript / TypeScript
            r"|[^/]*\.test\.[jt]sx?",
            r"|[^/]*\.spec\.[jt]sx?",
            // Kotlin
            r"|[^/]*Tests?\.kt",
            r"|[^/]*Spec\.kt",
            // Swift
            r"|[^/]*Tests?\.swift",
            r"|[^/]*Spec\.swift",
            // C#
            r"|[^/]*Tests?\.cs",
            // C / C++
            r"|test_[^/]*\.cpp",
            r"|[^/]*_test\.cpp",
            r"|test_[^/]*\.c",
            r"|[^/]*_test\.c",
            // Scala
            r"|[^/]*Spec\.scala",
            r"|[^/]*Suite\.scala",
            r"|[^/]*Test\.scala",
            // Dart
            r"|[^/]*_test\.dart",
            r"|test_[^/]*\.dart",
            // Lua
            r"|[^/]*_spec\.lua",
            r"|[^/]*_test\.lua",
            r"|test_[^/]*\.lua",
            // Shared helper patterns (all languages)
            r"|test_helpers?[^/]*\.\w+",
            r")$",
        );
        Regex::new(pattern).expect("test-file regex compiles")
    })
}

/// Regex matching directories named `test`, `tests`, `__tests__`, `spec`,
/// `testing`. Mirrors `_TEST_DIR_RE` from `penalties.py:56`.
fn test_dir_re() -> &'static Regex {
    static RE: OnceLock<Regex> = OnceLock::new();
    RE.get_or_init(|| {
        Regex::new(r"(?:^|/)(?:tests?|__tests__|spec|testing)(?:/|$)")
            .expect("test-dir regex compiles")
    })
}

/// Regex matching compat / legacy directories. Mirrors `_COMPAT_DIR_RE`.
fn compat_dir_re() -> &'static Regex {
    static RE: OnceLock<Regex> = OnceLock::new();
    RE.get_or_init(|| {
        Regex::new(r"(?:^|/)(?:compat|_compat|legacy)(?:/|$)").expect("compat-dir regex compiles")
    })
}

/// Regex matching examples / docs-src directories. Mirrors
/// `_EXAMPLES_DIR_RE`.
fn examples_dir_re() -> &'static Regex {
    static RE: OnceLock<Regex> = OnceLock::new();
    RE.get_or_init(|| {
        Regex::new(r"(?:^|/)(?:_?examples?|docs?_src)(?:/|$)").expect("examples-dir regex compiles")
    })
}

/// Regex matching TypeScript `.d.ts` declaration stubs.
fn type_defs_re() -> &'static Regex {
    static RE: OnceLock<Regex> = OnceLock::new();
    RE.get_or_init(|| {
        RegexBuilder::new(r"\.d\.ts$")
            .build()
            .expect("dts regex compiles")
    })
}

/// Filenames that are re-export barrels or package-level metadata.
const REEXPORT_FILENAMES: &[&str] = &["__init__.py", "package-info.java"];

// ---------------------------------------------------------------------------
// Public API.
// ---------------------------------------------------------------------------

/// Combined multiplicative penalty for all applicable path patterns.
///
/// Always in `(0.0, 1.0]`. Returns `1.0` when no patterns match.
/// Each matched category multiplies into the result independently, so a
/// file matching both a test-dir pattern and a re-export filename pays
/// `STRONG_PENALTY * MODERATE_PENALTY`.
///
/// Backslashes in `file_path` are normalised to forward slashes before
/// matching, so Windows-style paths work too.
#[must_use]
pub fn file_path_penalty(file_path: &str) -> f32 {
    let normalised = file_path.replace('\\', "/");
    let mut penalty = 1.0_f32;
    if test_file_re().is_match(&normalised) || test_dir_re().is_match(&normalised) {
        penalty *= STRONG_PENALTY;
    }
    if let Some(filename) = Path::new(file_path).file_name().and_then(|f| f.to_str())
        && REEXPORT_FILENAMES.contains(&filename)
    {
        penalty *= MODERATE_PENALTY;
    }
    if compat_dir_re().is_match(&normalised) {
        penalty *= STRONG_PENALTY;
    }
    if examples_dir_re().is_match(&normalised) {
        penalty *= STRONG_PENALTY;
    }
    if type_defs_re().is_match(&normalised) {
        penalty *= MILD_PENALTY;
    }
    penalty
}

/// Select the top-k results with optional path penalties and
/// file-saturation decay.
///
/// Mirrors `rerank_topk` from `penalties.py:81`. The greedy pass uses
/// the early-exit optimisation: once `selected.len() >= top_k` and the
/// remaining candidate's penalised score cannot beat the current k-th
/// best, the loop terminates.
///
/// - `penalise_paths == true` — apply [`file_path_penalty`] before
///   sorting and selection. This matches semble's behaviour for hybrid
///   and BM25-led queries.
/// - `penalise_paths == false` — bypass path priors (pure-semantic
///   queries where the priors don't apply).
///
/// Returns `(chunk_index, effective_score)` pairs, highest first, with
/// length capped at `top_k`.
#[must_use]
pub fn rerank_topk(
    scores: &[(usize, f32)],
    chunks: &[CodeChunk],
    top_k: usize,
    penalise_paths: bool,
) -> Vec<(usize, f32)> {
    if scores.is_empty() || top_k == 0 {
        return Vec::new();
    }

    // Step 1: apply path penalties (or identity), producing (chunk_index, penalised_score).
    let mut penalty_cache: HashMap<String, f32> = HashMap::new();
    let mut penalised: Vec<(usize, f32)> = scores
        .iter()
        .map(|&(idx, score)| {
            if !penalise_paths {
                return (idx, score);
            }
            let path = &chunks[idx].file_path;
            let factor = *penalty_cache
                .entry(path.clone())
                .or_insert_with(|| file_path_penalty(path));
            (idx, score * factor)
        })
        .collect();

    // Step 2: sort descending by penalised score with stable tie-break by index.
    penalised.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));

    // Step 3: greedy pick with file-saturation decay and early-exit.
    let mut file_selected: HashMap<String, usize> = HashMap::new();
    let mut selected: Vec<(usize, f32)> = Vec::with_capacity(top_k.min(penalised.len()));
    let mut min_selected = f32::INFINITY;

    for &(idx, pen_score) in &penalised {
        if selected.len() >= top_k && pen_score <= min_selected {
            // Remaining (sorted-descending) cannot beat the current k-th best.
            break;
        }
        let path = chunks[idx].file_path.clone();
        let already = *file_selected.get(&path).unwrap_or(&0);
        let eff_score = if already >= FILE_SATURATION_THRESHOLD {
            let excess = already - FILE_SATURATION_THRESHOLD + 1;
            // Excess is bounded by the number of chunks in the front; the
            // saturation here is far below i32::MAX in practice. The clamp
            // is defence-in-depth.
            let excess_i32 = i32::try_from(excess).unwrap_or(i32::MAX);
            pen_score * FILE_SATURATION_DECAY.powi(excess_i32)
        } else {
            pen_score
        };
        selected.push((idx, eff_score));
        file_selected.insert(path, already + 1);

        if selected.len() >= top_k {
            // Refresh the min over the current selection.
            min_selected = selected
                .iter()
                .map(|(_, s)| *s)
                .fold(f32::INFINITY, f32::min);
        }
    }

    // Final descending sort on effective scores.
    selected.sort_unstable_by(|a, b| b.1.total_cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
    selected.truncate(top_k);
    selected
}

#[cfg(test)]
mod tests {
    use super::*;

    fn chunk_at(path: &str) -> CodeChunk {
        CodeChunk {
            file_path: path.to_string(),
            name: String::new(),
            kind: String::new(),
            start_line: 1,
            end_line: 1,
            content: String::new(),
            enriched_content: String::new(),
        }
    }

    /// `test:penalties-test-file-regex-14-langs` — every per-language
    /// test-file pattern from `penalties.py:8` matches a representative
    /// file path.
    #[test]
    fn penalties_test_file_regex_14_langs() {
        let cases: &[&str] = &[
            // Python
            "src/test_foo.py",
            "src/foo_test.py",
            // Go
            "pkg/foo_test.go",
            // Java
            "src/FooTest.java",
            "src/FooTests.java",
            // PHP
            "src/FooTest.php",
            // Ruby
            "spec/foo_spec.rb",
            "test/foo_test.rb",
            // JS/TS
            "src/foo.test.js",
            "src/foo.spec.ts",
            "src/foo.test.tsx",
            // Kotlin
            "src/FooTest.kt",
            "src/FooTests.kt",
            "src/FooSpec.kt",
            // Swift
            "src/FooTests.swift",
            "src/FooSpec.swift",
            // C#
            "src/FooTest.cs",
            "src/FooTests.cs",
            // C / C++
            "src/test_foo.cpp",
            "src/foo_test.cpp",
            "src/test_foo.c",
            "src/foo_test.c",
            // Scala
            "src/FooSpec.scala",
            "src/FooSuite.scala",
            "src/FooTest.scala",
            // Dart
            "src/foo_test.dart",
            "src/test_foo.dart",
            // Lua
            "src/foo_spec.lua",
            "src/foo_test.lua",
            "src/test_foo.lua",
            // Shared helpers
            "test/test_helper.rb",
            "test/test_helpers.go",
        ];
        for path in cases {
            assert!(
                test_file_re().is_match(path),
                "expected test_file_re to match {path:?}"
            );
            assert!(
                (file_path_penalty(path) - STRONG_PENALTY).abs() < 1e-6,
                "expected STRONG_PENALTY for {path:?}; got {}",
                file_path_penalty(path)
            );
        }
    }

    /// `test:penalties-compat-dir` — compat / legacy directories trigger
    /// the strong penalty.
    #[test]
    fn penalties_compat_dir() {
        for path in &["compat/foo.py", "src/_compat/bar.rs", "legacy/baz.go"] {
            assert!(
                compat_dir_re().is_match(path),
                "expected compat match for {path:?}"
            );
            assert!((file_path_penalty(path) - STRONG_PENALTY).abs() < 1e-6);
        }
    }

    /// `test:penalties-examples-dir` — example / docs-src directories
    /// trigger the strong penalty.
    #[test]
    fn penalties_examples_dir() {
        for path in &[
            "examples/foo.py",
            "_examples/bar.rs",
            "example/baz.go",
            "docs_src/x.md",
        ] {
            assert!(
                examples_dir_re().is_match(path),
                "expected examples match for {path:?}"
            );
            assert!((file_path_penalty(path) - STRONG_PENALTY).abs() < 1e-6);
        }
    }

    /// `test:penalties-init-py-reexport` — `__init__.py` carries the
    /// moderate re-export penalty (and Java's `package-info.java`).
    ///
    /// Java path deliberately avoids `/example/` since that segment
    /// triggers `examples_dir_re`'s strong penalty (we want to verify
    /// re-export in isolation here).
    #[test]
    fn penalties_init_py_reexport() {
        assert!((file_path_penalty("src/__init__.py") - MODERATE_PENALTY).abs() < 1e-6);
        assert!(
            (file_path_penalty("src/com/myorg/package-info.java") - MODERATE_PENALTY).abs() < 1e-6
        );
    }

    /// `test:penalties-dts-stub` — `.d.ts` files take the mild penalty.
    #[test]
    fn penalties_dts_stub() {
        assert!((file_path_penalty("src/foo.d.ts") - MILD_PENALTY).abs() < 1e-6);
    }

    /// Non-penalized paths return penalty 1.0.
    #[test]
    fn non_penalized_path_is_identity() {
        assert!((file_path_penalty("src/foo.rs") - 1.0).abs() < 1e-6);
        assert!((file_path_penalty("lib/bar.py") - 1.0).abs() < 1e-6);
    }

    /// `test:rerank-topk-saturation-decay` — a third chunk from the
    /// same file is penalised by 0.5^2.
    #[test]
    fn rerank_topk_saturation_decay() {
        let chunks = vec![
            chunk_at("src/foo.rs"),
            chunk_at("src/foo.rs"),
            chunk_at("src/foo.rs"),
            chunk_at("src/bar.rs"),
        ];
        // All four chunks have identical raw scores. The path penalty is 1.0
        // for both files. Greedy order picks chunks in their submitted
        // ordering (stable tie-break by index).
        let scores = vec![(0, 1.0_f32), (1, 1.0), (2, 1.0), (3, 1.0)];
        let got = rerank_topk(&scores, &chunks, 4, true);
        assert_eq!(got.len(), 4);

        let scored: HashMap<usize, f32> = got.iter().copied().collect();
        // Chunk 0 (first from foo.rs): score 1.0.
        // Chunk 3 (first from bar.rs): score 1.0.
        // Chunk 1 (second from foo.rs): already=1, excess=1, score *= 0.5.
        // Chunk 2 (third from foo.rs):  already=2, excess=2, score *= 0.25.
        assert!(
            (scored[&0] - 1.0).abs() < 1e-6,
            "scored[0] = {}",
            scored[&0]
        );
        assert!(
            (scored[&3] - 1.0).abs() < 1e-6,
            "scored[3] = {}",
            scored[&3]
        );
        assert!(
            (scored[&1] - 0.5).abs() < 1e-6,
            "scored[1] = {}",
            scored[&1]
        );
        assert!(
            (scored[&2] - 0.25).abs() < 1e-6,
            "scored[2] = {}",
            scored[&2]
        );
    }

    /// `test:rerank-topk-greedy-early-exit` — top_k smaller than the
    /// input cuts the iteration short while preserving ordering.
    #[test]
    fn rerank_topk_greedy_early_exit() {
        let chunks: Vec<CodeChunk> = (0..10).map(|i| chunk_at(&format!("src/f{i}.rs"))).collect();
        let scores: Vec<(usize, f32)> = (0..10).map(|i| (i, 10.0 - i as f32)).collect();
        let got = rerank_topk(&scores, &chunks, 3, true);
        assert_eq!(got.len(), 3);
        let indices: Vec<usize> = got.iter().map(|(i, _)| *i).collect();
        assert_eq!(indices, vec![0, 1, 2]);
    }

    /// `property:penalty-regex-parity-python` — non-test paths in
    /// production layouts (src/, lib/, crates/, etc.) carry no penalty,
    /// matching Python's behaviour exactly.
    #[test]
    fn property_penalty_regex_parity_python() {
        let production_paths = &[
            "src/foo.py",
            "lib/parser.rs",
            "crates/ripvec-core/src/encoder/semble/tokens.rs",
            "pkg/server/handler.go",
            "app/models/user.rb",
            "main.c",
            "include/foo.h",
        ];
        for path in production_paths {
            assert!(
                (file_path_penalty(path) - 1.0).abs() < 1e-6,
                "non-test path {path:?} should have penalty 1.0; got {}",
                file_path_penalty(path)
            );
        }
        // And test-shaped paths always strong-penalize.
        for path in &["test_foo.py", "src/__init__.py", "src/foo.d.ts"] {
            assert!(
                file_path_penalty(path) < 1.0,
                "test-shaped path {path:?} should be penalised; got 1.0"
            );
        }
    }

    /// `penalise_paths == false` bypasses path priors but still applies
    /// file-saturation decay.
    #[test]
    fn rerank_topk_no_path_penalties_still_decays() {
        let chunks = vec![chunk_at("test/test_foo.py"), chunk_at("test/test_foo.py")];
        let scores = vec![(0, 1.0_f32), (1, 1.0)];
        // With penalise_paths=false, the test_ file's path penalty (0.3)
        // does NOT apply; but the second chunk from the same file still
        // gets the saturation decay (×0.5).
        let got = rerank_topk(&scores, &chunks, 2, false);
        let scored: HashMap<usize, f32> = got.iter().copied().collect();
        assert!((scored[&0] - 1.0).abs() < 1e-6);
        assert!((scored[&1] - 0.5).abs() < 1e-6);
    }
}