devboy_format_pipeline/trim/
mod.rs

1//! Tree trimming algorithms by budget.
2//!
3//! Algorithm selection depends on tree size:
4//! - < 100 nodes: Tree Knapsack DP (exact optimum)
5//! - 100-999: Greedy fractional (>= 63% of optimum, ~90%+ in practice)
6//! - 1000-9999: Hierarchical WFQ (proportionally fair)
7//! - >= 10000: Head+Tail linear (heuristic for logs)
8
9pub mod greedy;
10pub mod head_tail;
11pub mod knapsack;
12pub mod wfq;
13
14use crate::tree::TrimNode;
15
16/// Trim the tree to fit the budget, automatically selecting the optimal algorithm.
17///
18/// Budget is the maximum weight (in tokens) of the resulting included subtree.
19/// After the call, excluded nodes have `included = false`.
20pub fn trim(tree: &mut TrimNode, budget: usize) {
21    // If it already fits — do nothing
22    if tree.total_weight() <= budget {
23        return;
24    }
25
26    let count = tree.count_nodes();
27    match count {
28        0..=99 => knapsack::solve(tree, budget),
29        100..=999 => greedy::solve(tree, budget),
30        1000..=9999 => wfq::solve(tree, budget),
31        _ => head_tail::solve(tree, budget),
32    }
33}
34
35#[cfg(test)]
36mod tests {
37    use super::*;
38    use crate::tree::NodeKind;
39
40    fn make_tree(item_weights: &[usize], values: &[f64]) -> TrimNode {
41        let mut root = TrimNode::new(0, NodeKind::Root, 0);
42        for (i, (&w, &v)) in item_weights.iter().zip(values.iter()).enumerate() {
43            let mut node = TrimNode::new(i + 1, NodeKind::Item { index: i }, w);
44            node.value = v;
45            root.children.push(node);
46        }
47        root
48    }
49
50    #[test]
51    fn test_trim_no_op_when_fits() {
52        let mut tree = make_tree(&[10, 20, 30], &[1.0, 1.0, 1.0]);
53        trim(&mut tree, 100);
54        assert_eq!(tree.included_items_count(), 3);
55    }
56
57    #[test]
58    fn test_trim_reduces_weight() {
59        let mut tree = make_tree(&[100, 200, 300], &[1.0, 0.5, 0.2]);
60        let budget = 250;
61        trim(&mut tree, budget);
62        assert!(tree.total_weight() <= budget);
63        // Should keep at least one element
64        assert!(tree.included_items_count() >= 1);
65    }
66
67    #[test]
68    fn test_trim_prefers_high_value() {
69        let mut tree = make_tree(&[100, 100, 100], &[0.1, 0.9, 0.5]);
70        trim(&mut tree, 150);
71        // Item with value 0.9 (index 1) should be included
72        let included = tree.included_item_indices();
73        assert!(included.contains(&1), "High-value item should be kept");
74    }
75
76    #[test]
77    fn test_trim_selects_correct_algorithm() {
78        // < 100 nodes → knapsack
79        let mut small = make_tree(&[10; 5], &[1.0; 5]);
80        trim(&mut small, 30);
81        assert!(small.total_weight() <= 30);
82
83        // > 100 nodes → greedy
84        let weights: Vec<usize> = (0..120).map(|_| 10).collect();
85        let values: Vec<f64> = (0..120).map(|i| 1.0 - i as f64 * 0.005).collect();
86        let mut large = make_tree(&weights, &values);
87        trim(&mut large, 500);
88        assert!(large.total_weight() <= 500);
89    }
90
91    #[test]
92    fn test_trim_wfq_range() {
93        // 1000-9999 nodes → WFQ
94        let weights: Vec<usize> = (0..1500).map(|_| 5).collect();
95        let values: Vec<f64> = (0..1500).map(|i| 1.0 - i as f64 * 0.0005).collect();
96        let mut tree = make_tree(&weights, &values);
97        let budget = 2000;
98        trim(&mut tree, budget);
99        assert!(tree.total_weight() <= budget);
100        assert!(tree.included_items_count() > 0);
101    }
102
103    #[test]
104    fn test_trim_head_tail_range() {
105        // >= 10000 nodes → head_tail
106        let weights: Vec<usize> = (0..10500).map(|_| 1).collect();
107        let values: Vec<f64> = (0..10500).map(|_| 1.0).collect();
108        let mut tree = make_tree(&weights, &values);
109        let budget = 500;
110        trim(&mut tree, budget);
111        assert!(tree.total_weight() <= budget);
112    }
113}
devboy_format_pipeline/trim/mod.rs

devboy_format_pipeline/trim/
mod.rs
