torsh_profiler/
custom_export.rs

1//! Custom export formats for profiling data
2
3use crate::{ProfileEvent, TorshResult};
4use serde_json::{json, Value};
5use std::collections::HashMap;
6use std::fs::File;
7use std::io::Write;
8use torsh_core::TorshError;
9
10/// Custom export format configuration
11#[derive(Debug, Clone)]
12pub struct CustomExportFormat {
13    pub name: String,
14    pub description: String,
15    pub file_extension: String,
16    pub schema: ExportSchema,
17}
18
19/// Export schema definition
20#[derive(Debug, Clone)]
21pub enum ExportSchema {
22    /// JSON with custom field mapping
23    Json {
24        field_mapping: HashMap<String, String>,
25        include_metadata: bool,
26        pretty_print: bool,
27    },
28    /// CSV with custom column configuration
29    Csv {
30        columns: Vec<CsvColumn>,
31        delimiter: char,
32        include_header: bool,
33    },
34    /// XML format
35    Xml {
36        root_element: String,
37        event_element: String,
38        field_mapping: HashMap<String, String>,
39    },
40    /// Custom text format with template
41    Text { template: String, separator: String },
42}
43
44/// CSV column configuration
45#[derive(Debug, Clone)]
46pub struct CsvColumn {
47    pub name: String,
48    pub field: String,
49    pub formatter: Option<CsvFormatter>,
50}
51
52/// CSV field formatters
53#[derive(Debug, Clone)]
54pub enum CsvFormatter {
55    Duration(DurationFormat),
56    Memory(MemoryFormat),
57    Number(NumberFormat),
58    Text(TextFormat),
59}
60
61#[derive(Debug, Clone)]
62pub enum DurationFormat {
63    Microseconds,
64    Milliseconds,
65    Seconds,
66    HumanReadable,
67}
68
69#[derive(Debug, Clone)]
70pub enum MemoryFormat {
71    Bytes,
72    Kilobytes,
73    Megabytes,
74    Gigabytes,
75    HumanReadable,
76}
77
78#[derive(Debug, Clone)]
79pub enum NumberFormat {
80    Default,
81    Scientific,
82    Percentage,
83    WithCommas,
84}
85
86#[derive(Debug, Clone)]
87pub enum TextFormat {
88    Default,
89    Uppercase,
90    Lowercase,
91    Truncate(usize),
92}
93
94/// Custom export engine
95#[derive(Debug, Clone)]
96pub struct CustomExporter {
97    formats: HashMap<String, CustomExportFormat>,
98}
99
100impl CustomExporter {
101    /// Create a new custom exporter
102    pub fn new() -> Self {
103        let mut exporter = Self {
104            formats: HashMap::new(),
105        };
106
107        // Register default custom formats
108        exporter.register_default_formats();
109        exporter
110    }
111
112    /// Register a custom export format
113    pub fn register_format(&mut self, format: CustomExportFormat) {
114        self.formats.insert(format.name.clone(), format);
115    }
116
117    /// Get available format names
118    pub fn get_format_names(&self) -> Vec<String> {
119        self.formats.keys().cloned().collect()
120    }
121
122    /// Export events using a registered format
123    pub fn export(
124        &self,
125        events: &[ProfileEvent],
126        format_name: &str,
127        path: &str,
128    ) -> TorshResult<()> {
129        let format = self.formats.get(format_name).ok_or_else(|| {
130            TorshError::InvalidArgument(format!("Unknown export format: {format_name}"))
131        })?;
132
133        match &format.schema {
134            ExportSchema::Json {
135                field_mapping,
136                include_metadata,
137                pretty_print,
138            } => self.export_json(
139                events,
140                field_mapping,
141                *include_metadata,
142                *pretty_print,
143                path,
144            ),
145            ExportSchema::Csv {
146                columns,
147                delimiter,
148                include_header,
149            } => self.export_csv(events, columns, *delimiter, *include_header, path),
150            ExportSchema::Xml {
151                root_element,
152                event_element,
153                field_mapping,
154            } => self.export_xml(events, root_element, event_element, field_mapping, path),
155            ExportSchema::Text {
156                template,
157                separator,
158            } => self.export_text(events, template, separator, path),
159        }
160    }
161
162    /// Register default custom formats
163    fn register_default_formats(&mut self) {
164        // Compact JSON format
165        let compact_json = CustomExportFormat {
166            name: "compact_json".to_string(),
167            description: "Compact JSON format with minimal fields".to_string(),
168            file_extension: "json".to_string(),
169            schema: ExportSchema::Json {
170                field_mapping: [
171                    ("name".to_string(), "n".to_string()),
172                    ("duration_us".to_string(), "d".to_string()),
173                    ("category".to_string(), "c".to_string()),
174                ]
175                .iter()
176                .cloned()
177                .collect(),
178                include_metadata: false,
179                pretty_print: false,
180            },
181        };
182        self.register_format(compact_json);
183
184        // Performance-focused CSV
185        let perf_csv = CustomExportFormat {
186            name: "performance_csv".to_string(),
187            description: "CSV focused on performance metrics".to_string(),
188            file_extension: "csv".to_string(),
189            schema: ExportSchema::Csv {
190                columns: vec![
191                    CsvColumn {
192                        name: "Event".to_string(),
193                        field: "name".to_string(),
194                        formatter: None,
195                    },
196                    CsvColumn {
197                        name: "Duration (ms)".to_string(),
198                        field: "duration_us".to_string(),
199                        formatter: Some(CsvFormatter::Duration(DurationFormat::Milliseconds)),
200                    },
201                    CsvColumn {
202                        name: "FLOPS".to_string(),
203                        field: "flops".to_string(),
204                        formatter: Some(CsvFormatter::Number(NumberFormat::WithCommas)),
205                    },
206                    CsvColumn {
207                        name: "Bandwidth (MB)".to_string(),
208                        field: "bytes_transferred".to_string(),
209                        formatter: Some(CsvFormatter::Memory(MemoryFormat::Megabytes)),
210                    },
211                ],
212                delimiter: ',',
213                include_header: true,
214            },
215        };
216        self.register_format(perf_csv);
217
218        // Simple text format
219        let simple_text = CustomExportFormat {
220            name: "simple_text".to_string(),
221            description: "Simple text format for quick viewing".to_string(),
222            file_extension: "txt".to_string(),
223            schema: ExportSchema::Text {
224                template: "{name}: {duration_us}μs ({category})".to_string(),
225                separator: "\n".to_string(),
226            },
227        };
228        self.register_format(simple_text);
229    }
230
231    /// Export as custom JSON
232    fn export_json(
233        &self,
234        events: &[ProfileEvent],
235        field_mapping: &HashMap<String, String>,
236        include_metadata: bool,
237        pretty_print: bool,
238        path: &str,
239    ) -> TorshResult<()> {
240        let mut mapped_events = Vec::new();
241
242        for event in events {
243            let mut mapped_event = serde_json::Map::new();
244
245            // Map fields according to configuration
246            self.map_field(&mut mapped_event, "name", &event.name, field_mapping);
247            self.map_field(
248                &mut mapped_event,
249                "category",
250                &event.category,
251                field_mapping,
252            );
253            self.map_field(
254                &mut mapped_event,
255                "start_us",
256                &event.start_us,
257                field_mapping,
258            );
259            self.map_field(
260                &mut mapped_event,
261                "duration_us",
262                &event.duration_us,
263                field_mapping,
264            );
265            self.map_field(
266                &mut mapped_event,
267                "thread_id",
268                &event.thread_id,
269                field_mapping,
270            );
271
272            if let Some(ops) = event.operation_count {
273                self.map_field(&mut mapped_event, "operation_count", &ops, field_mapping);
274            }
275            if let Some(flops) = event.flops {
276                self.map_field(&mut mapped_event, "flops", &flops, field_mapping);
277            }
278            if let Some(bytes) = event.bytes_transferred {
279                self.map_field(
280                    &mut mapped_event,
281                    "bytes_transferred",
282                    &bytes,
283                    field_mapping,
284                );
285            }
286            if let Some(ref stack_trace) = event.stack_trace {
287                self.map_field(&mut mapped_event, "stack_trace", stack_trace, field_mapping);
288            }
289
290            mapped_events.push(Value::Object(mapped_event));
291        }
292
293        let mut output = serde_json::Map::new();
294        output.insert("events".to_string(), Value::Array(mapped_events));
295
296        if include_metadata {
297            let metadata = json!({
298                "export_timestamp": chrono::Utc::now().to_rfc3339(),
299                "event_count": events.len(),
300                "format": "custom_json"
301            });
302            output.insert("metadata".to_string(), metadata);
303        }
304
305        let json_output = Value::Object(output);
306        let mut file = File::create(path).map_err(|e| {
307            TorshError::InvalidArgument(format!("Failed to create file {path}: {e}"))
308        })?;
309
310        let json_string = if pretty_print {
311            serde_json::to_string_pretty(&json_output)
312        } else {
313            serde_json::to_string(&json_output)
314        }
315        .map_err(|e| TorshError::InvalidArgument(format!("Failed to serialize JSON: {e}")))?;
316
317        file.write_all(json_string.as_bytes())
318            .map_err(|e| TorshError::InvalidArgument(format!("Failed to write file: {e}")))?;
319
320        Ok(())
321    }
322
323    /// Export as custom CSV
324    fn export_csv(
325        &self,
326        events: &[ProfileEvent],
327        columns: &[CsvColumn],
328        delimiter: char,
329        include_header: bool,
330        path: &str,
331    ) -> TorshResult<()> {
332        let mut file = File::create(path).map_err(|e| {
333            TorshError::InvalidArgument(format!("Failed to create file {path}: {e}"))
334        })?;
335
336        // Write header
337        if include_header {
338            let header: Vec<String> = columns.iter().map(|col| col.name.clone()).collect();
339            writeln!(file, "{}", header.join(&delimiter.to_string()))
340                .map_err(|e| TorshError::InvalidArgument(format!("Failed to write header: {e}")))?;
341        }
342
343        // Write events
344        for event in events {
345            let row: Vec<String> = columns
346                .iter()
347                .map(|col| self.format_field_value(event, &col.field, &col.formatter))
348                .collect();
349            writeln!(file, "{}", row.join(&delimiter.to_string()))
350                .map_err(|e| TorshError::InvalidArgument(format!("Failed to write row: {e}")))?;
351        }
352
353        Ok(())
354    }
355
356    /// Export as XML
357    fn export_xml(
358        &self,
359        events: &[ProfileEvent],
360        root_element: &str,
361        event_element: &str,
362        field_mapping: &HashMap<String, String>,
363        path: &str,
364    ) -> TorshResult<()> {
365        let mut file = File::create(path).map_err(|e| {
366            TorshError::InvalidArgument(format!("Failed to create file {path}: {e}"))
367        })?;
368
369        writeln!(file, "<?xml version=\"1.0\" encoding=\"UTF-8\"?>")
370            .map_err(|e| TorshError::InvalidArgument(format!("Failed to write XML header: {e}")))?;
371
372        writeln!(file, "<{root_element}>").map_err(|e| {
373            TorshError::InvalidArgument(format!("Failed to write root element: {e}"))
374        })?;
375
376        for event in events {
377            writeln!(file, "  <{event_element}>").map_err(|e| {
378                TorshError::InvalidArgument(format!("Failed to write event element: {e}"))
379            })?;
380
381            self.write_xml_field(&mut file, "name", &event.name, field_mapping)?;
382            self.write_xml_field(&mut file, "category", &event.category, field_mapping)?;
383            self.write_xml_field(
384                &mut file,
385                "start_us",
386                &event.start_us.to_string(),
387                field_mapping,
388            )?;
389            self.write_xml_field(
390                &mut file,
391                "duration_us",
392                &event.duration_us.to_string(),
393                field_mapping,
394            )?;
395            self.write_xml_field(
396                &mut file,
397                "thread_id",
398                &event.thread_id.to_string(),
399                field_mapping,
400            )?;
401
402            if let Some(ops) = event.operation_count {
403                self.write_xml_field(
404                    &mut file,
405                    "operation_count",
406                    &ops.to_string(),
407                    field_mapping,
408                )?;
409            }
410            if let Some(flops) = event.flops {
411                self.write_xml_field(&mut file, "flops", &flops.to_string(), field_mapping)?;
412            }
413            if let Some(bytes) = event.bytes_transferred {
414                self.write_xml_field(
415                    &mut file,
416                    "bytes_transferred",
417                    &bytes.to_string(),
418                    field_mapping,
419                )?;
420            }
421
422            writeln!(file, "  </{event_element}>").map_err(|e| {
423                TorshError::InvalidArgument(format!("Failed to write closing event element: {e}"))
424            })?;
425        }
426
427        writeln!(file, "</{root_element}>").map_err(|e| {
428            TorshError::InvalidArgument(format!("Failed to write closing root element: {e}"))
429        })?;
430
431        Ok(())
432    }
433
434    /// Export as custom text
435    fn export_text(
436        &self,
437        events: &[ProfileEvent],
438        template: &str,
439        separator: &str,
440        path: &str,
441    ) -> TorshResult<()> {
442        let mut file = File::create(path).map_err(|e| {
443            TorshError::InvalidArgument(format!("Failed to create file {path}: {e}"))
444        })?;
445
446        for (i, event) in events.iter().enumerate() {
447            let formatted = self.format_template(template, event);
448            write!(file, "{formatted}")
449                .map_err(|e| TorshError::InvalidArgument(format!("Failed to write event: {e}")))?;
450
451            if i < events.len() - 1 {
452                write!(file, "{separator}").map_err(|e| {
453                    TorshError::InvalidArgument(format!("Failed to write separator: {e}"))
454                })?;
455            }
456        }
457
458        Ok(())
459    }
460
461    /// Helper function to map JSON fields
462    fn map_field<T: serde::Serialize>(
463        &self,
464        object: &mut serde_json::Map<String, Value>,
465        original_field: &str,
466        value: &T,
467        field_mapping: &HashMap<String, String>,
468    ) {
469        let field_name = field_mapping
470            .get(original_field)
471            .unwrap_or(&original_field.to_string())
472            .clone();
473        object.insert(field_name, json!(value));
474    }
475
476    /// Helper function to write XML fields
477    fn write_xml_field(
478        &self,
479        file: &mut File,
480        original_field: &str,
481        value: &str,
482        field_mapping: &HashMap<String, String>,
483    ) -> TorshResult<()> {
484        let default_field = original_field.to_string();
485        let field_name = field_mapping.get(original_field).unwrap_or(&default_field);
486
487        writeln!(file, "    <{field_name}>{value}</{field_name}>")
488            .map_err(|e| TorshError::InvalidArgument(format!("Failed to write XML field: {e}")))
489    }
490
491    /// Format field value according to formatter
492    fn format_field_value(
493        &self,
494        event: &ProfileEvent,
495        field: &str,
496        formatter: &Option<CsvFormatter>,
497    ) -> String {
498        let raw_value = match field {
499            "name" => event.name.clone(),
500            "category" => event.category.clone(),
501            "start_us" => event.start_us.to_string(),
502            "duration_us" => event.duration_us.to_string(),
503            "thread_id" => event.thread_id.to_string(),
504            "operation_count" => event
505                .operation_count
506                .map_or("".to_string(), |v| v.to_string()),
507            "flops" => event.flops.map_or("".to_string(), |v| v.to_string()),
508            "bytes_transferred" => event
509                .bytes_transferred
510                .map_or("".to_string(), |v| v.to_string()),
511            "stack_trace" => event.stack_trace.as_deref().unwrap_or("").to_string(),
512            _ => "".to_string(),
513        };
514
515        if let Some(formatter) = formatter {
516            self.apply_formatter(&raw_value, field, formatter)
517        } else {
518            raw_value
519        }
520    }
521
522    /// Apply formatter to value
523    fn apply_formatter(&self, value: &str, _field: &str, formatter: &CsvFormatter) -> String {
524        match formatter {
525            CsvFormatter::Duration(format) => {
526                if let Ok(duration_us) = value.parse::<u64>() {
527                    match format {
528                        DurationFormat::Microseconds => format!("{duration_us}μs"),
529                        DurationFormat::Milliseconds => {
530                            format!("{:.3}ms", duration_us as f64 / 1000.0)
531                        }
532                        DurationFormat::Seconds => {
533                            format!("{:.6}s", duration_us as f64 / 1_000_000.0)
534                        }
535                        DurationFormat::HumanReadable => {
536                            if duration_us < 1000 {
537                                format!("{duration_us}μs")
538                            } else if duration_us < 1_000_000 {
539                                format!("{:.2}ms", duration_us as f64 / 1000.0)
540                            } else {
541                                format!("{:.3}s", duration_us as f64 / 1_000_000.0)
542                            }
543                        }
544                    }
545                } else {
546                    value.to_string()
547                }
548            }
549            CsvFormatter::Memory(format) => {
550                if let Ok(bytes) = value.parse::<u64>() {
551                    match format {
552                        MemoryFormat::Bytes => format!("{bytes}B"),
553                        MemoryFormat::Kilobytes => format!("{:.2}KB", bytes as f64 / 1024.0),
554                        MemoryFormat::Megabytes => format!("{:.2}MB", bytes as f64 / 1_048_576.0),
555                        MemoryFormat::Gigabytes => {
556                            format!("{:.2}GB", bytes as f64 / 1_073_741_824.0)
557                        }
558                        MemoryFormat::HumanReadable => {
559                            if bytes < 1024 {
560                                format!("{bytes}B")
561                            } else if bytes < 1_048_576 {
562                                format!("{:.2}KB", bytes as f64 / 1024.0)
563                            } else if bytes < 1_073_741_824 {
564                                format!("{:.2}MB", bytes as f64 / 1_048_576.0)
565                            } else {
566                                format!("{:.2}GB", bytes as f64 / 1_073_741_824.0)
567                            }
568                        }
569                    }
570                } else {
571                    value.to_string()
572                }
573            }
574            CsvFormatter::Number(format) => {
575                if let Ok(num) = value.parse::<f64>() {
576                    match format {
577                        NumberFormat::Default => value.to_string(),
578                        NumberFormat::Scientific => format!("{num:.2e}"),
579                        NumberFormat::Percentage => format!("{:.2}%", num * 100.0),
580                        NumberFormat::WithCommas => {
581                            let parts: Vec<&str> = value.split('.').collect();
582                            let integer_part = parts[0];
583                            let mut formatted = String::new();
584                            for (i, c) in integer_part.chars().rev().enumerate() {
585                                if i > 0 && i % 3 == 0 {
586                                    formatted.insert(0, ',');
587                                }
588                                formatted.insert(0, c);
589                            }
590                            if parts.len() > 1 {
591                                formatted.push('.');
592                                formatted.push_str(parts[1]);
593                            }
594                            formatted
595                        }
596                    }
597                } else {
598                    value.to_string()
599                }
600            }
601            CsvFormatter::Text(format) => match format {
602                TextFormat::Default => value.to_string(),
603                TextFormat::Uppercase => value.to_uppercase(),
604                TextFormat::Lowercase => value.to_lowercase(),
605                TextFormat::Truncate(len) => {
606                    if value.len() > *len {
607                        format!("{}...", &value[..*len])
608                    } else {
609                        value.to_string()
610                    }
611                }
612            },
613        }
614    }
615
616    /// Format template string with event data
617    fn format_template(&self, template: &str, event: &ProfileEvent) -> String {
618        template
619            .replace("{name}", &event.name)
620            .replace("{category}", &event.category)
621            .replace("{start_us}", &event.start_us.to_string())
622            .replace("{duration_us}", &event.duration_us.to_string())
623            .replace("{thread_id}", &event.thread_id.to_string())
624            .replace(
625                "{operation_count}",
626                &event
627                    .operation_count
628                    .map_or("".to_string(), |v| v.to_string()),
629            )
630            .replace(
631                "{flops}",
632                &event.flops.map_or("".to_string(), |v| v.to_string()),
633            )
634            .replace(
635                "{bytes_transferred}",
636                &event
637                    .bytes_transferred
638                    .map_or("".to_string(), |v| v.to_string()),
639            )
640            .replace("{stack_trace}", event.stack_trace.as_deref().unwrap_or(""))
641    }
642}
643
644impl Default for CustomExporter {
645    fn default() -> Self {
646        Self::new()
647    }
648}
649
650#[cfg(test)]
651mod tests {
652    use super::*;
653
654    fn create_test_event() -> ProfileEvent {
655        ProfileEvent {
656            name: "test_event".to_string(),
657            category: "test".to_string(),
658            start_us: 1000,
659            duration_us: 5000,
660            thread_id: 123,
661            operation_count: Some(100),
662            flops: Some(1000000),
663            bytes_transferred: Some(1024),
664            stack_trace: None,
665        }
666    }
667
668    #[test]
669    fn test_custom_exporter_creation() {
670        let exporter = CustomExporter::new();
671        let formats = exporter.get_format_names();
672
673        assert!(formats.contains(&"compact_json".to_string()));
674        assert!(formats.contains(&"performance_csv".to_string()));
675        assert!(formats.contains(&"simple_text".to_string()));
676    }
677
678    #[test]
679    fn test_custom_json_export() {
680        let exporter = CustomExporter::new();
681        let events = vec![create_test_event()];
682
683        let compact_path = std::env::temp_dir().join("test_compact.json");
684        let compact_str = compact_path.display().to_string();
685        let result = exporter.export(&events, "compact_json", &compact_str);
686        assert!(result.is_ok());
687
688        // Clean up
689        let _ = std::fs::remove_file(&compact_path);
690    }
691
692    #[test]
693    fn test_custom_csv_export() {
694        let exporter = CustomExporter::new();
695        let events = vec![create_test_event()];
696
697        let perf_path = std::env::temp_dir().join("test_perf.csv");
698        let perf_str = perf_path.display().to_string();
699        let result = exporter.export(&events, "performance_csv", &perf_str);
700        assert!(result.is_ok());
701
702        // Clean up
703        let _ = std::fs::remove_file(&perf_path);
704    }
705
706    #[test]
707    fn test_custom_text_export() {
708        let exporter = CustomExporter::new();
709        let events = vec![create_test_event()];
710
711        let simple_path = std::env::temp_dir().join("test_simple.txt");
712        let simple_str = simple_path.display().to_string();
713        let result = exporter.export(&events, "simple_text", &simple_str);
714        assert!(result.is_ok());
715
716        // Clean up
717        let _ = std::fs::remove_file(&simple_path);
718    }
719
720    #[test]
721    fn test_duration_formatting() {
722        let exporter = CustomExporter::new();
723
724        let microseconds = exporter.apply_formatter(
725            "5000",
726            "duration_us",
727            &CsvFormatter::Duration(DurationFormat::Microseconds),
728        );
729        assert_eq!(microseconds, "5000μs");
730
731        let milliseconds = exporter.apply_formatter(
732            "5000",
733            "duration_us",
734            &CsvFormatter::Duration(DurationFormat::Milliseconds),
735        );
736        assert_eq!(milliseconds, "5.000ms");
737
738        let human_readable = exporter.apply_formatter(
739            "5000",
740            "duration_us",
741            &CsvFormatter::Duration(DurationFormat::HumanReadable),
742        );
743        assert_eq!(human_readable, "5.00ms");
744    }
745
746    #[test]
747    fn test_memory_formatting() {
748        let exporter = CustomExporter::new();
749
750        let bytes = exporter.apply_formatter(
751            "1024",
752            "bytes_transferred",
753            &CsvFormatter::Memory(MemoryFormat::Bytes),
754        );
755        assert_eq!(bytes, "1024B");
756
757        let kilobytes = exporter.apply_formatter(
758            "1024",
759            "bytes_transferred",
760            &CsvFormatter::Memory(MemoryFormat::Kilobytes),
761        );
762        assert_eq!(kilobytes, "1.00KB");
763
764        let human_readable = exporter.apply_formatter(
765            "1048576",
766            "bytes_transferred",
767            &CsvFormatter::Memory(MemoryFormat::HumanReadable),
768        );
769        assert_eq!(human_readable, "1.00MB");
770    }
771
772    #[test]
773    fn test_register_custom_format() {
774        let mut exporter = CustomExporter::new();
775
776        let custom_format = CustomExportFormat {
777            name: "test_format".to_string(),
778            description: "Test format".to_string(),
779            file_extension: "test".to_string(),
780            schema: ExportSchema::Text {
781                template: "Event: {name}".to_string(),
782                separator: " | ".to_string(),
783            },
784        };
785
786        exporter.register_format(custom_format);
787        let formats = exporter.get_format_names();
788        assert!(formats.contains(&"test_format".to_string()));
789    }
790}
torsh_profiler/custom_export.rs

torsh_profiler/
custom_export.rs
