#[cfg(feature = "serde")]
use serde::{Deserialize, Serialize};
use khive_score::DeterministicScore;
use crate::error::FoldError;
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct SelectorInput<T> {
pub id: String,
pub content: T,
pub size: usize,
pub score: f32,
#[cfg_attr(feature = "serde", serde(default))]
pub category: Option<String>,
#[cfg_attr(feature = "serde", serde(default))]
pub information_gain: Option<f32>,
#[cfg_attr(feature = "serde", serde(default))]
pub rank_score: Option<f64>,
}
#[derive(Debug, Clone)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct SelectorOutput<T> {
pub selected: Vec<SelectorInput<T>>,
pub total_size: usize,
pub budget: usize,
}
#[derive(Debug, Clone, Default)]
#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
pub struct SelectorWeights {
pub category_weights: std::collections::BTreeMap<String, f32>,
pub min_score: f32,
pub diversity_bias: f32,
#[cfg_attr(feature = "serde", serde(default))]
pub epistemic_weight: f32,
}
pub trait Selector<T>: Send + Sync {
fn select(
&self,
inputs: Vec<SelectorInput<T>>,
budget: usize,
weights: &SelectorWeights,
) -> Result<SelectorOutput<T>, FoldError>;
}
#[derive(Debug, Clone, Copy, Default)]
pub struct GreedySelector;
#[inline]
fn rank_base<T>(item: &SelectorInput<T>) -> f64 {
item.rank_score.unwrap_or(item.score as f64)
}
#[inline]
fn pragmatic_plus_epistemic<T>(item: &SelectorInput<T>, epistemic_weight: f32) -> f64 {
let base = rank_base(item);
if epistemic_weight == 0.0 {
return base;
}
base + epistemic_weight as f64 * item.information_gain.unwrap_or(0.0) as f64
}
fn effective_score<T>(
item: &SelectorInput<T>,
counts: &std::collections::BTreeMap<String, usize>,
bias: f32,
epistemic_weight: f32,
) -> f64 {
let base = pragmatic_plus_epistemic(item, epistemic_weight);
if bias == 0.0 {
return base;
}
let count = item
.category
.as_ref()
.and_then(|c| counts.get(c))
.copied()
.unwrap_or(0);
base * (1.0 - bias as f64 * count as f64 / (count as f64 + 1.0))
}
fn validate_selector_weights(weights: &SelectorWeights) -> Result<(), FoldError> {
if !weights.min_score.is_finite() {
return Err(FoldError::InvalidInput(
"SelectorWeights.min_score must be finite".to_string(),
));
}
if !weights.diversity_bias.is_finite() {
return Err(FoldError::InvalidInput(
"SelectorWeights.diversity_bias must be finite".to_string(),
));
}
if !weights.epistemic_weight.is_finite() {
return Err(FoldError::InvalidInput(
"SelectorWeights.epistemic_weight must be finite".to_string(),
));
}
for (category, weight) in &weights.category_weights {
if !weight.is_finite() {
return Err(FoldError::InvalidInput(format!(
"SelectorWeights.category_weights['{category}'] must be finite"
)));
}
}
Ok(())
}
impl<T: Clone> Selector<T> for GreedySelector {
fn select(
&self,
mut inputs: Vec<SelectorInput<T>>,
budget: usize,
weights: &SelectorWeights,
) -> Result<SelectorOutput<T>, FoldError> {
validate_selector_weights(weights)?;
for input in &inputs {
if let Some(gain) = input.information_gain {
if !gain.is_finite() {
return Err(FoldError::InvalidInput(format!(
"information_gain for '{}' must be finite",
input.id
)));
}
}
if let Some(rank_score) = input.rank_score {
if !rank_score.is_finite() {
return Err(FoldError::InvalidInput(format!(
"rank_score for '{}' must be finite",
input.id
)));
}
}
}
inputs.retain(|i| i.score.is_finite() && i.score >= weights.min_score);
if !weights.category_weights.is_empty() {
for item in &mut inputs {
if let Some(ref cat) = item.category {
if let Some(&w) = weights.category_weights.get(cat.as_str()) {
let w = w.max(0.0);
item.score *= w;
if let Some(rank_score) = item.rank_score {
item.rank_score = Some(rank_score * w as f64);
}
}
}
}
inputs.retain(|i| i.score.is_finite() && i.score >= weights.min_score);
}
let ew = weights.epistemic_weight;
let mut ranked = Vec::with_capacity(inputs.len());
for input in inputs {
let effective = pragmatic_plus_epistemic(&input, ew);
if !effective.is_finite() {
return Err(FoldError::InvalidInput(format!(
"effective score for '{}' must be finite",
input.id
)));
}
let det_score = DeterministicScore::from_f64(effective);
ranked.push((input, det_score));
}
ranked.sort_by(|(a, a_det), (b, b_det)| {
b_det
.cmp(a_det)
.then_with(|| a.size.cmp(&b.size))
.then_with(|| a.id.cmp(&b.id))
});
let inputs: Vec<_> = ranked.into_iter().map(|(input, _)| input).collect();
let mut selected = Vec::new();
let mut total_size = 0usize;
if weights.diversity_bias == 0.0 {
for input in inputs {
if input.size <= budget.saturating_sub(total_size) {
total_size += input.size;
selected.push(input);
}
}
} else {
let mut remaining = inputs;
let mut category_counts: std::collections::BTreeMap<String, usize> =
std::collections::BTreeMap::new();
while !remaining.is_empty() && total_size < budget {
let mut candidates = Vec::with_capacity(remaining.len());
for (i, item) in remaining.iter().enumerate() {
if item.size > budget.saturating_sub(total_size) {
continue;
}
let eff = effective_score(item, &category_counts, weights.diversity_bias, ew);
if !eff.is_finite() {
return Err(FoldError::InvalidInput(format!(
"effective score for '{}' must be finite",
item.id
)));
}
candidates.push((i, DeterministicScore::from_f64(eff)));
}
let best_idx = candidates
.into_iter()
.max_by(|&(i, a_det), &(j, b_det)| {
a_det
.cmp(&b_det)
.then_with(|| remaining[j].size.cmp(&remaining[i].size))
.then_with(|| remaining[i].id.cmp(&remaining[j].id))
})
.map(|(i, _)| i);
match best_idx {
Some(idx) => {
let item = remaining.swap_remove(idx);
if let Some(ref cat) = item.category {
*category_counts.entry(cat.clone()).or_default() += 1;
}
total_size += item.size;
selected.push(item);
}
None => break,
}
}
}
Ok(SelectorOutput {
selected,
total_size,
budget,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
fn input(id: &str, size: usize, score: f32) -> SelectorInput<()> {
SelectorInput {
id: id.to_string(),
content: (),
size,
score,
category: None,
information_gain: None,
rank_score: None,
}
}
fn input_cat(id: &str, size: usize, score: f32, cat: &str) -> SelectorInput<()> {
SelectorInput {
id: id.to_string(),
content: (),
size,
score,
category: Some(cat.to_string()),
information_gain: None,
rank_score: None,
}
}
fn weights(min_score: f32) -> SelectorWeights {
SelectorWeights {
min_score,
..Default::default()
}
}
#[test]
fn empty_input() {
let inputs: Vec<SelectorInput<()>> = vec![];
let out = GreedySelector.select(inputs, 1000, &weights(0.0)).unwrap();
assert!(out.selected.is_empty());
assert_eq!(out.total_size, 0);
assert_eq!(out.budget, 1000);
}
#[test]
fn packs_highest_scores_first() {
let inputs = vec![
input("a", 100, 0.5),
input("b", 100, 0.9),
input("c", 100, 0.7),
];
let out = GreedySelector.select(inputs, 200, &weights(0.0)).unwrap();
assert_eq!(out.selected.len(), 2);
assert_eq!(out.selected[0].id, "b");
assert_eq!(out.selected[1].id, "c");
assert_eq!(out.total_size, 200);
}
#[test]
fn respects_budget() {
let inputs = vec![
input("a", 300, 0.9),
input("b", 300, 0.8),
input("c", 300, 0.7),
];
let out = GreedySelector.select(inputs, 500, &weights(0.0)).unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(out.selected[0].id, "a");
assert_eq!(out.total_size, 300);
}
#[test]
fn filters_below_min_score() {
let inputs = vec![
input("a", 10, 0.8),
input("b", 10, 0.1),
input("c", 10, 0.5),
];
let out = GreedySelector.select(inputs, 1000, &weights(0.3)).unwrap();
assert_eq!(out.selected.len(), 2);
assert_eq!(out.selected[0].id, "a");
assert_eq!(out.selected[1].id, "c");
}
#[test]
fn filters_nan_and_inf() {
let inputs = vec![
input("nan", 10, f32::NAN),
input("inf", 10, f32::INFINITY),
input("neg_inf", 10, f32::NEG_INFINITY),
input("ok", 10, 0.5),
];
let out = GreedySelector.select(inputs, 1000, &weights(0.0)).unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(out.selected[0].id, "ok");
}
#[test]
fn tie_break_size_ascending() {
let inputs = vec![input("big", 200, 0.5), input("small", 50, 0.5)];
let out = GreedySelector.select(inputs, 1000, &weights(0.0)).unwrap();
assert_eq!(out.selected[0].id, "small");
assert_eq!(out.selected[1].id, "big");
}
#[test]
fn tie_break_id_ascending() {
let inputs = vec![input("z", 100, 0.5), input("a", 100, 0.5)];
let out = GreedySelector.select(inputs, 1000, &weights(0.0)).unwrap();
assert_eq!(out.selected[0].id, "a");
assert_eq!(out.selected[1].id, "z");
}
#[test]
fn skips_oversized_items_takes_smaller() {
let inputs = vec![
input("huge", 900, 0.9),
input("small1", 40, 0.3),
input("small2", 40, 0.2),
];
let out = GreedySelector.select(inputs, 100, &weights(0.0)).unwrap();
assert_eq!(out.selected.len(), 2);
assert_eq!(out.selected[0].id, "small1");
assert_eq!(out.selected[1].id, "small2");
assert_eq!(out.total_size, 80);
}
#[test]
fn zero_budget() {
let inputs = vec![input("a", 1, 0.9)];
let out = GreedySelector.select(inputs, 0, &weights(0.0)).unwrap();
assert!(out.selected.is_empty());
}
#[test]
fn deterministic_across_input_order() {
let a = vec![
input("x", 50, 0.7),
input("y", 50, 0.7),
input("z", 50, 0.7),
];
let b = vec![
input("z", 50, 0.7),
input("x", 50, 0.7),
input("y", 50, 0.7),
];
let out_a = GreedySelector.select(a, 100, &weights(0.0)).unwrap();
let out_b = GreedySelector.select(b, 100, &weights(0.0)).unwrap();
let ids_a: Vec<&str> = out_a.selected.iter().map(|i| i.id.as_str()).collect();
let ids_b: Vec<&str> = out_b.selected.iter().map(|i| i.id.as_str()).collect();
assert_eq!(ids_a, ids_b);
assert_eq!(ids_a, vec!["x", "y"]);
}
#[test]
fn exact_budget_fit() {
let inputs = vec![input("a", 50, 0.9), input("b", 50, 0.8)];
let out = GreedySelector.select(inputs, 100, &weights(0.0)).unwrap();
assert_eq!(out.selected.len(), 2);
assert_eq!(out.total_size, 100);
}
#[test]
fn category_weights_boost_preferred_category() {
let inputs = vec![
input_cat("a", 100, 0.9, "low"),
input_cat("b", 100, 0.5, "high"),
];
let w = SelectorWeights {
category_weights: [("high".to_string(), 2.0f32), ("low".to_string(), 1.0f32)]
.into_iter()
.collect(),
..Default::default()
};
let out = GreedySelector.select(inputs, 100, &w).unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(out.selected[0].id, "b");
}
#[test]
fn category_weights_can_push_below_min_score() {
let inputs = vec![
input_cat("a", 10, 0.4, "bad"),
input_cat("b", 10, 0.8, "good"),
];
let w = SelectorWeights {
min_score: 0.3,
category_weights: [("bad".to_string(), 0.5f32)].into_iter().collect(),
..Default::default()
};
let out = GreedySelector.select(inputs, 1000, &w).unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(out.selected[0].id, "b");
}
#[test]
fn diversity_bias_zero_identical_to_greedy() {
let make = || {
vec![
input_cat("a", 100, 0.9, "x"),
input_cat("b", 100, 0.8, "x"),
input_cat("c", 100, 0.7, "y"),
]
};
let w_greedy = SelectorWeights {
..Default::default()
};
let w_bias0 = SelectorWeights {
diversity_bias: 0.0,
..Default::default()
};
let out_g = GreedySelector.select(make(), 200, &w_greedy).unwrap();
let out_b = GreedySelector.select(make(), 200, &w_bias0).unwrap();
let ids_g: Vec<&str> = out_g.selected.iter().map(|i| i.id.as_str()).collect();
let ids_b: Vec<&str> = out_b.selected.iter().map(|i| i.id.as_str()).collect();
assert_eq!(ids_g, ids_b);
}
#[test]
fn diversity_bias_prefers_different_categories() {
let inputs = vec![
input_cat("a", 100, 0.9, "x"),
input_cat("b", 100, 0.8, "x"),
input_cat("c", 100, 0.7, "y"),
];
let w = SelectorWeights {
diversity_bias: 1.0,
..Default::default()
};
let out = GreedySelector.select(inputs, 200, &w).unwrap();
assert_eq!(out.selected.len(), 2);
let ids: Vec<&str> = out.selected.iter().map(|i| i.id.as_str()).collect();
assert!(ids.contains(&"a"), "a should always be selected");
assert!(
ids.contains(&"c"),
"c should be preferred over b due to diversity"
);
}
#[test]
fn no_overflow_near_usize_max() {
let large = usize::MAX - 1;
let inputs = vec![
SelectorInput {
id: "a".to_string(),
content: (),
size: large,
score: 0.9,
category: None,
information_gain: None,
rank_score: None,
},
SelectorInput {
id: "b".to_string(),
content: (),
size: 10,
score: 0.8,
category: None,
information_gain: None,
rank_score: None,
},
];
let out = GreedySelector.select(inputs, 100, &weights(0.0)).unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(out.selected[0].id, "b");
}
#[test]
fn diversity_bias_no_categories_unaffected() {
let inputs = vec![
input("a", 100, 0.9),
input("b", 100, 0.8),
input("c", 100, 0.7),
];
let w = SelectorWeights {
diversity_bias: 1.0,
..Default::default()
};
let out = GreedySelector.select(inputs, 200, &w).unwrap();
assert_eq!(out.selected.len(), 2);
assert_eq!(out.selected[0].id, "a");
assert_eq!(out.selected[1].id, "b");
}
fn input_with_gain(id: &str, size: usize, score: f32, gain: f32) -> SelectorInput<()> {
SelectorInput {
id: id.to_string(),
content: (),
size,
score,
category: None,
information_gain: Some(gain),
rank_score: None,
}
}
#[test]
fn epistemic_weight_zero_preserves_behavior() {
let make = || {
vec![
input_with_gain("a", 100, 0.9, 10.0),
input_with_gain("b", 100, 0.8, 0.0),
input_with_gain("c", 100, 0.7, 5.0),
]
};
let w_default = SelectorWeights {
..Default::default()
};
let w_zero = SelectorWeights {
epistemic_weight: 0.0,
..Default::default()
};
let out_d = GreedySelector.select(make(), 200, &w_default).unwrap();
let out_z = GreedySelector.select(make(), 200, &w_zero).unwrap();
let ids_d: Vec<&str> = out_d.selected.iter().map(|i| i.id.as_str()).collect();
let ids_z: Vec<&str> = out_z.selected.iter().map(|i| i.id.as_str()).collect();
assert_eq!(ids_d, ids_z);
assert_eq!(ids_d, vec!["a", "b"]);
}
#[test]
fn epistemic_weight_positive_reorders_by_gain() {
let inputs = vec![
input_with_gain("a", 100, 0.5, 10.0),
input_with_gain("b", 100, 0.9, 0.0),
];
let w = SelectorWeights {
epistemic_weight: 1.0,
..Default::default()
};
let out = GreedySelector.select(inputs, 100, &w).unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(out.selected[0].id, "a");
}
#[test]
fn information_gain_none_equivalent_to_zero() {
let with_none = vec![
input("a", 100, 0.9), input("b", 100, 0.8),
];
let with_zero = vec![
input_with_gain("a", 100, 0.9, 0.0),
input_with_gain("b", 100, 0.8, 0.0),
];
let w = SelectorWeights {
epistemic_weight: 1.0,
..Default::default()
};
let out_none = GreedySelector.select(with_none, 200, &w).unwrap();
let out_zero = GreedySelector.select(with_zero, 200, &w).unwrap();
let ids_none: Vec<&str> = out_none.selected.iter().map(|i| i.id.as_str()).collect();
let ids_zero: Vec<&str> = out_zero.selected.iter().map(|i| i.id.as_str()).collect();
assert_eq!(ids_none, ids_zero);
}
#[test]
fn epistemic_weight_works_with_diversity_bias() {
let inputs = vec![
{
let mut i = input_with_gain("a", 100, 0.5, 10.0);
i.category = Some("x".to_string());
i
},
{
let mut i = input_with_gain("b", 100, 0.8, 0.0);
i.category = Some("x".to_string());
i
},
{
let mut i = input_with_gain("c", 100, 0.3, 0.0);
i.category = Some("y".to_string());
i
},
];
let w = SelectorWeights {
epistemic_weight: 1.0,
diversity_bias: 0.5,
..Default::default()
};
let out = GreedySelector.select(inputs, 200, &w).unwrap();
assert_eq!(out.selected.len(), 2);
assert_eq!(out.selected[0].id, "a");
assert_eq!(out.selected[1].id, "b");
}
#[test]
fn greedy_selector_rejects_nan_information_gain() {
let inputs = vec![
input_with_gain("a", 100, 0.1, f32::NAN),
input_with_gain("b", 100, 0.9, 0.0),
];
let w = SelectorWeights {
epistemic_weight: 1.0,
..Default::default()
};
let err = GreedySelector.select(inputs, 100, &w).unwrap_err();
assert!(matches!(err, FoldError::InvalidInput(_)));
}
#[test]
fn greedy_selector_rejects_non_finite_epistemic_weight() {
for bad in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
let inputs = vec![input("a", 100, 0.5)];
let w = SelectorWeights {
epistemic_weight: bad,
..Default::default()
};
let err = GreedySelector.select(inputs, 100, &w).unwrap_err();
assert!(matches!(err, FoldError::InvalidInput(_)));
}
}
#[test]
fn greedy_selector_rejects_non_finite_diversity_bias() {
for bad in [f32::NAN, f32::INFINITY, f32::NEG_INFINITY] {
let inputs = vec![input("a", 100, 0.5)];
let w = SelectorWeights {
diversity_bias: bad,
..Default::default()
};
let err = GreedySelector.select(inputs, 100, &w).unwrap_err();
assert!(matches!(err, FoldError::InvalidInput(_)));
}
}
#[test]
fn greedy_selector_rejects_non_finite_category_weight() {
let inputs = vec![input_cat("a", 100, 0.5, "x")];
let w = SelectorWeights {
category_weights: [("x".to_string(), f32::NAN)].into_iter().collect(),
..Default::default()
};
let err = GreedySelector.select(inputs, 100, &w).unwrap_err();
assert!(matches!(err, FoldError::InvalidInput(_)));
}
#[test]
fn greedy_selector_handles_extreme_f32_products_without_overflow() {
let inputs = vec![input_with_gain("a", 100, f32::MAX, f32::MAX)];
let w = SelectorWeights {
epistemic_weight: f32::MAX,
..Default::default()
};
let out = GreedySelector.select(inputs, 100, &w).unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(out.selected[0].id, "a");
}
#[test]
fn rejects_non_finite_rank_score() {
for bad in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
let mut item = input("a", 100, 0.5);
item.rank_score = Some(bad);
let err = GreedySelector
.select(vec![item], 100, &weights(0.0))
.unwrap_err();
assert!(matches!(err, FoldError::InvalidInput(_)));
}
}
#[test]
fn rank_score_saturates_at_deterministic_score_max_without_panic() {
let mut big = input("big", 200, 0.0);
big.rank_score = Some(f64::MAX);
let mut small = input("small", 50, 0.0);
small.rank_score = Some(f64::MAX / 2.0);
let out = GreedySelector
.select(vec![big, small], 1000, &weights(0.0))
.unwrap();
assert_eq!(out.selected.len(), 2);
assert_eq!(out.selected[0].id, "small");
assert_eq!(out.selected[1].id, "big");
}
#[test]
fn rank_score_distinguishes_values_within_f32_ulp_of_one() {
let a_score = 1.0_f32;
let b_score = 1.0_f32; assert_eq!(a_score.to_bits(), b_score.to_bits());
let mut a = input("a", 100, a_score);
a.rank_score = Some(1.0);
let mut b = input("b", 100, b_score);
b.rank_score = Some(1.000_000_04);
let out = GreedySelector
.select(vec![a, b], 100, &weights(0.0))
.unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(
out.selected[0].id, "b",
"higher rank_score must win despite tied f32 score"
);
}
#[test]
fn category_weights_reorder_candidates_when_rank_score_present() {
let mut a = input_cat("a", 100, 0.9, "low");
a.rank_score = Some(0.9);
let mut b = input_cat("b", 100, 0.5, "high");
b.rank_score = Some(0.5);
let w = SelectorWeights {
category_weights: [("high".to_string(), 2.0f32), ("low".to_string(), 1.0f32)]
.into_iter()
.collect(),
..Default::default()
};
let out = GreedySelector.select(vec![a, b], 100, &w).unwrap();
assert_eq!(out.selected.len(), 1);
assert_eq!(
out.selected[0].id, "b",
"category weight must still reorder candidates when rank_score is present"
);
}
#[test]
fn rank_score_zero_ties_break_deterministically() {
let mut a = input("z", 100, 0.0);
a.rank_score = Some(0.0);
let mut b = input("a", 100, 0.0);
b.rank_score = Some(0.0);
let out = GreedySelector
.select(vec![a, b], 1000, &weights(0.0))
.unwrap();
assert_eq!(out.selected.len(), 2);
assert_eq!(out.selected[0].id, "a");
assert_eq!(out.selected[1].id, "z");
}
}