x-pipe-rs 0.1.0

//! Candidate scoring: assigning relevance scores.
//!
//! A [`Scorer`] assigns a [`Score`] to each candidate item.
//! Multiple scorers can be combined additively via
//! [`combine_scorers`].

use std::ops::Add;
use std::sync::Arc;

use comp_cat_rs::effect::io::Io;

use crate::score::{Score, ScoredCandidate};
use crate::stage::Stage;

/// Assigns a relevance [`Score`] to a candidate item.
pub trait Scorer<I, E> {
    /// Score a single candidate.
    fn score(&self, item: &I) -> Io<E, Score>;
}

/// Convert a [`Scorer`] into a [`Stage`] that wraps items
/// in [`ScoredCandidate`].
pub fn scorer_stage<S, I, E>(scorer: S) -> Stage<E, Vec<I>, Vec<ScoredCandidate<I>>>
where
    S: Scorer<I, E> + Send + Sync + 'static,
    I: Send + 'static,
    E: Send + 'static,
{
    let scorer = Arc::new(scorer);
    Stage::new(move |items: Vec<I>| {
        items.into_iter().fold(
            Io::pure(Vec::new()),
            |acc_io, item| {
                let s = Arc::clone(&scorer);
                acc_io.flat_map(move |acc| {
                    s.score(&item).map(move |score| {
                        let sc = ScoredCandidate::new(item, score);
                        acc.into_iter().chain(std::iter::once(sc)).collect()
                    })
                })
            },
        )
    })
}

/// Combine two scorers additively.
///
/// The combined score is `s1.score(item).add(s2.score(item))`.
pub struct CombinedScorer<S1, S2> {
    first: S1,
    second: S2,
}

impl<S1, S2> CombinedScorer<S1, S2> {
    /// Create a combined scorer from two scorers.
    pub fn new(first: S1, second: S2) -> Self {
        Self { first, second }
    }
}

impl<S1, S2, I, E> Scorer<I, E> for CombinedScorer<S1, S2>
where
    S1: Scorer<I, E> + Send + Sync + 'static,
    S2: Scorer<I, E> + Send + Sync + 'static,
    I: Send + Sync + 'static,
    E: Send + 'static,
{
    fn score(&self, item: &I) -> Io<E, Score> {
        // We need both scores, but Io::zip requires ownership.
        // Since we have &self, we produce two Io values and zip them.
        let io1 = self.first.score(item);
        let io2 = self.second.score(item);
        io1.zip(io2).map(|(a, b)| a.add(b))
    }
}