Skip to main content

ghostscope_ui/components/command_panel/
trace_persistence.rs

1//! Trace persistence module for saving and loading trace configurations
2//!
3//! This module provides functionality to save active traces to script files
4//! and load them back, preserving their state (enabled/disabled) and full
5//! script content.
6
7use chrono::Local;
8use std::collections::HashMap;
9use std::fs;
10use std::io;
11use std::path::{Path, PathBuf};
12
13use crate::events::{TraceDefinition, TraceStatus};
14
15/// Represents a single trace configuration for persistence
16#[derive(Debug, Clone)]
17pub struct TraceConfig {
18    pub id: u32,
19    pub target: String,                // Function name or file:line
20    pub script: String,                // Full script content
21    pub status: TraceStatus,           // Active, Disabled, or Failed
22    pub binary_path: String,           // Associated binary
23    pub selected_index: Option<usize>, // Optional selected index for multi-address targets
24}
25
26/// Filter options for saving traces
27#[derive(Debug, Clone, Copy, PartialEq)]
28pub enum SaveFilter {
29    All,      // Save all traces
30    Enabled,  // Save only enabled traces
31    Disabled, // Save only disabled traces
32}
33
34/// Result of a save operation
35#[derive(Debug)]
36pub struct SaveResult {
37    pub filename: PathBuf,
38    pub saved_count: usize,
39    pub total_count: usize,
40}
41
42/// Result of a load operation
43#[derive(Debug)]
44pub struct LoadResult {
45    pub filename: PathBuf,
46    pub loaded_count: usize,
47    pub enabled_count: usize,
48    pub disabled_count: usize,
49}
50
51/// Main trace persistence handler
52pub struct TracePersistence {
53    /// Current trace configurations indexed by ID
54    traces: HashMap<u32, TraceConfig>,
55    /// Binary path for the current session
56    binary_path: Option<String>,
57    /// Process ID for the current session
58    pid: Option<u32>,
59}
60
61impl Default for TracePersistence {
62    fn default() -> Self {
63        Self::new()
64    }
65}
66
67impl TracePersistence {
68    /// Create a new trace persistence handler
69    pub fn new() -> Self {
70        Self {
71            traces: HashMap::new(),
72            binary_path: None,
73            pid: None,
74        }
75    }
76
77    /// Update binary path for the session
78    pub fn set_binary_path(&mut self, path: String) {
79        self.binary_path = Some(path);
80    }
81
82    /// Update process ID for the session
83    pub fn set_pid(&mut self, pid: u32) {
84        self.pid = Some(pid);
85    }
86
87    /// Add or update a trace configuration
88    pub fn add_trace(&mut self, config: TraceConfig) {
89        self.traces.insert(config.id, config);
90    }
91
92    /// Remove a trace configuration
93    pub fn remove_trace(&mut self, id: u32) -> Option<TraceConfig> {
94        self.traces.remove(&id)
95    }
96
97    /// Update trace status
98    pub fn update_trace_status(&mut self, id: u32, status: TraceStatus) {
99        if let Some(trace) = self.traces.get_mut(&id) {
100            trace.status = status;
101        }
102    }
103
104    /// Get all traces matching the filter
105    pub fn get_filtered_traces(&self, filter: SaveFilter) -> Vec<&TraceConfig> {
106        self.traces
107            .values()
108            .filter(|t| match filter {
109                SaveFilter::All => true,
110                SaveFilter::Enabled => matches!(t.status, TraceStatus::Active),
111                SaveFilter::Disabled => matches!(t.status, TraceStatus::Disabled),
112            })
113            .collect()
114    }
115
116    /// Save traces to a file
117    pub fn save_traces(
118        &self,
119        filename: Option<&str>,
120        filter: SaveFilter,
121    ) -> io::Result<SaveResult> {
122        // Use provided filename or generate default
123        let path = if let Some(name) = filename {
124            // Use filename exactly as provided - no extension added
125            PathBuf::from(name)
126        } else {
127            // Generate default filename with .gs extension
128            self.generate_default_filename()
129        };
130
131        // Get traces to save
132        let traces = self.get_filtered_traces(filter);
133        if traces.is_empty() {
134            return Err(io::Error::new(
135                io::ErrorKind::InvalidInput,
136                "No traces to save",
137            ));
138        }
139
140        // Generate file content
141        let content = self.generate_save_content(&traces, filter);
142
143        // Write to file
144        fs::write(&path, content)?;
145
146        Ok(SaveResult {
147            filename: path,
148            saved_count: traces.len(),
149            total_count: self.traces.len(),
150        })
151    }
152
153    /// Generate default filename with timestamp
154    fn generate_default_filename(&self) -> PathBuf {
155        let timestamp = Local::now().format("%Y%m%d_%H%M%S");
156        let binary_name = self
157            .binary_path
158            .as_ref()
159            .and_then(|p| Path::new(p).file_name())
160            .and_then(|n| n.to_str())
161            .unwrap_or("program");
162
163        PathBuf::from(format!("traces_{binary_name}_{timestamp}.gs"))
164    }
165
166    /// Generate the content for the save file
167    fn generate_save_content(&self, traces: &[&TraceConfig], filter: SaveFilter) -> String {
168        let mut content = String::new();
169
170        // Write header
171        content.push_str(&self.generate_header(traces.len(), filter));
172        content.push('\n');
173
174        // Write each trace
175        for (idx, trace) in traces.iter().enumerate() {
176            if idx > 0 {
177                content.push('\n');
178            }
179            content.push_str(&self.generate_trace_section(trace));
180        }
181
182        content
183    }
184
185    /// Generate file header with metadata
186    fn generate_header(&self, trace_count: usize, filter: SaveFilter) -> String {
187        let mut header = String::new();
188
189        // File identification
190        header.push_str("// GhostScope Trace Save File v1.0\n");
191
192        // Timestamp
193        let timestamp = Local::now().format("%Y-%m-%d %H:%M:%S");
194        header.push_str(&format!("// Generated: {timestamp}\n"));
195
196        // Binary information
197        if let Some(ref binary) = self.binary_path {
198            header.push_str(&format!("// Binary: {binary}\n"));
199        }
200
201        // PID information (if available)
202        if let Some(pid) = self.pid {
203            header.push_str(&format!("// PID: {pid}\n"));
204        }
205
206        // Filter information
207        let filter_desc = match filter {
208            SaveFilter::All => "all",
209            SaveFilter::Enabled => "enabled only",
210            SaveFilter::Disabled => "disabled only",
211        };
212        header.push_str(&format!("// Filter: {filter_desc}\n"));
213
214        // Trace count summary
215        let enabled_count = self
216            .traces
217            .values()
218            .filter(|t| matches!(t.status, TraceStatus::Active))
219            .count();
220        let disabled_count = self
221            .traces
222            .values()
223            .filter(|t| matches!(t.status, TraceStatus::Disabled))
224            .count();
225
226        header.push_str(&format!(
227            "// Traces: {trace_count} total ({enabled_count} enabled, {disabled_count} disabled)\n"
228        ));
229
230        header
231    }
232
233    /// Generate a single trace section
234    fn generate_trace_section(&self, trace: &TraceConfig) -> String {
235        let mut section = String::new();
236
237        // Section separator
238        section.push_str("// ========================================\n");
239
240        // Trace metadata
241        let status_str = match trace.status {
242            TraceStatus::Active => "ENABLED",
243            TraceStatus::Disabled => "DISABLED",
244            TraceStatus::Failed => "FAILED",
245        };
246
247        section.push_str(&format!(
248            "// Trace {}: {} [{}]\n",
249            trace.id, trace.target, status_str
250        ));
251        section.push_str(&format!("// Target: {}\n", trace.target));
252        section.push_str(&format!("// Status: {}\n", trace.status));
253        if let Some(idx) = trace.selected_index {
254            section.push_str(&format!("// Index: {idx}\n"));
255        }
256        section.push_str("// ========================================\n");
257
258        // Add disabled marker if needed
259        if matches!(trace.status, TraceStatus::Disabled) {
260            section.push_str("//@disabled\n");
261        }
262
263        // Trace command and script
264        let script_block = Self::format_trace_block(&trace.script, &trace.target);
265        section.push_str(&script_block);
266
267        section
268    }
269
270    /// Wrap raw script body with a trace header/brace pair
271    fn wrap_script_body(target: &str, body: &str) -> String {
272        let mut wrapped = String::new();
273        wrapped.push_str(&format!("trace {target} {{\n"));
274
275        if body.trim().is_empty() {
276            wrapped.push_str("}\n");
277            return wrapped;
278        }
279
280        for line in body.lines() {
281            wrapped.push_str("    ");
282            wrapped.push_str(line);
283            wrapped.push('\n');
284        }
285
286        wrapped.push_str("}\n");
287        wrapped
288    }
289
290    /// Normalize script content into canonical trace block format
291    fn format_trace_block(script: &str, target: &str) -> String {
292        if let Some(body) = Self::extract_trace_body(script) {
293            let dedented = Self::dedent_body(&body);
294            let trimmed = dedented.trim_end_matches(['\r', '\n']);
295            Self::wrap_script_body(target, trimmed)
296        } else {
297            Self::wrap_script_body(target, script)
298        }
299    }
300
301    /// Extract the body of an existing trace block, tolerating inline braces
302    fn extract_trace_body(script: &str) -> Option<String> {
303        let trimmed = script.trim();
304        if !trimmed.starts_with("trace ") {
305            return None;
306        }
307
308        let bytes = trimmed.as_bytes();
309        let mut start_brace = None;
310        for (idx, &b) in bytes.iter().enumerate() {
311            if b == b'{' {
312                start_brace = Some(idx);
313                break;
314            }
315        }
316        let start = start_brace?;
317        let mut depth = 1usize;
318        let mut i = start + 1;
319        while i < bytes.len() {
320            match bytes[i] {
321                b'{' => depth += 1,
322                b'}' => {
323                    depth -= 1;
324                    if depth == 0 {
325                        let raw_body = &trimmed[start + 1..i];
326                        let normalized = Self::trim_wrapped_body(raw_body);
327                        return Some(normalized.to_string());
328                    }
329                }
330                _ => {}
331            }
332            i += 1;
333        }
334        None
335    }
336
337    /// Trim surrounding whitespace/newlines around an extracted body
338    fn trim_wrapped_body(body: &str) -> &str {
339        let mut slice = body.trim_end_matches(['\r', '\n', ' ', '\t']);
340        loop {
341            if slice.starts_with("\r\n") {
342                slice = &slice[2..];
343            } else if slice.starts_with('\n') || slice.starts_with('\r') {
344                slice = &slice[1..];
345            } else {
346                break;
347            }
348        }
349        slice
350    }
351
352    /// Remove common indentation so we can re-indent consistently in the save file
353    fn dedent_body(body: &str) -> String {
354        let lines: Vec<&str> = body.lines().collect();
355        let indent = lines
356            .iter()
357            .filter_map(|line| {
358                let trimmed = line.trim();
359                if trimmed.is_empty() {
360                    None
361                } else {
362                    Some(
363                        line.as_bytes()
364                            .iter()
365                            .take_while(|&&b| b == b' ' || b == b'\t')
366                            .count(),
367                    )
368                }
369            })
370            .min()
371            .unwrap_or(0);
372
373        if indent == 0 {
374            return body.to_string();
375        }
376
377        let mut result = String::new();
378        for (idx, line) in lines.iter().enumerate() {
379            let line = *line;
380            if idx > 0 {
381                result.push('\n');
382            }
383            if line.trim().is_empty() {
384                continue;
385            }
386            let skip = indent.min(line.len());
387            let content = line.get(skip..).unwrap_or("");
388            result.push_str(content.trim_end_matches('\r'));
389        }
390
391        result
392    }
393
394    /// Parse a saved trace file for loading
395    pub fn parse_trace_file(content: &str) -> io::Result<Vec<TraceDefinition>> {
396        let mut traces = Vec::new();
397        let mut current_target: Option<String> = None;
398        let mut in_script = false;
399        let mut script_lines = Vec::new();
400        let mut pending_disabled = false;
401        let mut pending_index: Option<usize> = None;
402        // Track nested braces so inner blocks (e.g., if { ... }) don't terminate the trace section
403        let mut brace_depth: usize = 0;
404
405        for line in content.lines() {
406            let trimmed = line.trim();
407
408            // Check for disabled marker
409            if trimmed == "//@disabled" {
410                pending_disabled = true;
411                continue;
412            }
413
414            // Parse optional index metadata line (e.g., "// Index: 3")
415            if let Some(rest) = trimmed.strip_prefix("// Index:") {
416                let val = rest.trim();
417                if let Ok(idx) = val.parse::<usize>() {
418                    pending_index = Some(idx);
419                }
420                continue;
421            }
422
423            // Check for trace command start
424            if trimmed.starts_with("trace ") && trimmed.ends_with(" {") {
425                // Extract target from trace command
426                let target = trimmed
427                    .strip_prefix("trace ")
428                    .and_then(|s| s.strip_suffix(" {"))
429                    .unwrap_or("")
430                    .to_string();
431
432                current_target = Some(target);
433                in_script = true;
434                script_lines.clear();
435                // Opening brace for the trace section
436                brace_depth = 1;
437                continue;
438            }
439
440            // Check for script end: only close when this '}' matches the outer trace block
441            if in_script && trimmed == "}" && brace_depth == 1 {
442                if let Some(target) = current_target.take() {
443                    let script = script_lines.join("\n");
444                    traces.push(TraceDefinition {
445                        target,
446                        script,
447                        enabled: !pending_disabled,
448                        selected_index: pending_index,
449                    });
450                    pending_disabled = false;
451                    pending_index = None;
452                }
453                in_script = false;
454                brace_depth = 0;
455                continue;
456            }
457
458            // Collect script lines
459            if in_script {
460                // Remove leading indentation (4 spaces)
461                let script_line = if let Some(stripped) = line.strip_prefix("    ") {
462                    stripped
463                } else {
464                    line
465                };
466                script_lines.push(script_line.to_string());
467
468                // Update brace depth based on current line content so nested '}' are preserved
469                // Note: naïve count, acceptable because braces rarely appear in string literals in our scripts
470                let opens = script_line.chars().filter(|&c| c == '{').count();
471                let closes = script_line.chars().filter(|&c| c == '}').count();
472                // Saturating arithmetic to avoid underflow on malformed input
473                brace_depth = brace_depth.saturating_add(opens).saturating_sub(closes);
474            }
475        }
476
477        Ok(traces)
478    }
479
480    /// Load traces from a file
481    pub fn load_traces_from_file(filename: &str) -> io::Result<Vec<TraceDefinition>> {
482        let content = fs::read_to_string(filename)?;
483        Self::parse_trace_file(&content)
484    }
485}
486
487/// Extension trait for command parsing
488pub trait CommandParser {
489    fn parse_save_traces_command(&self) -> Option<(Option<String>, SaveFilter)>;
490}
491
492impl CommandParser for str {
493    fn parse_save_traces_command(&self) -> Option<(Option<String>, SaveFilter)> {
494        let parts: Vec<&str> = self.split_whitespace().collect();
495
496        if parts.len() < 2 || parts[0] != "save" || parts[1] != "traces" {
497            return None;
498        }
499
500        match parts.len() {
501            2 => {
502                // save traces
503                Some((None, SaveFilter::All))
504            }
505            3 => {
506                // save traces <filename> or save traces enabled/disabled
507                match parts[2] {
508                    "enabled" => Some((None, SaveFilter::Enabled)),
509                    "disabled" => Some((None, SaveFilter::Disabled)),
510                    filename => Some((Some(filename.to_string()), SaveFilter::All)),
511                }
512            }
513            4 => {
514                // save traces enabled/disabled <filename>
515                let filter = match parts[2] {
516                    "enabled" => SaveFilter::Enabled,
517                    "disabled" => SaveFilter::Disabled,
518                    _ => return None,
519                };
520                Some((Some(parts[3].to_string()), filter))
521            }
522            _ => None,
523        }
524    }
525}
526
527#[cfg(test)]
528mod tests {
529    use super::*;
530
531    #[test]
532    fn test_parse_save_command() {
533        // Basic save
534        let (file, filter) = "save traces".parse_save_traces_command().unwrap();
535        assert_eq!(file, None);
536        assert_eq!(filter, SaveFilter::All);
537
538        // Save with filename
539        let (file, filter) = "save traces session.gs"
540            .parse_save_traces_command()
541            .unwrap();
542        assert_eq!(file, Some("session.gs".to_string()));
543        assert_eq!(filter, SaveFilter::All);
544
545        // Save enabled only
546        let (file, filter) = "save traces enabled".parse_save_traces_command().unwrap();
547        assert_eq!(file, None);
548        assert_eq!(filter, SaveFilter::Enabled);
549
550        // Save disabled with filename
551        let (file, filter) = "save traces disabled debug.gs"
552            .parse_save_traces_command()
553            .unwrap();
554        assert_eq!(file, Some("debug.gs".to_string()));
555        assert_eq!(filter, SaveFilter::Disabled);
556    }
557
558    #[test]
559    fn test_parse_trace_file() {
560        let content = r#"// Header
561//@disabled
562trace main {
563    print "hello";
564    print "world";
565}
566
567trace foo {
568    print "foo";
569}"#;
570
571        let traces = TracePersistence::parse_trace_file(content).unwrap();
572        assert_eq!(traces.len(), 2);
573
574        assert_eq!(traces[0].target, "main");
575        assert!(!traces[0].enabled); // disabled trace
576        assert_eq!(traces[0].script, "print \"hello\";\nprint \"world\";");
577
578        assert_eq!(traces[1].target, "foo");
579        assert!(traces[1].enabled); // enabled trace
580        assert_eq!(traces[1].script, "print \"foo\";");
581    }
582
583    #[test]
584    fn test_save_traces_avoids_double_wrapping() {
585        use std::time::{SystemTime, UNIX_EPOCH};
586
587        let mut persistence = TracePersistence::new();
588        persistence.add_trace(TraceConfig {
589            id: 1,
590            target: "main".to_string(),
591            script: "trace main {\n    print \"hello\";\n}".to_string(),
592            status: TraceStatus::Active,
593            binary_path: "/bin/app".to_string(),
594            selected_index: None,
595        });
596
597        let filename = std::env::temp_dir().join(format!(
598            "ghostscope_trace_test_{}.gs",
599            SystemTime::now()
600                .duration_since(UNIX_EPOCH)
601                .unwrap()
602                .as_nanos()
603        ));
604        let filename_str = filename.to_string_lossy().to_string();
605
606        let result = persistence
607            .save_traces(Some(&filename_str), SaveFilter::All)
608            .expect("save traces succeeds");
609
610        let saved = std::fs::read_to_string(&result.filename).expect("saved trace file readable");
611
612        assert!(
613            saved.contains("trace main {\n    print \"hello\";\n}\n"),
614            "saved trace missing expected block:\n{saved}"
615        );
616        assert!(
617            !saved.contains("trace main {\n    trace main"),
618            "trace block was double wrapped:\n{saved}"
619        );
620
621        let parsed = TracePersistence::parse_trace_file(&saved).expect("saved file parses");
622        assert_eq!(parsed.len(), 1);
623        assert_eq!(parsed[0].target, "main");
624        assert_eq!(parsed[0].script, "print \"hello\";");
625
626        let _ = std::fs::remove_file(&result.filename);
627    }
628}