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