cel_brief/

budget.rs

//! [`PruneStrategy`] + [`apply_budget`].
//!
//! The builder calls [`apply_budget`] once per
//! turn against the fanned-out, tokenized contributions; it returns a `(kept,
//! dropped)` split that honours both the global [`crate::types::TokenBudget`]
//! and any per-priority floors set on it.
//!
//! The crate ships two strategies:
//! - [`PruneStrategy::ImportanceFirst`] — default; drops the lowest
//!   `(priority, importance)` items first.
//! - [`PruneStrategy::RoundRobin`] — sweeps lowest priority across all
//!   sources before touching the next-higher bucket. Useful when several
//!   sources of equal priority should suffer pruning symmetrically rather
//!   than one source losing all of its contributions.
//!
//! Custom strategies are out of scope for Phase 2; the enum stays
//! non-exhaustive so we can add variants without a breaking change.

use std::collections::HashMap;

use serde::{Deserialize, Serialize};

use crate::receipt::{DropReason, DroppedContribution};
use crate::source::Contribution;
use crate::types::{Priority, SourceId, TokenBudget};

/// How [`apply_budget`] orders dropped contributions when the assembled set
/// exceeds the prompt budget.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum PruneStrategy {
    /// Default. Sort dropped candidates by priority ascending, then
    /// importance ascending — lowest-importance items in the lowest priority
    /// bucket are dropped first.
    #[default]
    ImportanceFirst,
    /// Sweep one source at a time within each priority bucket, dropping the
    /// least-important contribution from each source in turn. Useful when
    /// several sources share a priority and you don't want one source to
    /// keep all of its content while another loses everything.
    RoundRobin,
}

/// One contribution as the budget layer sees it: the original
/// [`Contribution`], the source that produced it, the ground-truth
/// `actual_tokens` count from the [`crate::tokenizer::Tokenizer`], and the
/// source's priority at admission time.
///
/// `actual_tokens` is the value the budget uses for ceiling enforcement —
/// `Contribution::estimated_tokens` is a hint only.
#[derive(Debug, Clone)]
pub struct WeightedContribution {
    /// Source that produced the contribution.
    pub source: SourceId,
    /// Source's priority at admission time.
    pub priority: Priority,
    /// Ground-truth token count from the active tokenizer.
    pub actual_tokens: usize,
    /// The contribution itself.
    pub contribution: Contribution,
    /// Stable index inside the source's contribution list. Used by
    /// `apply_budget` to preserve source-defined order within priority.
    pub source_index: usize,
}

/// Apply the configured prune strategy to bring `contributions` under
/// `budget.prompt_budget()`, honouring any per-priority floors.
///
/// Returns `(kept, dropped)`:
/// - `kept` is the post-prune list **in input order** (so callers can rely
///   on source-defined ordering surviving the budget pass).
/// - `dropped` is the list of [`DroppedContribution`] records for the
///   receipt.
///
/// Behaviour:
/// - Items in [`Priority::Critical`] are never dropped (Critical is the
///   "must-keep" bucket). If Critical alone exceeds the
///   budget, every non-Critical item is dropped and `kept` may still be
///   over-budget — the builder turns that into a
///   [`crate::error::BriefError::BudgetUnsatisfiable`].
/// - Floors in `budget.floor_per_priority` are honoured **best-effort**:
///   when over budget, the pruner will break a lower priority's floor
///   before dropping a higher-priority item, matching the docstring on
///   [`TokenBudget::floor_per_priority`] ("a higher priority can borrow
///   from a lower floor when over budget").
///
/// Algorithm:
/// 1. Iterate non-Critical priority buckets from [`Priority::Low`] →
///    [`Priority::High`].
/// 2. Within each bucket, drop items in strategy-defined order, honouring
///    the bucket's floor (items that would breach the floor are skipped
///    and marked as floored).
/// 3. Before moving up to the next-higher bucket, run a borrow pass over
///    every floored item at this or lower priorities; drop until under
///    budget or out of floored candidates.
/// 4. Critical bucket is left untouched. The caller turns any residual
///    over-budget total into [`crate::error::BriefError::BudgetUnsatisfiable`].
pub fn apply_budget(
    contributions: Vec<WeightedContribution>,
    budget: &TokenBudget,
    strategy: PruneStrategy,
) -> (Vec<WeightedContribution>, Vec<DroppedContribution>) {
    let prompt_budget = budget.prompt_budget();
    let total_tokens: usize = contributions.iter().map(|c| c.actual_tokens).sum();

    // Fast path: everything fits.
    if total_tokens <= prompt_budget {
        return (contributions, Vec::new());
    }

    let indexed: Vec<(usize, WeightedContribution)> =
        contributions.into_iter().enumerate().collect();

    let mut dropped_records: Vec<DroppedContribution> = Vec::new();
    let mut dropped_set: std::collections::HashSet<usize> = std::collections::HashSet::new();
    let mut bucket_tokens = bucket_token_map(&indexed);
    let mut running_total = total_tokens;
    let floors = &budget.floor_per_priority;

    // Process buckets low → high so lower-priority items drop first.
    // `Priority::ALL` is `[Low, Normal, High, Critical]`; we exclude
    // Critical.
    let droppable_priorities = [Priority::Low, Priority::Normal, Priority::High];

    // Track which items got skipped because of a floor — eligible for
    // the borrow pass below.
    let mut floored_skips: Vec<usize> = Vec::new();

    for priority in droppable_priorities {
        if running_total <= prompt_budget {
            break;
        }

        // Compute drop order restricted to this bucket.
        let bucket_indices = drop_order_for_priority(&indexed, priority, strategy);

        for original_idx in bucket_indices {
            if running_total <= prompt_budget {
                break;
            }
            let Some((_, entry)) = indexed.iter().find(|(i, _)| *i == original_idx) else {
                continue;
            };
            let tokens = entry.actual_tokens;
            let source = entry.source.clone();

            // Honour the floor for this bucket (skip if breaching).
            if let Some(&floor) = floors.get(&priority) {
                let bucket = bucket_tokens.get(&priority).copied().unwrap_or(0);
                if bucket.saturating_sub(tokens) < floor {
                    floored_skips.push(original_idx);
                    continue;
                }
            }

            dropped_set.insert(original_idx);
            dropped_records.push(DroppedContribution {
                source,
                reason: DropReason::OverBudget,
                tokens,
            });
            running_total -= tokens;
            if let Some(bucket) = bucket_tokens.get_mut(&priority) {
                *bucket = bucket.saturating_sub(tokens);
            }
        }

        // Before moving to the next-higher bucket, see if we can satisfy
        // the budget by borrowing from floored items at this or lower
        // priorities. This implements the "higher priority can borrow
        // from a lower floor" guarantee — we'd rather break a low-prio
        // floor than drop a high-prio item.
        if running_total > prompt_budget {
            let borrow_indices: Vec<usize> = floored_skips
                .iter()
                .copied()
                .filter(|i| !dropped_set.contains(i))
                .collect();

            for original_idx in borrow_indices {
                if running_total <= prompt_budget {
                    break;
                }
                let Some((_, entry)) = indexed.iter().find(|(i, _)| *i == original_idx) else {
                    continue;
                };
                let tokens = entry.actual_tokens;
                let source = entry.source.clone();
                let p = entry.priority;

                dropped_set.insert(original_idx);
                dropped_records.push(DroppedContribution {
                    source,
                    reason: DropReason::OverBudget,
                    tokens,
                });
                running_total -= tokens;
                if let Some(bucket) = bucket_tokens.get_mut(&p) {
                    *bucket = bucket.saturating_sub(tokens);
                }
            }
        }
    }

    let kept: Vec<WeightedContribution> = indexed
        .into_iter()
        .filter_map(|(i, c)| {
            if dropped_set.contains(&i) {
                None
            } else {
                Some(c)
            }
        })
        .collect();

    (kept, dropped_records)
}

/// Returns the list of original indices for items at exactly `priority`,
/// in the order they should be dropped according to `strategy`.
///
/// Critical-priority items are filtered out by the caller — calling this
/// with `Priority::Critical` returns an empty vector.
fn drop_order_for_priority(
    indexed: &[(usize, WeightedContribution)],
    priority: Priority,
    strategy: PruneStrategy,
) -> Vec<usize> {
    if priority == Priority::Critical {
        return Vec::new();
    }
    let mut candidates: Vec<&(usize, WeightedContribution)> = indexed
        .iter()
        .filter(|(_, c)| c.priority == priority)
        .collect();

    match strategy {
        PruneStrategy::ImportanceFirst => {
            // Sort by (importance asc, source_index desc).
            // NaN importance sorts to the front (drop first) — defensive.
            candidates.sort_by(|a, b| {
                a.1.contribution
                    .importance
                    .partial_cmp(&b.1.contribution.importance)
                    .unwrap_or(std::cmp::Ordering::Less)
                    .then_with(|| b.1.source_index.cmp(&a.1.source_index))
            });
            candidates.into_iter().map(|(i, _)| *i).collect()
        }
        PruneStrategy::RoundRobin => {
            // Group by source within this priority.
            let mut by_source: HashMap<SourceId, Vec<&(usize, WeightedContribution)>> =
                HashMap::new();
            for entry in &candidates {
                by_source
                    .entry(entry.1.source.clone())
                    .or_default()
                    .push(entry);
            }
            // Within each source: sort by importance asc, source_index
            // desc (later entries in a source drop before earlier ones
            // for round-robin fairness — `pop()` below pulls from the
            // end).
            for entries in by_source.values_mut() {
                entries.sort_by(|a, b| {
                    b.1.contribution
                        .importance
                        .partial_cmp(&a.1.contribution.importance)
                        .unwrap_or(std::cmp::Ordering::Less)
                        .then_with(|| a.1.source_index.cmp(&b.1.source_index))
                });
            }
            // Stable order for sources — sort by SourceId so test output
            // is deterministic.
            let mut source_ids: Vec<SourceId> = by_source.keys().cloned().collect();
            source_ids.sort();

            let mut ordered: Vec<usize> = Vec::with_capacity(candidates.len());
            let mut exhausted_sources = 0;
            while exhausted_sources < source_ids.len() {
                exhausted_sources = 0;
                for sid in &source_ids {
                    if let Some(entries) = by_source.get_mut(sid) {
                        if let Some(entry) = entries.pop() {
                            ordered.push(entry.0);
                        } else {
                            exhausted_sources += 1;
                        }
                    } else {
                        exhausted_sources += 1;
                    }
                }
            }
            ordered
        }
    }
}

fn bucket_token_map(indexed: &[(usize, WeightedContribution)]) -> HashMap<Priority, usize> {
    let mut out: HashMap<Priority, usize> = HashMap::new();
    for (_, entry) in indexed {
        *out.entry(entry.priority).or_insert(0) += entry.actual_tokens;
    }
    out
}

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

    use crate::source::{Contribution, ContributionContent};
    use crate::types::{Role, SourceId};

    fn wc(
        source: &str,
        priority: Priority,
        tokens: usize,
        importance: f32,
        source_index: usize,
    ) -> WeightedContribution {
        WeightedContribution {
            source: SourceId::new(source),
            priority,
            actual_tokens: tokens,
            source_index,
            contribution: Contribution {
                content: ContributionContent::Text {
                    role: Role::User,
                    content: "x".repeat(tokens * 4),
                },
                estimated_tokens: tokens,
                importance,
                redactable: true,
                tags: Vec::new(),
            },
        }
    }

    #[test]
    fn under_budget_keeps_everything() {
        let items = vec![
            wc("a", Priority::Normal, 10, 0.5, 0),
            wc("b", Priority::Normal, 20, 0.3, 1),
        ];
        let budget = TokenBudget::new(1000, 0);
        let (kept, dropped) = apply_budget(items, &budget, PruneStrategy::default());
        assert_eq!(kept.len(), 2);
        assert!(dropped.is_empty());
    }

    #[test]
    fn over_budget_drops_lowest_importance_first() {
        let items = vec![
            wc("a", Priority::Normal, 100, 0.9, 0),
            wc("b", Priority::Normal, 100, 0.1, 1),
            wc("c", Priority::Normal, 100, 0.5, 2),
        ];
        // Budget for ~200 tokens of content.
        let budget = TokenBudget::new(200, 0);
        let (kept, dropped) = apply_budget(items, &budget, PruneStrategy::ImportanceFirst);
        assert_eq!(kept.len(), 2);
        assert_eq!(dropped.len(), 1);
        // The 0.1 importance entry should be the dropped one.
        let kept_sources: Vec<&str> = kept.iter().map(|c| c.source.as_str()).collect();
        assert!(kept_sources.contains(&"a"));
        assert!(kept_sources.contains(&"c"));
        assert_eq!(dropped[0].source.as_str(), "b");
    }

    #[test]
    fn critical_priority_is_never_dropped() {
        let items = vec![
            wc("crit", Priority::Critical, 500, 0.1, 0),
            wc("low", Priority::Low, 100, 0.9, 1),
        ];
        let budget = TokenBudget::new(200, 0);
        let (kept, dropped) = apply_budget(items, &budget, PruneStrategy::default());
        // Critical kept; low dropped because budget is 200 and crit alone
        // is 500 (over budget but critical can't be touched).
        assert_eq!(kept.len(), 1);
        assert_eq!(kept[0].source.as_str(), "crit");
        assert_eq!(dropped.len(), 1);
        assert_eq!(dropped[0].source.as_str(), "low");
    }

    #[test]
    fn priority_floor_protects_a_bucket() {
        // Floor of 100 on Normal means the pruner stops removing Normal
        // contributions once that bucket hits 100 tokens.
        let items = vec![
            wc("h", Priority::High, 200, 0.9, 0),
            wc("n1", Priority::Normal, 80, 0.5, 1),
            wc("n2", Priority::Normal, 80, 0.3, 2),
            wc("l", Priority::Low, 100, 0.7, 3),
        ];
        // Budget for 200 tokens. Need to drop 260 tokens.
        let budget = TokenBudget::new(200, 0).with_floor(Priority::Normal, 100);
        let (kept, dropped) = apply_budget(items, &budget, PruneStrategy::ImportanceFirst);
        // After best-effort and borrow pass: total is High(200) + maybe
        // some normal/low. Floors are best-effort, so the borrow pass
        // can take normal below the floor when over budget. We don't
        // assert exact keep set — we assert the floor was honoured on
        // the best-effort pass by checking the drop order included low
        // before normal.
        let kept_sources: Vec<&str> = kept.iter().map(|c| c.source.as_str()).collect();
        // High must be kept.
        assert!(kept_sources.contains(&"h"));
        // Low should have been dropped (lower priority).
        let dropped_sources: Vec<&str> = dropped.iter().map(|c| c.source.as_str()).collect();
        assert!(dropped_sources.contains(&"l"));
    }

    #[test]
    fn round_robin_spreads_drops_across_sources() {
        let items = vec![
            wc("a", Priority::Normal, 50, 0.5, 0),
            wc("a", Priority::Normal, 50, 0.4, 1),
            wc("a", Priority::Normal, 50, 0.3, 2),
            wc("b", Priority::Normal, 50, 0.5, 0),
            wc("b", Priority::Normal, 50, 0.4, 1),
            wc("b", Priority::Normal, 50, 0.3, 2),
        ];
        // Total 300; budget 150 → drop 150 tokens (3 items).
        let budget = TokenBudget::new(150, 0);
        let (kept, dropped) = apply_budget(items, &budget, PruneStrategy::RoundRobin);
        assert_eq!(kept.len(), 3);
        assert_eq!(dropped.len(), 3);
        // Round robin should drop ~the same count from each source.
        let drop_count_a = dropped.iter().filter(|d| d.source.as_str() == "a").count();
        let drop_count_b = dropped.iter().filter(|d| d.source.as_str() == "b").count();
        // With 3 drops across 2 sources, we expect (2,1) or (1,2), never
        // (3,0) or (0,3).
        assert!(
            drop_count_a == 1 && drop_count_b == 2 || drop_count_a == 2 && drop_count_b == 1,
            "round-robin should spread: got a={drop_count_a} b={drop_count_b}"
        );
    }

    #[test]
    fn kept_preserves_input_order_within_priority() {
        let items = vec![
            wc("a", Priority::Normal, 10, 0.9, 0),
            wc("a", Priority::Normal, 10, 0.5, 1),
            wc("a", Priority::Normal, 10, 0.7, 2),
            wc("a", Priority::Normal, 10, 0.1, 3),
        ];
        // Drop 20 tokens (2 items). Expect the 0.1 and 0.5 to drop, 0.9
        // and 0.7 to remain *in original order* (index 0, 2).
        let budget = TokenBudget::new(20, 0);
        let (kept, _dropped) = apply_budget(items, &budget, PruneStrategy::ImportanceFirst);
        assert_eq!(kept.len(), 2);
        assert_eq!(kept[0].source_index, 0);
        assert_eq!(kept[1].source_index, 2);
    }
}