baobao_codegen/pipeline/
snapshot.rs

1//! Pipeline snapshot plugin for visualization and debugging.
2//!
3//! This module provides a plugin that captures the pipeline state after each phase,
4//! enabling visualization of the compilation process.
5
6use std::{
7    fs,
8    path::{Path, PathBuf},
9    sync::RwLock,
10};
11
12use eyre::Result;
13use serde::Serialize;
14
15use super::{CompilationContext, Diagnostic, Plugin};
16use crate::schema::ComputedData;
17
18/// A snapshot of the pipeline state at a specific phase.
19#[derive(Debug, Clone, Serialize)]
20pub struct PhaseSnapshot {
21    /// The phase that just completed.
22    pub phase: String,
23
24    /// The Application IR (available after "lower" phase).
25    #[serde(skip_serializing_if = "Option::is_none")]
26    pub ir: Option<baobao_ir::AppIR>,
27
28    /// Pre-computed analysis data (available after "analyze" phase).
29    #[serde(skip_serializing_if = "Option::is_none")]
30    pub computed: Option<ComputedData>,
31
32    /// Diagnostics collected so far.
33    pub diagnostics: Vec<Diagnostic>,
34}
35
36/// A plugin that captures pipeline state after each phase.
37///
38/// Use this plugin with the `--visualize` flag to output intermediate
39/// representations for debugging and understanding the pipeline.
40///
41/// # Example
42///
43/// ```ignore
44/// let snapshot_plugin = SnapshotPlugin::new();
45/// let pipeline = Pipeline::new().plugin(snapshot_plugin.clone());
46/// let ctx = pipeline.run(manifest)?;
47///
48/// // Write snapshots to disk
49/// snapshot_plugin.write_to_dir(".bao/debug")?;
50/// ```
51pub struct SnapshotPlugin {
52    /// Collected snapshots.
53    snapshots: RwLock<Vec<PhaseSnapshot>>,
54    /// Output directory for snapshots.
55    output_dir: Option<PathBuf>,
56}
57
58impl SnapshotPlugin {
59    /// Create a new snapshot plugin.
60    pub fn new() -> Self {
61        Self {
62            snapshots: RwLock::new(Vec::new()),
63            output_dir: None,
64        }
65    }
66
67    /// Create a new snapshot plugin that writes to a directory.
68    pub fn with_output_dir(output_dir: impl Into<PathBuf>) -> Self {
69        Self {
70            snapshots: RwLock::new(Vec::new()),
71            output_dir: Some(output_dir.into()),
72        }
73    }
74
75    /// Get all collected snapshots.
76    pub fn snapshots(&self) -> Vec<PhaseSnapshot> {
77        self.snapshots.read().unwrap().clone()
78    }
79
80    /// Write all snapshots to the configured output directory.
81    pub fn write_to_dir(&self, dir: impl AsRef<Path>) -> Result<()> {
82        let dir = dir.as_ref();
83        fs::create_dir_all(dir)?;
84
85        for snapshot in self.snapshots.read().unwrap().iter() {
86            let filename = format!("{}.json", snapshot.phase);
87            let path = dir.join(&filename);
88            let json = serde_json::to_string_pretty(snapshot)?;
89            fs::write(&path, json)?;
90        }
91
92        Ok(())
93    }
94
95    fn capture_snapshot(&self, phase: &str, ctx: &CompilationContext) {
96        let snapshot = PhaseSnapshot {
97            phase: phase.to_string(),
98            ir: ctx.ir.clone(),
99            computed: ctx.computed.clone(),
100            diagnostics: ctx.diagnostics.clone(),
101        };
102        self.snapshots.write().unwrap().push(snapshot);
103    }
104}
105
106impl Default for SnapshotPlugin {
107    fn default() -> Self {
108        Self::new()
109    }
110}
111
112impl Plugin for SnapshotPlugin {
113    fn name(&self) -> &'static str {
114        "snapshot"
115    }
116
117    fn on_after_phase(&self, phase: &str, ctx: &mut CompilationContext) -> Result<()> {
118        self.capture_snapshot(phase, ctx);
119
120        // If output directory is configured, write immediately
121        if let Some(ref dir) = self.output_dir {
122            let filename = format!("{}.json", phase);
123            let path = dir.join(&filename);
124
125            if let Some(snapshot) = self.snapshots.read().unwrap().last() {
126                fs::create_dir_all(dir)?;
127                let json = serde_json::to_string_pretty(snapshot)?;
128                fs::write(&path, json)?;
129            }
130        }
131
132        Ok(())
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn test_snapshot_plugin_creation() {
142        let plugin = SnapshotPlugin::new();
143        assert!(plugin.snapshots().is_empty());
144    }
145}
baobao_codegen/pipeline/snapshot.rs

baobao_codegen/pipeline/
snapshot.rs
