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        section.push_str(&format!("trace {} {{\n", trace.target));
265
266        // Indent script content
267        for line in trace.script.lines() {
268            section.push_str("    ");
269            section.push_str(line);
270            section.push('\n');
271        }
272
273        section.push_str("}\n");
274
275        section
276    }
277
278    /// Parse a saved trace file for loading
279    pub fn parse_trace_file(content: &str) -> io::Result<Vec<TraceDefinition>> {
280        let mut traces = Vec::new();
281        let mut current_target: Option<String> = None;
282        let mut in_script = false;
283        let mut script_lines = Vec::new();
284        let mut pending_disabled = false;
285        let mut pending_index: Option<usize> = None;
286        // Track nested braces so inner blocks (e.g., if { ... }) don't terminate the trace section
287        let mut brace_depth: usize = 0;
288
289        for line in content.lines() {
290            let trimmed = line.trim();
291
292            // Check for disabled marker
293            if trimmed == "//@disabled" {
294                pending_disabled = true;
295                continue;
296            }
297
298            // Parse optional index metadata line (e.g., "// Index: 3")
299            if let Some(rest) = trimmed.strip_prefix("// Index:") {
300                let val = rest.trim();
301                if let Ok(idx) = val.parse::<usize>() {
302                    pending_index = Some(idx);
303                }
304                continue;
305            }
306
307            // Check for trace command start
308            if trimmed.starts_with("trace ") && trimmed.ends_with(" {") {
309                // Extract target from trace command
310                let target = trimmed
311                    .strip_prefix("trace ")
312                    .and_then(|s| s.strip_suffix(" {"))
313                    .unwrap_or("")
314                    .to_string();
315
316                current_target = Some(target);
317                in_script = true;
318                script_lines.clear();
319                // Opening brace for the trace section
320                brace_depth = 1;
321                continue;
322            }
323
324            // Check for script end: only close when this '}' matches the outer trace block
325            if in_script && trimmed == "}" && brace_depth == 1 {
326                if let Some(target) = current_target.take() {
327                    let script = script_lines.join("\n");
328                    traces.push(TraceDefinition {
329                        target,
330                        script,
331                        enabled: !pending_disabled,
332                        selected_index: pending_index,
333                    });
334                    pending_disabled = false;
335                    pending_index = None;
336                }
337                in_script = false;
338                brace_depth = 0;
339                continue;
340            }
341
342            // Collect script lines
343            if in_script {
344                // Remove leading indentation (4 spaces)
345                let script_line = if let Some(stripped) = line.strip_prefix("    ") {
346                    stripped
347                } else {
348                    line
349                };
350                script_lines.push(script_line.to_string());
351
352                // Update brace depth based on current line content so nested '}' are preserved
353                // Note: naïve count, acceptable because braces rarely appear in string literals in our scripts
354                let opens = script_line.chars().filter(|&c| c == '{').count();
355                let closes = script_line.chars().filter(|&c| c == '}').count();
356                // Saturating arithmetic to avoid underflow on malformed input
357                brace_depth = brace_depth.saturating_add(opens).saturating_sub(closes);
358            }
359        }
360
361        Ok(traces)
362    }
363
364    /// Load traces from a file
365    pub fn load_traces_from_file(filename: &str) -> io::Result<Vec<TraceDefinition>> {
366        let content = fs::read_to_string(filename)?;
367        Self::parse_trace_file(&content)
368    }
369}
370
371/// Extension trait for command parsing
372pub trait CommandParser {
373    fn parse_save_traces_command(&self) -> Option<(Option<String>, SaveFilter)>;
374}
375
376impl CommandParser for str {
377    fn parse_save_traces_command(&self) -> Option<(Option<String>, SaveFilter)> {
378        let parts: Vec<&str> = self.split_whitespace().collect();
379
380        if parts.len() < 2 || parts[0] != "save" || parts[1] != "traces" {
381            return None;
382        }
383
384        match parts.len() {
385            2 => {
386                // save traces
387                Some((None, SaveFilter::All))
388            }
389            3 => {
390                // save traces <filename> or save traces enabled/disabled
391                match parts[2] {
392                    "enabled" => Some((None, SaveFilter::Enabled)),
393                    "disabled" => Some((None, SaveFilter::Disabled)),
394                    filename => Some((Some(filename.to_string()), SaveFilter::All)),
395                }
396            }
397            4 => {
398                // save traces enabled/disabled <filename>
399                let filter = match parts[2] {
400                    "enabled" => SaveFilter::Enabled,
401                    "disabled" => SaveFilter::Disabled,
402                    _ => return None,
403                };
404                Some((Some(parts[3].to_string()), filter))
405            }
406            _ => None,
407        }
408    }
409}
410
411#[cfg(test)]
412mod tests {
413    use super::*;
414
415    #[test]
416    fn test_parse_save_command() {
417        // Basic save
418        let (file, filter) = "save traces".parse_save_traces_command().unwrap();
419        assert_eq!(file, None);
420        assert_eq!(filter, SaveFilter::All);
421
422        // Save with filename
423        let (file, filter) = "save traces session.gs"
424            .parse_save_traces_command()
425            .unwrap();
426        assert_eq!(file, Some("session.gs".to_string()));
427        assert_eq!(filter, SaveFilter::All);
428
429        // Save enabled only
430        let (file, filter) = "save traces enabled".parse_save_traces_command().unwrap();
431        assert_eq!(file, None);
432        assert_eq!(filter, SaveFilter::Enabled);
433
434        // Save disabled with filename
435        let (file, filter) = "save traces disabled debug.gs"
436            .parse_save_traces_command()
437            .unwrap();
438        assert_eq!(file, Some("debug.gs".to_string()));
439        assert_eq!(filter, SaveFilter::Disabled);
440    }
441
442    #[test]
443    fn test_parse_trace_file() {
444        let content = r#"// Header
445//@disabled
446trace main {
447    print "hello";
448    print "world";
449}
450
451trace foo {
452    print "foo";
453}"#;
454
455        let traces = TracePersistence::parse_trace_file(content).unwrap();
456        assert_eq!(traces.len(), 2);
457
458        assert_eq!(traces[0].target, "main");
459        assert!(!traces[0].enabled); // disabled trace
460        assert_eq!(traces[0].script, "print \"hello\";\nprint \"world\";");
461
462        assert_eq!(traces[1].target, "foo");
463        assert!(traces[1].enabled); // enabled trace
464        assert_eq!(traces[1].script, "print \"foo\";");
465    }
466}