aft/compress/
caps.rs

1use std::collections::BTreeMap;
2
3pub const CAP_ERRORS: usize = 20;
4pub const CAP_WARNINGS: usize = 10;
5pub const CAP_LIST: usize = 20;
6pub const CAP_INVENTORY: usize = 50;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
9pub enum DropClass {
10    Error,
11    Warning,
12    Failure,
13    Issue,
14    List,
15    Inventory,
16    Timing,
17}
18
19impl DropClass {
20    pub fn default_cap(self) -> usize {
21        match self {
22            Self::Error | Self::Failure => CAP_ERRORS,
23            Self::Warning => CAP_WARNINGS,
24            Self::Issue | Self::List | Self::Timing => CAP_LIST,
25            Self::Inventory => CAP_INVENTORY,
26        }
27    }
28
29    pub fn singular(self) -> &'static str {
30        match self {
31            Self::Error => "error",
32            Self::Warning => "warning",
33            Self::Failure => "failure",
34            Self::Issue => "issue",
35            Self::List => "list item",
36            Self::Inventory => "inventory item",
37            Self::Timing => "timing line",
38        }
39    }
40
41    pub fn plural(self) -> &'static str {
42        match self {
43            Self::Error => "errors",
44            Self::Warning => "warnings",
45            Self::Failure => "failures",
46            Self::Issue => "issues",
47            Self::List => "list items",
48            Self::Inventory => "inventory items",
49            Self::Timing => "timing lines",
50        }
51    }
52}
53
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct ClassifiedBlock {
56    pub class: Option<DropClass>,
57    pub text: String,
58}
59
60impl ClassifiedBlock {
61    pub fn new(class: DropClass, text: impl Into<String>) -> Self {
62        Self {
63            class: Some(class),
64            text: text.into(),
65        }
66    }
67
68    pub fn unclassified(text: impl Into<String>) -> Self {
69        Self {
70            class: None,
71            text: text.into(),
72        }
73    }
74}
75
76#[derive(Debug, Clone, PartialEq, Eq, Default)]
77pub struct ClassCapResult {
78    pub text: String,
79    pub dropped_by_class: BTreeMap<DropClass, usize>,
80}
81
82pub fn cap_classified_blocks(blocks: Vec<ClassifiedBlock>) -> ClassCapResult {
83    cap_classified_blocks_with(blocks, DropClass::default_cap)
84}
85
86pub fn cap_classified_blocks_with<F>(blocks: Vec<ClassifiedBlock>, cap_for: F) -> ClassCapResult
87where
88    F: Fn(DropClass) -> usize,
89{
90    let mut seen_by_class: BTreeMap<DropClass, usize> = BTreeMap::new();
91    let mut dropped_by_class: BTreeMap<DropClass, usize> = BTreeMap::new();
92    let mut kept = Vec::new();
93
94    for block in blocks {
95        let Some(class) = block.class else {
96            kept.push(block.text);
97            continue;
98        };
99
100        let seen = seen_by_class.entry(class).or_default();
101        *seen += 1;
102        if *seen <= cap_for(class) {
103            kept.push(block.text);
104        } else {
105            *dropped_by_class.entry(class).or_default() += 1;
106        }
107    }
108
109    ClassCapResult {
110        text: join_blocks(kept),
111        dropped_by_class,
112    }
113}
114
115pub fn join_blocks(blocks: Vec<String>) -> String {
116    blocks
117        .into_iter()
118        .map(|block| block.trim_end().to_string())
119        .filter(|block| !block.is_empty())
120        .collect::<Vec<_>>()
121        .join("\n")
122}
123
124#[cfg(test)]
125mod tests {
126    use super::*;
127
128    #[test]
129    fn keeps_error_and_warning_blocks_spread_through_long_stream() {
130        let mut blocks = Vec::new();
131        for index in 0..40 {
132            blocks.push(ClassifiedBlock::new(
133                DropClass::Error,
134                format!("error {index}\n  error context {index}"),
135            ));
136            if index < 20 {
137                blocks.push(ClassifiedBlock::new(
138                    DropClass::Warning,
139                    format!("warning {index}\n  warning context {index}"),
140                ));
141            }
142            blocks.push(ClassifiedBlock::unclassified(format!("progress {index}")));
143        }
144
145        let capped = cap_classified_blocks(blocks);
146
147        assert_eq!(capped.text.matches("error context").count(), CAP_ERRORS);
148        assert_eq!(capped.text.matches("warning context").count(), CAP_WARNINGS);
149        assert_eq!(capped.dropped_by_class.get(&DropClass::Error), Some(&20));
150        assert_eq!(capped.dropped_by_class.get(&DropClass::Warning), Some(&10));
151        assert!(capped.text.contains("progress 39"));
152        assert!(!capped.text.contains("error 39\n  error context 39"));
153    }
154
155    #[test]
156    fn caps_by_class_without_splitting_blocks() {
157        let blocks = vec![
158            ClassifiedBlock::new(DropClass::Error, "error 1\n  context"),
159            ClassifiedBlock::new(DropClass::Warning, "warning 1\n  context"),
160            ClassifiedBlock::new(DropClass::Error, "error 2\n  context"),
161        ];
162
163        let capped = cap_classified_blocks_with(blocks, |class| match class {
164            DropClass::Error => 1,
165            DropClass::Warning => 10,
166            _ => 0,
167        });
168
169        assert!(capped.text.contains("error 1\n  context"));
170        assert!(capped.text.contains("warning 1\n  context"));
171        assert!(!capped.text.contains("error 2"));
172        assert_eq!(capped.dropped_by_class.get(&DropClass::Error), Some(&1));
173    }
174}
aft/compress/caps.rs

aft/compress/
caps.rs
