use crate::CupelError;
use crate::model::{ContextBudget, ContextItem, ScoredItem};
use crate::slicer::Slicer;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
pub struct GreedySlice;
impl Slicer for GreedySlice {
fn slice(
&self,
sorted: &[ScoredItem],
budget: &ContextBudget,
) -> Result<Vec<ContextItem>, CupelError> {
if sorted.is_empty() || budget.target_tokens() <= 0 {
return Ok(Vec::new());
}
let mut densities: Vec<(f64, usize)> = sorted
.iter()
.enumerate()
.map(|(i, si)| {
let tokens = si.item.tokens();
let density = if tokens == 0 {
f64::MAX
} else {
si.score / tokens as f64
};
(density, i)
})
.collect();
densities.sort_by(|a, b| b.0.total_cmp(&a.0).then_with(|| a.1.cmp(&b.1)));
let mut remaining = budget.target_tokens();
let mut selected = Vec::new();
for &(_, orig_idx) in &densities {
let item = &sorted[orig_idx].item;
let tokens = item.tokens();
if tokens == 0 {
selected.push(item.clone());
} else if tokens <= remaining {
selected.push(item.clone());
remaining -= tokens;
}
}
Ok(selected)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ContextItemBuilder;
use std::collections::HashMap;
fn budget(target: i64) -> ContextBudget {
ContextBudget::new(target * 2, target, 0, HashMap::new(), 0.0).unwrap()
}
fn item(content: &str, tokens: i64) -> ContextItem {
ContextItemBuilder::new(content, tokens).build().unwrap()
}
fn scored(item: ContextItem, score: f64) -> ScoredItem {
ScoredItem { item, score }
}
#[test]
fn equal_density_preserves_input_order() {
let a = item("A", 100);
let b = item("B", 50);
let c = item("C", 200);
let items = vec![
scored(a, 1.0), scored(b, 0.5), scored(c, 2.0), ];
let result = GreedySlice.slice(&items, &budget(400)).unwrap();
assert_eq!(result.len(), 3);
assert_eq!(result[0].content(), "A");
assert_eq!(result[1].content(), "B");
assert_eq!(result[2].content(), "C");
}
#[test]
fn zero_token_items_tied_preserves_input_order() {
let x = item("X", 0);
let y = item("Y", 0);
let z = item("Z", 0);
let items = vec![scored(x, 0.3), scored(y, 0.9), scored(z, 0.1)];
let result = GreedySlice.slice(&items, &budget(50)).unwrap();
assert_eq!(result.len(), 3);
assert_eq!(result[0].content(), "X");
assert_eq!(result[1].content(), "Y");
assert_eq!(result[2].content(), "Z");
}
#[test]
fn equal_density_budget_constraint_drops_last_in_input_order() {
let a = item("first", 30);
let b = item("second", 30);
let c = item("third", 30);
let items = vec![
scored(a, 0.6), scored(b, 0.6), scored(c, 0.6), ];
let result = GreedySlice.slice(&items, &budget(60)).unwrap();
assert_eq!(result.len(), 2);
assert_eq!(result[0].content(), "first");
assert_eq!(result[1].content(), "second");
}
#[test]
fn deterministic_tiebreak_is_idempotent() {
for _ in 0..10 {
let a = item("A", 20);
let b = item("B", 20);
let c = item("C", 20);
let items = vec![scored(a, 0.4), scored(b, 0.4), scored(c, 0.4)];
let result = GreedySlice.slice(&items, &budget(40)).unwrap();
assert_eq!(result.len(), 2);
assert_eq!(result[0].content(), "A");
assert_eq!(result[1].content(), "B");
}
}
}