dfir_lang/graph/
graph_write.rs

Help
1#![warn(missing_docs)]
2
3use std::borrow::Cow;
4use std::error::Error;
5
6use auto_impl::auto_impl;
7use slotmap::Key;
8
9use super::ops::DelayType;
10use super::{Color, GraphNodeId, GraphSubgraphId};
11
12/// Trait for writing textual representations of graphs, i.e. mermaid or dot graphs.
13#[auto_impl(&mut, Box)]
14pub(crate) trait GraphWrite {
15    /// Error type emitted by writing.
16    type Err: Error;
17
18    /// Begin the graph. First method called.
19    fn write_prologue(&mut self) -> Result<(), Self::Err>;
20
21    /// Write a node, with styling.
22    fn write_node(
23        &mut self,
24        node_id: GraphNodeId,
25        node: &str,
26        node_color: Option<Color>,
27    ) -> Result<(), Self::Err>;
28
29    /// Write an edge, with styling.
30    fn write_edge(
31        &mut self,
32        src_id: GraphNodeId,
33        dst_id: GraphNodeId,
34        delay_type: Option<DelayType>,
35        label: Option<&str>,
36        is_reference: bool,
37    ) -> Result<(), Self::Err>;
38
39    /// Begin writing a subgraph.
40    fn write_subgraph_start(
41        &mut self,
42        sg_id: GraphSubgraphId,
43        stratum: usize,
44        subgraph_nodes: impl Iterator<Item = GraphNodeId>,
45    ) -> Result<(), Self::Err>;
46    /// Write the nodes associated with a single variable name, within a subgraph.
47    fn write_varname(
48        &mut self,
49        varname: &str,
50        varname_nodes: impl Iterator<Item = GraphNodeId>,
51        sg_id: Option<GraphSubgraphId>,
52    ) -> Result<(), Self::Err>;
53    /// End writing a subgraph.
54    fn write_subgraph_end(&mut self) -> Result<(), Self::Err>;
55
56    /// End the graph. Last method called.
57    fn write_epilogue(&mut self) -> Result<(), Self::Err>;
58}
59
60/// Escapes a string for use in a mermaid graph label.
61pub fn escape_mermaid(string: &str) -> String {
62    string
63        .replace('&', "&amp;")
64        .replace('<', "&lt;")
65        .replace('>', "&gt;")
66        .replace('"', "&quot;")
67        // Mermaid entity codes
68        // https://mermaid.js.org/syntax/flowchart.html#entity-codes-to-escape-characters
69        .replace('#', "&num;")
70        // Not really needed, newline literals seem to work
71        .replace('\n', "<br>")
72        // Mermaid font awesome fa
73        // https://github.com/mermaid-js/mermaid/blob/e4d2118d4bfa023628a020b7ab1f8c491e6dc523/packages/mermaid/src/diagrams/flowchart/flowRenderer-v2.js#L62
74        .replace("fa:fa", "fa:<wbr>fa")
75        .replace("fab:fa", "fab:<wbr>fa")
76        .replace("fal:fa", "fal:<wbr>fa")
77        .replace("far:fa", "far:<wbr>fa")
78        .replace("fas:fa", "fas:<wbr>fa")
79}
80
81pub struct Mermaid<W> {
82    write: W,
83    // How many links have been written, for styling
84    // https://mermaid.js.org/syntax/flowchart.html#styling-links
85    link_count: usize,
86}
87impl<W> Mermaid<W> {
88    pub fn new(write: W) -> Self {
89        Self {
90            write,
91            link_count: 0,
92        }
93    }
94}
95impl<W> GraphWrite for Mermaid<W>
96where
97    W: std::fmt::Write,
98{
99    type Err = std::fmt::Error;
100
101    fn write_prologue(&mut self) -> Result<(), Self::Err> {
102        writeln!(
103            self.write,
104            r"%%{{init:{{'theme':'base','themeVariables':{{'clusterBkg':'#ddd','clusterBorder':'#888'}}}}}}%%",
105        )?;
106        writeln!(self.write, "flowchart TD")?;
107        writeln!(
108            self.write,
109            "classDef pullClass fill:#8af,stroke:#000,text-align:left,white-space:pre",
110        )?;
111        writeln!(
112            self.write,
113            "classDef pushClass fill:#ff8,stroke:#000,text-align:left,white-space:pre",
114        )?;
115        writeln!(
116            self.write,
117            "classDef otherClass fill:#fdc,stroke:#000,text-align:left,white-space:pre",
118        )?;
119
120        writeln!(self.write, "linkStyle default stroke:#aaa")?;
121        Ok(())
122    }
123
124    fn write_node(
125        &mut self,
126        node_id: GraphNodeId,
127        node: &str,
128        node_color: Option<Color>,
129    ) -> Result<(), Self::Err> {
130        let class_str = match node_color {
131            Some(Color::Push) => "pushClass",
132            Some(Color::Pull) => "pullClass",
133            _ => "otherClass",
134        };
135        let label = format!(
136            r#"{node_id:?}{lbracket}"{node_label} <code>{code}</code>"{rbracket}:::{class}"#,
137            node_id = node_id.data(),
138            node_label = if node.contains('\n') {
139                format!("<div style=text-align:center>({:?})</div>", node_id.data())
140            } else {
141                format!("({:?})", node_id.data())
142            },
143            class = class_str,
144            lbracket = match node_color {
145                Some(Color::Push) => r"[/",
146                Some(Color::Pull) => r"[\",
147                _ => "[",
148            },
149            code = escape_mermaid(node),
150            rbracket = match node_color {
151                Some(Color::Push) => r"\]",
152                Some(Color::Pull) => r"/]",
153                _ => "]",
154            },
155        );
156        writeln!(self.write, "{}", label)?;
157        Ok(())
158    }
159
160    fn write_edge(
161        &mut self,
162        src_id: GraphNodeId,
163        dst_id: GraphNodeId,
164        delay_type: Option<DelayType>,
165        label: Option<&str>,
166        _is_reference: bool,
167    ) -> Result<(), Self::Err> {
168        let src_str = format!("{:?}", src_id.data());
169        let dest_str = format!("{:?}", dst_id.data());
170        #[expect(clippy::write_literal, reason = "code readability")]
171        write!(
172            self.write,
173            "{src}{arrow_body}{arrow_head}{label}{dst}",
174            src = src_str.trim(),
175            arrow_body = "--",
176            arrow_head = match delay_type {
177                None | Some(DelayType::MonotoneAccum) => ">",
178                Some(DelayType::Stratum) => "x",
179                Some(DelayType::Tick | DelayType::TickLazy) => "o",
180            },
181            label = if let Some(label) = &label {
182                Cow::Owned(format!("|{}|", escape_mermaid(label.trim())))
183            } else {
184                Cow::Borrowed("")
185            },
186            dst = dest_str.trim(),
187        )?;
188        if let Some(delay_type) = delay_type {
189            write!(
190                self.write,
191                "; linkStyle {} stroke:{}",
192                self.link_count,
193                match delay_type {
194                    DelayType::Stratum | DelayType::Tick | DelayType::TickLazy => "red",
195                    DelayType::MonotoneAccum => "#060",
196                }
197            )?;
198        }
199        writeln!(self.write)?;
200        self.link_count += 1;
201        Ok(())
202    }
203
204    fn write_subgraph_start(
205        &mut self,
206        sg_id: GraphSubgraphId,
207        stratum: usize,
208        subgraph_nodes: impl Iterator<Item = GraphNodeId>,
209    ) -> Result<(), Self::Err> {
210        writeln!(
211            self.write,
212            "subgraph sg_{sg:?} [\"sg_{sg:?} stratum {:?}\"]",
213            stratum,
214            sg = sg_id.data(),
215        )?;
216        for node_id in subgraph_nodes {
217            writeln!(self.write, "    {node_id:?}", node_id = node_id.data())?;
218        }
219        Ok(())
220    }
221
222    fn write_varname(
223        &mut self,
224        varname: &str,
225        varname_nodes: impl Iterator<Item = GraphNodeId>,
226        sg_id: Option<GraphSubgraphId>,
227    ) -> Result<(), Self::Err> {
228        let pad = if let Some(sg_id) = sg_id {
229            writeln!(
230                self.write,
231                "    subgraph sg_{sg:?}_var_{var} [\"var <tt>{var}</tt>\"]",
232                sg = sg_id.data(),
233                var = varname,
234            )?;
235            "    "
236        } else {
237            writeln!(
238                self.write,
239                "subgraph var_{0} [\"var <tt>{0}</tt>\"]",
240                varname,
241            )?;
242            writeln!(self.write, "style var_{} fill:transparent", varname)?;
243            ""
244        };
245        for local_named_node in varname_nodes {
246            writeln!(self.write, "    {}{:?}", pad, local_named_node.data())?;
247        }
248        writeln!(self.write, "{}end", pad)?;
249        Ok(())
250    }
251
252    fn write_subgraph_end(&mut self) -> Result<(), Self::Err> {
253        writeln!(self.write, "end")?;
254        Ok(())
255    }
256
257    fn write_epilogue(&mut self) -> Result<(), Self::Err> {
258        // No-op.
259        Ok(())
260    }
261}
262
263/// Escapes a string for use in a DOT graph label.
264///
265/// Newline can be:
266/// * "\\n" for newline.
267/// * "\\l" for left-aligned newline.
268/// * "\\r" for right-aligned newline.
269pub fn escape_dot(string: &str, newline: &str) -> String {
270    string.replace('"', "\\\"").replace('\n', newline)
271}
272
273pub struct Dot<W> {
274    write: W,
275}
276impl<W> Dot<W> {
277    pub fn new(write: W) -> Self {
278        Self { write }
279    }
280}
281impl<W> GraphWrite for Dot<W>
282where
283    W: std::fmt::Write,
284{
285    type Err = std::fmt::Error;
286
287    fn write_prologue(&mut self) -> Result<(), Self::Err> {
288        writeln!(self.write, "digraph {{")?;
289        const FONTS: &str = "\"Monaco,Menlo,Consolas,&quot;Droid Sans Mono&quot;,Inconsolata,&quot;Courier New&quot;,monospace\"";
290        writeln!(self.write, "    node [fontname={}, style=filled];", FONTS)?;
291        writeln!(self.write, "    edge [fontname={}];", FONTS)?;
292        Ok(())
293    }
294
295    fn write_node(
296        &mut self,
297        node_id: GraphNodeId,
298        node: &str,
299        node_color: Option<Color>,
300    ) -> Result<(), Self::Err> {
301        let nm = escape_dot(node, "\\l");
302        let label = format!("n{:?}", node_id.data());
303        let shape_str = match node_color {
304            Some(Color::Push) => "house",
305            Some(Color::Pull) => "invhouse",
306            Some(Color::Hoff) => "parallelogram",
307            Some(Color::Comp) => "circle",
308            None => "rectangle",
309        };
310        let color_str = match node_color {
311            Some(Color::Push) => "\"#ffff88\"",
312            Some(Color::Pull) => "\"#88aaff\"",
313            Some(Color::Hoff) => "\"#ddddff\"",
314            Some(Color::Comp) => "white",
315            None => "\"#ddddff\"",
316        };
317        write!(
318            self.write,
319            "    {} [label=\"({}) {}{}\"",
320            label,
321            label,
322            nm,
323            // if contains linebreak left-justify by appending another "\\l"
324            if nm.contains("\\l") { "\\l" } else { "" },
325        )?;
326        write!(self.write, ", shape={}, fillcolor={}", shape_str, color_str)?;
327        writeln!(self.write, "]")?;
328        Ok(())
329    }
330
331    fn write_edge(
332        &mut self,
333        src_id: GraphNodeId,
334        dst_id: GraphNodeId,
335        delay_type: Option<DelayType>,
336        label: Option<&str>,
337        _is_reference: bool,
338    ) -> Result<(), Self::Err> {
339        let mut properties = Vec::<Cow<'static, str>>::new();
340        if let Some(label) = label {
341            properties.push(format!("label=\"{}\"", escape_dot(label, "\\n")).into());
342        };
343        // Color
344        if delay_type.is_some() {
345            properties.push("color=red".into());
346        }
347
348        write!(
349            self.write,
350            "    n{:?} -> n{:?}",
351            src_id.data(),
352            dst_id.data(),
353        )?;
354        if !properties.is_empty() {
355            write!(self.write, " [")?;
356            for prop in itertools::Itertools::intersperse(properties.into_iter(), ", ".into()) {
357                write!(self.write, "{}", prop)?;
358            }
359            write!(self.write, "]")?;
360        }
361        writeln!(self.write)?;
362        Ok(())
363    }
364
365    fn write_subgraph_start(
366        &mut self,
367        sg_id: GraphSubgraphId,
368        stratum: usize,
369        subgraph_nodes: impl Iterator<Item = GraphNodeId>,
370    ) -> Result<(), Self::Err> {
371        writeln!(
372            self.write,
373            "    subgraph \"cluster n{:?}\" {{",
374            sg_id.data(),
375        )?;
376        writeln!(self.write, "        fillcolor=\"#dddddd\"")?;
377        writeln!(self.write, "        style=filled")?;
378        writeln!(
379            self.write,
380            "        label = \"sg_{:?}\\nstratum {}\"",
381            sg_id.data(),
382            stratum,
383        )?;
384        for node_id in subgraph_nodes {
385            writeln!(self.write, "        n{:?}", node_id.data(),)?;
386        }
387        Ok(())
388    }
389
390    fn write_varname(
391        &mut self,
392        varname: &str,
393        varname_nodes: impl Iterator<Item = GraphNodeId>,
394        sg_id: Option<GraphSubgraphId>,
395    ) -> Result<(), Self::Err> {
396        let pad = if let Some(sg_id) = sg_id {
397            writeln!(
398                self.write,
399                "        subgraph \"cluster_sg_{sg:?}_var_{var}\" {{",
400                sg = sg_id.data(),
401                var = varname,
402            )?;
403            "    "
404        } else {
405            writeln!(
406                self.write,
407                "    subgraph \"cluster_var_{var}\" {{",
408                var = varname,
409            )?;
410            ""
411        };
412        writeln!(
413            self.write,
414            "        {}label=\"var {var}\"",
415            pad,
416            var = varname
417        )?;
418        for local_named_node in varname_nodes {
419            writeln!(self.write, "        {}n{:?}", pad, local_named_node.data())?;
420        }
421        writeln!(self.write, "    {}}}", pad)?;
422        Ok(())
423    }
424
425    fn write_subgraph_end(&mut self) -> Result<(), Self::Err> {
426        // subgraph footer
427        writeln!(self.write, "    }}")?;
428        Ok(())
429    }
430
431    fn write_epilogue(&mut self) -> Result<(), Self::Err> {
432        writeln!(self.write, "}}")?;
433        Ok(())
434    }
435}
dfir_lang/graph/graph_write.rs

dfir_lang/graph/
graph_write.rs
