tokenburn_core/features/report/
drill.rs1use std::borrow::Cow;
12
13use super::chart::{bucket_label_of, chart_bucket, pie, time_series, Pie, TimeSeries};
14use super::ops::{build_report, by_project};
15use super::types::{Query, ReportRow};
16use crate::features::usage::{Row, Summary};
17
18#[derive(Debug, Clone, Default, PartialEq, Eq)]
20pub struct Selection {
21 pub bucket: Option<String>,
23 pub project: Option<String>,
25}
26
27impl Selection {
28 pub fn is_empty(&self) -> bool {
29 self.bucket.is_none() && self.project.is_none()
30 }
31
32 pub fn clear(&mut self) {
33 *self = Self::default();
34 }
35
36 pub fn toggle_bucket(&mut self, label: &str) {
38 self.bucket = toggled(self.bucket.take(), label);
39 }
40
41 pub fn toggle_project(&mut self, key: &str) {
43 self.project = toggled(self.project.take(), key);
44 }
45
46 fn row_in_project(&self, row: &Row) -> bool {
47 self.project
48 .as_deref()
49 .is_none_or(|p| project_key(row) == p)
50 }
51
52 fn row_in_bucket(&self, row: &Row, query: &Query) -> bool {
53 self.bucket
54 .as_deref()
55 .is_none_or(|b| bucket_label_of(row, chart_bucket(query)) == b)
56 }
57
58 pub fn describe(&self) -> Vec<String> {
60 let mut out = Vec::new();
61 if let Some(b) = &self.bucket {
62 out.push(b.clone());
63 }
64 if let Some(p) = &self.project {
65 out.push(super::chart::short_name(p));
66 }
67 out
68 }
69}
70
71fn toggled(current: Option<String>, label: &str) -> Option<String> {
72 if current.as_deref() == Some(label) {
73 None
74 } else {
75 Some(label.to_string())
76 }
77}
78
79pub fn project_key(row: &Row) -> String {
81 format!("{}:{}", row.tool.as_str(), row.project)
82}
83
84#[derive(Debug, Clone)]
86pub struct Drill {
87 pub report: Vec<ReportRow>,
88 pub total: Summary,
89 pub projects: Vec<(String, Summary)>,
90 pub series: TimeSeries,
92 pub pie: Pie,
94 pub selection: Selection,
96}
97
98pub fn drill(rows: &[Row], query: &Query, selection: &Selection) -> Drill {
101 let full = time_series(rows, query);
102 let mut sel = selection.clone();
103 if sel
104 .bucket
105 .as_ref()
106 .is_some_and(|b| !full.labels.iter().any(|l| l == b))
107 {
108 sel.bucket = None;
109 }
110 if sel
111 .project
112 .as_ref()
113 .is_some_and(|p| !rows.iter().any(|r| &project_key(r) == p))
114 {
115 sel.project = None;
116 }
117
118 let in_project: Cow<[Row]> = if sel.project.is_none() {
122 Cow::Borrowed(rows)
123 } else {
124 Cow::Owned(
125 rows.iter()
126 .filter(|r| sel.row_in_project(r))
127 .cloned()
128 .collect(),
129 )
130 };
131 let in_bucket: Cow<[Row]> = if sel.bucket.is_none() {
132 Cow::Borrowed(rows)
133 } else {
134 Cow::Owned(
135 rows.iter()
136 .filter(|r| sel.row_in_bucket(r, query))
137 .cloned()
138 .collect(),
139 )
140 };
141 let both: Cow<[Row]> = match (sel.project.is_some(), sel.bucket.is_some()) {
142 (false, false) => Cow::Borrowed(rows),
143 (true, false) => Cow::Borrowed(&in_project),
144 (false, true) => Cow::Borrowed(&in_bucket),
145 (true, true) => Cow::Owned(
146 in_project
147 .iter()
148 .filter(|r| sel.row_in_bucket(r, query))
149 .cloned()
150 .collect(),
151 ),
152 };
153
154 let series = if sel.project.is_some() {
156 restrict(&full, &in_project, query)
157 } else {
158 full
159 };
160 let projects = by_project(&both);
161 Drill {
162 report: build_report(&both, query),
163 total: Summary::of(both.iter()),
164 series,
165 pie: pie(&by_project(&in_bucket), 8),
166 projects,
167 selection: sel,
168 }
169}
170
171fn restrict(full: &TimeSeries, rows: &[Row], query: &Query) -> TimeSeries {
173 let mut out = TimeSeries {
174 bucket: full.bucket,
175 labels: full.labels.clone(),
176 tokens: vec![[0; 8]; full.labels.len()],
177 cost: vec![[0.0; 8]; full.labels.len()],
178 };
179 let bucket = chart_bucket(query);
180 for r in rows {
181 let l = bucket_label_of(r, bucket);
182 if let Some(i) = out.labels.iter().position(|x| *x == l) {
183 let t = r.tool.series_index();
184 out.tokens[i][t] += r.input + r.output + r.cache_read + r.cache_write;
185 out.cost[i][t] += r.cost;
186 }
187 }
188 out
189}
190
191#[cfg(test)]
192mod tests {
193 use super::*;
194 use crate::features::report::types::{Bucket, Window};
195 use crate::features::usage::Tool;
196
197 fn row(tool: Tool, project: &str, ts: &str, tokens: u64) -> Row {
198 Row {
199 tool,
200 project: project.into(),
201 id: "x".into(),
202 ts: ts.parse().unwrap(),
203 input: tokens,
204 output: 0,
205 cache_read: 0,
206 cache_write: 0,
207 cost: tokens as f64 / 1000.0,
208 }
209 }
210
211 fn rows() -> Vec<Row> {
212 vec![
213 row(Tool::Pi, "a", "2026-05-01T12:00:00Z", 1000),
214 row(Tool::Pi, "b", "2026-05-01T13:00:00Z", 500),
215 row(Tool::Claude, "a", "2026-05-03T12:00:00Z", 2000),
216 ]
217 }
218
219 fn q() -> Query {
220 Query {
221 bucket: Bucket::Day,
222 window: Window::All,
223 ..Query::default()
224 }
225 }
226
227 fn sel(bucket: Option<&str>, project: Option<&str>) -> Selection {
228 Selection {
229 bucket: bucket.map(Into::into),
230 project: project.map(Into::into),
231 }
232 }
233
234 #[test]
235 fn no_selection_is_the_whole_dataset() {
236 let d = drill(&rows(), &q(), &Selection::default());
237 assert_eq!(d.total.total_tokens, 3500);
238 assert_eq!(d.series.labels.len(), 3, "May 1, 2, 3 (gap filled)");
239 assert!(d.selection.is_empty());
240 }
241
242 #[test]
243 fn a_bucket_scopes_totals_report_and_donut_but_not_the_bars() {
244 let all = drill(&rows(), &q(), &Selection::default());
245 let first = all.series.labels[0].clone();
246 let d = drill(&rows(), &q(), &sel(Some(&first), None));
247 assert_eq!(d.total.total_tokens, 1500);
248 assert!(d.report.iter().all(|r| r.bucket.as_deref() == Some(&first)));
249 assert_eq!(d.pie.slices.len(), 2, "only the two projects of that day");
250 assert_eq!(d.series.labels, all.series.labels, "x-axis unchanged");
251 assert_eq!(
252 d.series.bucket_total(crate::Metric::Tokens, 2),
253 2000.0,
254 "bars still show the other days"
255 );
256 }
257
258 #[test]
259 fn a_project_scopes_totals_report_and_bars_but_not_the_donut() {
260 let all = drill(&rows(), &q(), &Selection::default());
261 let d = drill(&rows(), &q(), &sel(None, Some("pi:a")));
262 assert_eq!(d.total.total_tokens, 1000);
263 assert_eq!(d.series.labels, all.series.labels, "same x-axis");
264 assert_eq!(d.series.bucket_total(crate::Metric::Tokens, 0), 1000.0);
265 assert_eq!(d.series.bucket_total(crate::Metric::Tokens, 2), 0.0);
266 assert_eq!(
267 d.pie.slices.len(),
268 all.pie.slices.len(),
269 "donut shows everything"
270 );
271 }
272
273 #[test]
274 fn both_intersect() {
275 let all = drill(&rows(), &q(), &Selection::default());
276 let first = all.series.labels[0].clone();
277 let d = drill(&rows(), &q(), &sel(Some(&first), Some("pi:b")));
278 assert_eq!(d.total.total_tokens, 500);
279 }
280
281 #[test]
282 fn toggling_selects_then_clears() {
283 let mut s = Selection::default();
284 s.toggle_bucket("2026-05-01");
285 assert_eq!(s.bucket.as_deref(), Some("2026-05-01"));
286 s.toggle_bucket("2026-05-03");
287 assert_eq!(
288 s.bucket.as_deref(),
289 Some("2026-05-03"),
290 "another bar replaces it"
291 );
292 s.toggle_bucket("2026-05-03");
293 assert!(s.bucket.is_none(), "same bar again clears");
294 s.toggle_project("pi:a");
295 assert!(!s.is_empty());
296 s.clear();
297 assert!(s.is_empty());
298 }
299
300 #[test]
301 fn stale_selections_are_dropped() {
302 let d = drill(&rows(), &q(), &sel(Some("1999-01-01"), Some("pi:ghost")));
303 assert!(d.selection.is_empty());
304 assert_eq!(d.total.total_tokens, 3500, "falls back to everything");
305 }
306
307 #[test]
308 fn works_when_the_report_is_not_bucketed() {
309 let query = Query {
311 bucket: Bucket::None,
312 window: Window::Month,
313 ..Query::default()
314 };
315 let all = drill(&rows(), &query, &Selection::default());
316 let day = all.series.labels[0].clone();
317 let d = drill(&rows(), &query, &sel(Some(&day), None));
318 assert_eq!(d.total.total_tokens, 1500);
319 assert!(d.report.iter().any(|r| r.tool == "TOTAL"));
320 }
321
322 #[test]
323 fn describe_lists_the_chips() {
324 assert!(Selection::default().describe().is_empty());
325 let d = sel(Some("2026-05-01"), Some("pi:Users-me-atlas")).describe();
326 assert_eq!(d, vec!["2026-05-01".to_string(), "pi:atlas".to_string()]);
327 }
328
329 #[test]
330 fn slice_keys_identify_projects_and_other_is_unselectable() {
331 let d = drill(&rows(), &q(), &Selection::default());
332 assert!(d.pie.slices.iter().all(|s| s.key.is_some()));
333 assert_eq!(project_key(&rows()[0]), "pi:a");
334 }
335}