arr/crates/arr-core/src/subs.rs

//! Ranking subtitle candidates against a media file (`DESIGN.md` §15).
//!
//! Pure and IO-free, alongside the release scoring in §5.5: the input is the
//! facts already known about the file plus a list of candidates as reported
//! by a provider, the output is those candidates ordered best first, each
//! carrying the verdict that put it there.
//!
//! The verdict vocabulary mirrors the release path — [`SubtitleVerdict`]
//! carries `Eligible` or `Rejected` a named [`SubtitleRule`], not a string
//! built for a log line — so §9.3's manual view needs no second concept for
//! subtitles. There is no `waived` bucket: nothing about a subtitle is worth
//! overriding by hand.

use crate::Source;

/// What is known about the file a subtitle candidate is ranked against.
#[derive(Clone, Copy, Debug, Default)]
pub struct SubtitleTarget<'a> {
    /// The file's own `moviehash`, when computed.
    pub moviehash: Option<&'a str>,
    /// The exact name of the release that produced the file.
    pub release_name: Option<&'a str>,
    /// The release group claimed by that release name.
    pub release_group: Option<&'a str>,
    /// The source tier claimed by that release name.
    pub source: Option<Source>,
}

/// One subtitle candidate, as reported by a provider.
#[derive(Clone, Copy, Debug, Default)]
pub struct SubtitleCandidate<'a> {
    /// Covers foreign-language lines and on-screen signs only; never
    /// satisfies a want (§15's "Forced and SDH").
    pub forced: bool,
    /// Complete and satisfies a want, but ranks below a plain subtitle for
    /// the same language.
    pub hearing_impaired: bool,
    pub moviehash: Option<&'a str>,
    pub release_name: Option<&'a str>,
    pub release_group: Option<&'a str>,
    pub source: Option<Source>,
    pub uploader_rating: f64,
    pub download_count: u64,
}

/// Why a subtitle candidate was rejected.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SubtitleRule {
    /// A forced track covers foreign-language lines only and cannot satisfy
    /// a want (§15).
    Forced,
}

impl SubtitleRule {
    /// The stable name shared by the API and UI, so one rule reads the same
    /// everywhere (mirrors [`crate::Rule::name`]).
    #[must_use]
    pub const fn name(&self) -> &'static str {
        match self {
            Self::Forced => "forced",
        }
    }
}

/// The verdict for one subtitle candidate.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum SubtitleVerdict {
    Eligible,
    Rejected(SubtitleRule),
}

/// One candidate paired with the verdict that placed it in the ranking.
#[derive(Clone, Copy, Debug)]
pub struct RankedSubtitle<'a> {
    /// Position of this candidate in the slice handed to [`rank`]. Ranking
    /// reorders, and a candidate carries no identity of its own, so this is
    /// how a caller maps a result back to the provider candidate it came
    /// from — and therefore to the id it must ask the provider to download.
    pub index: usize,
    pub candidate: SubtitleCandidate<'a>,
    pub verdict: SubtitleVerdict,
}

/// The ordering key for an eligible candidate, most significant field first.
/// Every field is a boolean win or an integer tiebreaker, so `Ord` alone
/// gives the ranking — no float, no partial order.
type TierKey = (bool, bool, bool, bool, i64, u64);

fn tier_key(candidate: &SubtitleCandidate<'_>, target: &SubtitleTarget<'_>) -> TierKey {
    let plain = !candidate.hearing_impaired;
    let moviehash_match = matches_ci(candidate.moviehash, target.moviehash);
    let exact_release_match = matches_ci(candidate.release_name, target.release_name);
    let group_or_source_match = matches_ci(candidate.release_group, target.release_group)
        || matches_source(candidate.source, target.source);
    // Scaled to hundredths so the tiebreaker is an integer, not a float.
    // Ratings sit in 0..10, nowhere near i64's range.
    #[allow(clippy::cast_possible_truncation)]
    let rating = (candidate.uploader_rating * 100.0).round() as i64;

    (
        plain,
        moviehash_match,
        exact_release_match,
        group_or_source_match,
        rating,
        candidate.download_count,
    )
}

fn matches_ci(candidate: Option<&str>, target: Option<&str>) -> bool {
    match (candidate, target) {
        (Some(candidate), Some(target)) => candidate.eq_ignore_ascii_case(target),
        _ => false,
    }
}

fn matches_source(candidate: Option<Source>, target: Option<Source>) -> bool {
    matches!((candidate, target), (Some(candidate), Some(target)) if candidate == target)
}

/// Rank candidates against a target, best first.
///
/// A forced candidate is rejected outright and sorts after every eligible
/// one, in the order it was given (§15: it can never satisfy a want). Among
/// eligible candidates, a plain subtitle always outranks a hearing-impaired
/// one for the same language; within that split, a `moviehash` match wins
/// outright, then an exact release-name match, then a shared release group
/// or source, then uploader rating and download count as tiebreakers.
#[must_use]
pub fn rank<'a>(
    target: &SubtitleTarget<'_>,
    candidates: &[SubtitleCandidate<'a>],
) -> Vec<RankedSubtitle<'a>> {
    let mut ranked: Vec<RankedSubtitle<'a>> = candidates
        .iter()
        .enumerate()
        .map(|(index, candidate)| RankedSubtitle {
            index,
            candidate: *candidate,
            verdict: if candidate.forced {
                SubtitleVerdict::Rejected(SubtitleRule::Forced)
            } else {
                SubtitleVerdict::Eligible
            },
        })
        .collect();

    ranked.sort_by(|a, b| match (a.verdict, b.verdict) {
        (SubtitleVerdict::Rejected(_), SubtitleVerdict::Rejected(_)) => std::cmp::Ordering::Equal,
        (SubtitleVerdict::Rejected(_), SubtitleVerdict::Eligible) => std::cmp::Ordering::Greater,
        (SubtitleVerdict::Eligible, SubtitleVerdict::Rejected(_)) => std::cmp::Ordering::Less,
        (SubtitleVerdict::Eligible, SubtitleVerdict::Eligible) => {
            tier_key(&b.candidate, target).cmp(&tier_key(&a.candidate, target))
        }
    });

    ranked
}

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

    fn target() -> SubtitleTarget<'static> {
        SubtitleTarget {
            moviehash: Some("abc123"),
            release_name: Some("Movie.2024.1080p.WEB-DL-GROUP"),
            release_group: Some("GROUP"),
            source: Some(Source::WebDl),
        }
    }

    fn plain() -> SubtitleCandidate<'static> {
        SubtitleCandidate {
            uploader_rating: 5.0,
            download_count: 100,
            ..Default::default()
        }
    }

    #[test]
    fn ranking_reports_where_each_candidate_came_from() {
        // Ranking reorders, and a candidate carries no id of its own, so the
        // index is the only way back to the provider candidate — and so to
        // the id the provider is asked to download.
        let weak = plain();
        let strong = SubtitleCandidate {
            moviehash: Some("abc123"),
            ..plain()
        };
        let ranked = rank(&target(), &[weak, strong]);

        assert_eq!(ranked[0].index, 1);
        assert_eq!(ranked[1].index, 0);
    }

    #[test]
    fn a_rejected_candidate_still_reports_its_index() {
        let forced = SubtitleCandidate {
            forced: true,
            ..plain()
        };
        let ranked = rank(&target(), &[forced, plain()]);

        assert_eq!(ranked[0].index, 1);
        assert_eq!(ranked[1].index, 0);
        assert_eq!(
            ranked[1].verdict,
            SubtitleVerdict::Rejected(SubtitleRule::Forced)
        );
    }

    #[test]
    fn a_moviehash_match_wins_outright_over_every_other_tier() {
        let hash_match = SubtitleCandidate {
            moviehash: Some("abc123"),
            uploader_rating: 0.0,
            download_count: 0,
            ..plain()
        };
        let everything_else = SubtitleCandidate {
            release_name: Some("Movie.2024.1080p.WEB-DL-GROUP"),
            release_group: Some("GROUP"),
            source: Some(Source::WebDl),
            uploader_rating: 10.0,
            download_count: 1_000_000,
            ..plain()
        };
        let ranked = rank(&target(), &[everything_else, hash_match]);
        assert!(matches!(ranked[0].candidate.moviehash, Some("abc123")));
        assert_eq!(ranked[0].verdict, SubtitleVerdict::Eligible);
    }

    #[test]
    fn an_exact_release_name_match_outranks_group_or_source_alone() {
        let exact_name = SubtitleCandidate {
            release_name: Some("Movie.2024.1080p.WEB-DL-GROUP"),
            ..plain()
        };
        let group_only = SubtitleCandidate {
            release_group: Some("GROUP"),
            uploader_rating: 10.0,
            download_count: 1_000_000,
            ..plain()
        };
        let ranked = rank(&target(), &[group_only, exact_name]);
        assert_eq!(
            ranked[0].candidate.release_name,
            Some("Movie.2024.1080p.WEB-DL-GROUP")
        );
    }

    #[test]
    fn same_source_alone_outranks_no_match_at_all() {
        let same_source = SubtitleCandidate {
            source: Some(Source::WebDl),
            ..plain()
        };
        let no_match = SubtitleCandidate {
            uploader_rating: 10.0,
            download_count: 1_000_000,
            ..plain()
        };
        let ranked = rank(&target(), &[no_match, same_source]);
        assert_eq!(ranked[0].candidate.source, Some(Source::WebDl));
    }

    #[test]
    #[allow(clippy::float_cmp)]
    fn rating_then_downloads_break_ties() {
        let low = SubtitleCandidate {
            uploader_rating: 3.0,
            download_count: 50,
            ..plain()
        };
        let high_rating = SubtitleCandidate {
            uploader_rating: 8.0,
            download_count: 10,
            ..plain()
        };
        let ranked = rank(&target(), &[low, high_rating]);
        assert_eq!(ranked[0].candidate.uploader_rating, 8.0);

        let same_rating_more_downloads = SubtitleCandidate {
            uploader_rating: 8.0,
            download_count: 9_000,
            ..plain()
        };
        let ranked = rank(&target(), &[high_rating, same_rating_more_downloads]);
        assert_eq!(ranked[0].candidate.download_count, 9_000);
    }

    #[test]
    fn a_forced_candidate_is_never_eligible() {
        let forced = SubtitleCandidate {
            forced: true,
            moviehash: Some("abc123"),
            ..plain()
        };
        let ranked = rank(&target(), &[forced]);
        assert_eq!(
            ranked[0].verdict,
            SubtitleVerdict::Rejected(SubtitleRule::Forced)
        );
        assert_eq!(SubtitleRule::Forced.name(), "forced");
    }

    #[test]
    fn an_sdh_candidate_ranks_below_any_plain_candidate() {
        let sdh_with_hash_match = SubtitleCandidate {
            hearing_impaired: true,
            moviehash: Some("abc123"),
            uploader_rating: 10.0,
            download_count: 1_000_000,
            ..plain()
        };
        let plain_with_nothing = SubtitleCandidate {
            uploader_rating: 0.0,
            download_count: 0,
            ..plain()
        };
        let ranked = rank(&target(), &[sdh_with_hash_match, plain_with_nothing]);
        assert!(!ranked[0].candidate.hearing_impaired);
        assert!(ranked[1].candidate.hearing_impaired);
    }

    #[test]
    fn rejected_candidates_sort_after_every_eligible_one_in_input_order() {
        let forced_a = SubtitleCandidate {
            forced: true,
            download_count: 1,
            ..plain()
        };
        let forced_b = SubtitleCandidate {
            forced: true,
            download_count: 2,
            ..plain()
        };
        let eligible = plain();
        let ranked = rank(&target(), &[forced_a, eligible, forced_b]);
        assert_eq!(ranked[0].verdict, SubtitleVerdict::Eligible);
        assert_eq!(ranked[1].candidate.download_count, 1);
        assert_eq!(ranked[2].candidate.download_count, 2);
    }
}