Skip to main content

context_forge/analysis/
injection.rs

1use std::cmp::Ordering;
2
3use crate::analysis::scoring::ScoringConfig;
4
5/// Configuration for progressive injection policy.
6#[derive(Debug, Clone)]
7#[non_exhaustive]
8pub struct InjectionConfig {
9    /// Budget escalation factor per compaction level (default: 0.25).
10    pub escalation_factor: f64,
11    /// Optional absolute cap on scaled budget.
12    pub max_budget_cap: Option<usize>,
13}
14
15impl Default for InjectionConfig {
16    fn default() -> Self {
17        Self {
18            escalation_factor: 0.25,
19            max_budget_cap: None,
20        }
21    }
22}
23
24/// Compute the effective token budget based on compaction depth.
25///
26/// Formula: `base_budget * (1 + (compaction_count - 1) * escalation_factor)`
27/// Clamped to `max_budget_cap` if set.
28///
29/// `compaction_count` of 0 or `None` is treated as 1 (no scaling).
30#[must_use]
31#[allow(
32    clippy::cast_precision_loss,
33    clippy::cast_possible_truncation,
34    clippy::cast_sign_loss,
35    reason = "f64 scaling arithmetic is required by the policy formula"
36)]
37pub fn scale_budget(
38    base_budget: usize,
39    compaction_count: Option<i64>,
40    config: &InjectionConfig,
41) -> usize {
42    let count = compaction_count.unwrap_or(1);
43    if count <= 1 || config.escalation_factor <= 0.0 {
44        return base_budget;
45    }
46
47    let scale = 1.0 + (count - 1) as f64 * config.escalation_factor;
48    let scaled = base_budget as f64 * scale;
49    let ceiled = scaled.ceil();
50
51    let mut result = if !ceiled.is_finite() || ceiled <= 0.0 {
52        base_budget
53    } else if ceiled >= usize::MAX as f64 {
54        usize::MAX
55    } else {
56        ceiled as usize
57    };
58
59    if let Some(cap) = config.max_budget_cap {
60        result = result.min(cap);
61    }
62
63    result
64}
65
66/// Return a `ScoringConfig` with category weights redistributed
67/// according to the compaction-level priority table.
68///
69/// For count <= 1, weights are left unchanged.
70/// For count == 2 and count >= 3, existing category weights from `base`
71/// are sorted descending and reassigned by priority order:
72/// - Count 2: Reinforcing > Corrective > Stateful > Decisive
73/// - Count 3+: Reinforcing > Stateful > Corrective > Decisive
74///
75/// `uncategorized_weight` and `importance_half_life_secs` are preserved from `base`.
76#[must_use]
77pub fn adjust_weights(base: &ScoringConfig, compaction_count: Option<i64>) -> ScoringConfig {
78    let count = compaction_count.unwrap_or(1);
79    if count <= 1 {
80        return base.clone();
81    }
82
83    let mut adjusted = base.clone();
84    let mut weights = [
85        base.corrective_weight,
86        base.decisive_weight,
87        base.stateful_weight,
88        base.reinforcing_weight,
89    ];
90    weights.sort_by(|a, b| b.partial_cmp(a).unwrap_or(Ordering::Equal));
91
92    if count == 2 {
93        adjusted.reinforcing_weight = weights[0];
94        adjusted.corrective_weight = weights[1];
95        adjusted.stateful_weight = weights[2];
96        adjusted.decisive_weight = weights[3];
97        return adjusted;
98    }
99
100    adjusted.reinforcing_weight = weights[0];
101    adjusted.stateful_weight = weights[1];
102    adjusted.corrective_weight = weights[2];
103    adjusted.decisive_weight = weights[3];
104    adjusted
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    // Budget scaling tests
112    #[test]
113    fn scale_budget_no_scaling_at_count_one() {
114        let config = InjectionConfig::default();
115        assert_eq!(scale_budget(2048, Some(1), &config), 2048);
116    }
117
118    #[test]
119    fn scale_budget_no_scaling_at_count_zero() {
120        let config = InjectionConfig::default();
121        assert_eq!(scale_budget(2048, Some(0), &config), 2048);
122    }
123
124    #[test]
125    fn scale_budget_no_scaling_at_count_none() {
126        let config = InjectionConfig::default();
127        assert_eq!(scale_budget(2048, None, &config), 2048);
128    }
129
130    #[test]
131    fn scale_budget_scales_at_count_two() {
132        let config = InjectionConfig::default();
133        assert_eq!(scale_budget(2048, Some(2), &config), 2560); // 2048 * 1.25
134    }
135
136    #[test]
137    fn scale_budget_scales_at_count_three() {
138        let config = InjectionConfig::default();
139        assert_eq!(scale_budget(2048, Some(3), &config), 3072); // 2048 * 1.5
140    }
141
142    #[test]
143    fn scale_budget_scales_at_count_five() {
144        let config = InjectionConfig::default();
145        assert_eq!(scale_budget(2048, Some(5), &config), 4096); // 2048 * 2.0
146    }
147
148    #[test]
149    fn scale_budget_respects_max_cap() {
150        let config = InjectionConfig {
151            max_budget_cap: Some(3000),
152            ..InjectionConfig::default()
153        };
154        assert_eq!(scale_budget(2048, Some(5), &config), 3000);
155    }
156
157    #[test]
158    fn scale_budget_no_scaling_with_zero_escalation() {
159        let config = InjectionConfig {
160            escalation_factor: 0.0,
161            ..InjectionConfig::default()
162        };
163        assert_eq!(scale_budget(2048, Some(5), &config), 2048);
164    }
165
166    #[test]
167    fn scale_budget_no_scaling_with_negative_escalation() {
168        let config = InjectionConfig {
169            escalation_factor: -0.5,
170            ..InjectionConfig::default()
171        };
172        assert_eq!(scale_budget(2048, Some(5), &config), 2048);
173    }
174
175    // Weight adjustment tests
176    #[test]
177    fn adjust_weights_default_at_count_zero() {
178        let base = ScoringConfig::default();
179        let adjusted = adjust_weights(&base, Some(0));
180        assert!((adjusted.corrective_weight - 1.5).abs() < f64::EPSILON);
181        assert!((adjusted.decisive_weight - 1.3).abs() < f64::EPSILON);
182        assert!((adjusted.stateful_weight - 1.2).abs() < f64::EPSILON);
183        assert!((adjusted.reinforcing_weight - 1.0).abs() < f64::EPSILON);
184    }
185
186    #[test]
187    fn adjust_weights_default_at_count_one() {
188        let base = ScoringConfig::default();
189        let adjusted = adjust_weights(&base, Some(1));
190        assert!((adjusted.corrective_weight - 1.5).abs() < f64::EPSILON);
191        assert!((adjusted.decisive_weight - 1.3).abs() < f64::EPSILON);
192        assert!((adjusted.stateful_weight - 1.2).abs() < f64::EPSILON);
193        assert!((adjusted.reinforcing_weight - 1.0).abs() < f64::EPSILON);
194    }
195
196    #[test]
197    fn adjust_weights_drift_at_count_two() {
198        let base = ScoringConfig::default();
199        let adjusted = adjust_weights(&base, Some(2));
200        assert!((adjusted.reinforcing_weight - 1.5).abs() < f64::EPSILON);
201        assert!((adjusted.corrective_weight - 1.3).abs() < f64::EPSILON);
202        assert!((adjusted.stateful_weight - 1.2).abs() < f64::EPSILON);
203        assert!((adjusted.decisive_weight - 1.0).abs() < f64::EPSILON);
204    }
205
206    #[test]
207    fn adjust_weights_heavy_drift_at_count_three() {
208        let base = ScoringConfig::default();
209        let adjusted = adjust_weights(&base, Some(3));
210        assert!((adjusted.reinforcing_weight - 1.5).abs() < f64::EPSILON);
211        assert!((adjusted.stateful_weight - 1.3).abs() < f64::EPSILON);
212        assert!((adjusted.corrective_weight - 1.2).abs() < f64::EPSILON);
213        assert!((adjusted.decisive_weight - 1.0).abs() < f64::EPSILON);
214    }
215
216    #[test]
217    fn adjust_weights_heavy_drift_at_count_ten() {
218        let base = ScoringConfig::default();
219        let adjusted = adjust_weights(&base, Some(10));
220        assert!((adjusted.reinforcing_weight - 1.5).abs() < f64::EPSILON);
221        assert!((adjusted.stateful_weight - 1.3).abs() < f64::EPSILON);
222        assert!((adjusted.corrective_weight - 1.2).abs() < f64::EPSILON);
223        assert!((adjusted.decisive_weight - 1.0).abs() < f64::EPSILON);
224    }
225
226    #[test]
227    fn adjust_weights_preserves_uncategorized() {
228        let base = ScoringConfig {
229            uncategorized_weight: 0.75,
230            ..ScoringConfig::default()
231        };
232        let adjusted = adjust_weights(&base, Some(3));
233        assert!((adjusted.uncategorized_weight - 0.75).abs() < f64::EPSILON);
234    }
235
236    #[test]
237    fn adjust_weights_preserves_half_life() {
238        let base = ScoringConfig {
239            importance_half_life_secs: 123_456.0,
240            ..ScoringConfig::default()
241        };
242        let adjusted = adjust_weights(&base, Some(3));
243        assert!((adjusted.importance_half_life_secs - 123_456.0).abs() < f64::EPSILON);
244    }
245
246    #[test]
247    fn adjust_weights_none_uses_default_order() {
248        let base = ScoringConfig::default();
249        let adjusted = adjust_weights(&base, None);
250        assert!((adjusted.corrective_weight - 1.5).abs() < f64::EPSILON);
251        assert!((adjusted.decisive_weight - 1.3).abs() < f64::EPSILON);
252        assert!((adjusted.stateful_weight - 1.2).abs() < f64::EPSILON);
253        assert!((adjusted.reinforcing_weight - 1.0).abs() < f64::EPSILON);
254    }
255
256    #[test]
257    fn scale_budget_no_scaling_at_negative_count() {
258        let config = InjectionConfig::default();
259        assert_eq!(scale_budget(2048, Some(-1), &config), 2048);
260    }
261
262    #[test]
263    fn adjust_weights_default_at_negative_count() {
264        let base = ScoringConfig::default();
265        let adjusted = adjust_weights(&base, Some(-1));
266        assert!((adjusted.corrective_weight - 1.5).abs() < f64::EPSILON);
267        assert!((adjusted.decisive_weight - 1.3).abs() < f64::EPSILON);
268        assert!((adjusted.stateful_weight - 1.2).abs() < f64::EPSILON);
269        assert!((adjusted.reinforcing_weight - 1.0).abs() < f64::EPSILON);
270    }
271}