laddu_expr/
visualization.rs

1use std::{collections::HashSet, fmt};
2
3use crate::{ExprGraph, ExprId, ExprMetadata, ExprNode, expression::node_children};
4
5/// Controls how graph displays handle nodes reached through multiple paths.
6#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
7pub enum RepeatedSubtrees {
8    /// Render the complete subtree at every occurrence.
9    #[default]
10    Expand,
11    /// Render a later occurrence as a reference to the first.
12    Reference,
13}
14
15/// Node categories available to visualization style selectors.
16#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
17pub enum ExprNodeKind {
18    /// Real constant node.
19    RealConst,
20    /// Complex constant node.
21    ComplexConst,
22    /// Scalar parameter node.
23    ScalarParam,
24    /// Scalar event-data node.
25    EventScalar,
26    /// Four-momentum component event-data node.
27    EventP4Component,
28    /// Unary-operation node.
29    Unary,
30    /// Binary-operation node.
31    Binary,
32    /// N-ary addition node.
33    NaryAdd,
34    /// N-ary multiplication node.
35    NaryMul,
36    /// Complex-construction node.
37    Complex,
38    /// Vector-construction node.
39    Vector,
40    /// Matrix-construction node.
41    Matrix,
42    /// Vector-component node.
43    Component,
44    /// Matrix-element node.
45    MatrixElement,
46    /// Matrix-matrix multiplication node.
47    MatMul,
48    /// Matrix-vector multiplication node.
49    MatVec,
50    /// Dot-product node.
51    Dot,
52    /// Linear-system solution node.
53    Solve,
54}
55
56impl ExprNodeKind {
57    /// Returns the category corresponding to `node`.
58    pub fn of(node: &ExprNode) -> Self {
59        match node {
60            ExprNode::RealConst(_) => Self::RealConst,
61            ExprNode::ComplexConst(_) => Self::ComplexConst,
62            ExprNode::ScalarParam(_) => Self::ScalarParam,
63            ExprNode::EventScalar(_) => Self::EventScalar,
64            ExprNode::EventP4Component { .. } => Self::EventP4Component,
65            ExprNode::Unary { .. } => Self::Unary,
66            ExprNode::Binary { .. } => Self::Binary,
67            ExprNode::NaryAdd { .. } => Self::NaryAdd,
68            ExprNode::NaryMul { .. } => Self::NaryMul,
69            ExprNode::Complex { .. } => Self::Complex,
70            ExprNode::Vector { .. } => Self::Vector,
71            ExprNode::Matrix { .. } => Self::Matrix,
72            ExprNode::Component { .. } => Self::Component,
73            ExprNode::MatrixElement { .. } => Self::MatrixElement,
74            ExprNode::MatMul { .. } => Self::MatMul,
75            ExprNode::MatVec { .. } => Self::MatVec,
76            ExprNode::Dot { .. } => Self::Dot,
77            ExprNode::Solve { .. } => Self::Solve,
78        }
79    }
80}
81
82/// An RGB color used by tree and Graphviz displays.
83#[derive(Copy, Clone, Debug, PartialEq, Eq)]
84pub struct DisplayColor {
85    red: u8,
86    green: u8,
87    blue: u8,
88}
89
90impl DisplayColor {
91    /// Creates a color from red, green, and blue channels.
92    pub const fn rgb(red: u8, green: u8, blue: u8) -> Self {
93        Self { red, green, blue }
94    }
95
96    fn dot(self) -> String {
97        format!("#{:02x}{:02x}{:02x}", self.red, self.green, self.blue)
98    }
99
100    fn ansi_foreground(self) -> String {
101        format!("\x1b[38;2;{};{};{}m", self.red, self.green, self.blue)
102    }
103}
104
105/// Optional foreground, fill, and border colors for a displayed node.
106#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
107pub struct NodeStyle {
108    /// Text color.
109    pub foreground: Option<DisplayColor>,
110    /// Background or fill color.
111    pub fill: Option<DisplayColor>,
112    /// Outline color.
113    pub border: Option<DisplayColor>,
114}
115
116impl NodeStyle {
117    /// Creates a style with no color overrides.
118    pub const fn new() -> Self {
119        Self {
120            foreground: None,
121            fill: None,
122            border: None,
123        }
124    }
125
126    /// Sets the text color.
127    pub const fn with_foreground(mut self, color: DisplayColor) -> Self {
128        self.foreground = Some(color);
129        self
130    }
131
132    /// Sets the background or fill color.
133    pub const fn with_fill(mut self, color: DisplayColor) -> Self {
134        self.fill = Some(color);
135        self
136    }
137
138    /// Sets the outline color.
139    pub const fn with_border(mut self, color: DisplayColor) -> Self {
140        self.border = Some(color);
141        self
142    }
143
144    fn overlay(&mut self, other: Self) {
145        if other.foreground.is_some() {
146            self.foreground = other.foreground;
147        }
148        if other.fill.is_some() {
149            self.fill = other.fill;
150        }
151        if other.border.is_some() {
152            self.border = other.border;
153        }
154    }
155}
156
157/// Predicate selecting expression nodes for a [`NodeStyleRule`].
158#[derive(Clone, Debug, PartialEq, Eq)]
159pub enum NodeSelector {
160    /// Select every node.
161    Any,
162    /// Select nodes in a category.
163    Kind(ExprNodeKind),
164    /// Select nodes with a matching metadata or source name.
165    Name(String),
166    /// Select nodes carrying a metadata tag.
167    Tag(String),
168}
169
170impl NodeSelector {
171    fn matches(&self, node: &ExprNode, metadata: Option<&ExprMetadata>) -> bool {
172        match self {
173            Self::Any => true,
174            Self::Kind(kind) => *kind == ExprNodeKind::of(node),
175            Self::Name(name) => {
176                metadata.and_then(ExprMetadata::name) == Some(name.as_str())
177                    || match node {
178                        ExprNode::ScalarParam(parameter) => parameter.name() == name,
179                        ExprNode::EventScalar(node_name)
180                        | ExprNode::EventP4Component {
181                            name: node_name, ..
182                        } => node_name.as_ref() == name,
183                        _ => false,
184                    }
185            }
186            Self::Tag(tag) => metadata.is_some_and(|metadata| metadata.has_tag(tag)),
187        }
188    }
189}
190
191/// A selector and the style to overlay on matching nodes.
192#[derive(Clone, Debug, PartialEq, Eq)]
193pub struct NodeStyleRule {
194    /// Predicate used to select nodes.
195    pub selector: NodeSelector,
196    /// Style overlaid on selected nodes.
197    pub style: NodeStyle,
198}
199
200impl NodeStyleRule {
201    /// Creates a style rule from a selector and style.
202    pub fn new(selector: NodeSelector, style: NodeStyle) -> Self {
203        Self { selector, style }
204    }
205}
206
207/// Built-in color palette for expression graphs.
208#[derive(Copy, Clone, Debug, PartialEq, Eq)]
209pub enum ColorPreset {
210    /// Colors selected for light backgrounds.
211    Light,
212    /// Colors selected for dark backgrounds.
213    Dark,
214}
215
216#[derive(Clone, Debug, Default)]
217struct DisplayOptions {
218    repeated_subtrees: RepeatedSubtrees,
219    rules: Vec<NodeStyleRule>,
220}
221
222impl DisplayOptions {
223    fn with_preset(&mut self, preset: ColorPreset) {
224        let (constant, parameter, event, operation, linear_algebra) = match preset {
225            ColorPreset::Light => (
226                DisplayColor::rgb(88, 96, 105),
227                DisplayColor::rgb(0, 92, 197),
228                DisplayColor::rgb(3, 102, 214),
229                DisplayColor::rgb(130, 80, 223),
230                DisplayColor::rgb(207, 34, 46),
231            ),
232            ColorPreset::Dark => (
233                DisplayColor::rgb(139, 148, 158),
234                DisplayColor::rgb(88, 166, 255),
235                DisplayColor::rgb(121, 192, 255),
236                DisplayColor::rgb(210, 168, 255),
237                DisplayColor::rgb(255, 123, 114),
238            ),
239        };
240        let style = |color| NodeStyle::new().with_foreground(color).with_border(color);
241        for kind in [ExprNodeKind::RealConst, ExprNodeKind::ComplexConst] {
242            self.rules.push(NodeStyleRule::new(
243                NodeSelector::Kind(kind),
244                style(constant),
245            ));
246        }
247        self.rules.push(NodeStyleRule::new(
248            NodeSelector::Kind(ExprNodeKind::ScalarParam),
249            style(parameter),
250        ));
251        for kind in [ExprNodeKind::EventScalar, ExprNodeKind::EventP4Component] {
252            self.rules
253                .push(NodeStyleRule::new(NodeSelector::Kind(kind), style(event)));
254        }
255        for kind in [
256            ExprNodeKind::Unary,
257            ExprNodeKind::Binary,
258            ExprNodeKind::NaryAdd,
259            ExprNodeKind::NaryMul,
260            ExprNodeKind::Complex,
261            ExprNodeKind::Vector,
262            ExprNodeKind::Matrix,
263            ExprNodeKind::Component,
264            ExprNodeKind::MatrixElement,
265        ] {
266            self.rules.push(NodeStyleRule::new(
267                NodeSelector::Kind(kind),
268                style(operation),
269            ));
270        }
271        for kind in [
272            ExprNodeKind::MatMul,
273            ExprNodeKind::MatVec,
274            ExprNodeKind::Dot,
275            ExprNodeKind::Solve,
276        ] {
277            self.rules.push(NodeStyleRule::new(
278                NodeSelector::Kind(kind),
279                style(linear_algebra),
280            ));
281        }
282    }
283
284    fn resolve(&self, graph: &ExprGraph, id: ExprId, node: &ExprNode) -> NodeStyle {
285        let mut style = NodeStyle::default();
286        let metadata = graph.metadata(id);
287        for rule in &self.rules {
288            if rule.selector.matches(node, metadata) {
289                style.overlay(rule.style);
290            }
291        }
292        style
293    }
294}
295
296macro_rules! display_builder_methods {
297    () => {
298        /// Sets how nodes reached through multiple paths are rendered.
299        pub fn repeated_subtrees(mut self, repeated_subtrees: RepeatedSubtrees) -> Self {
300            self.options.repeated_subtrees = repeated_subtrees;
301            self
302        }
303
304        /// Chooses between fully expanding and referencing repeated subtrees.
305        pub fn expand_repeated(self, expand: bool) -> Self {
306            self.repeated_subtrees(if expand {
307                RepeatedSubtrees::Expand
308            } else {
309                RepeatedSubtrees::Reference
310            })
311        }
312
313        /// Adds the style rules from a built-in color palette.
314        pub fn with_preset(mut self, preset: ColorPreset) -> Self {
315            self.options.with_preset(preset);
316            self
317        }
318
319        /// Appends a node style rule.
320        ///
321        /// Later matching rules override fields set by earlier rules.
322        pub fn with_style_rule(mut self, rule: NodeStyleRule) -> Self {
323            self.options.rules.push(rule);
324            self
325        }
326    };
327}
328
329/// Configurable indented-tree display for an [`ExprGraph`].
330pub struct ExprGraphTreeDisplay<'a> {
331    graph: &'a ExprGraph,
332    options: DisplayOptions,
333}
334
335impl<'a> ExprGraphTreeDisplay<'a> {
336    pub(crate) fn new(graph: &'a ExprGraph) -> Self {
337        Self {
338            graph,
339            options: DisplayOptions::default(),
340        }
341    }
342
343    display_builder_methods!();
344
345    fn fmt_node(
346        &self,
347        f: &mut fmt::Formatter<'_>,
348        id: ExprId,
349        prefix: &str,
350        edge: Option<(&str, bool)>,
351        visited: &mut HashSet<ExprId>,
352    ) -> fmt::Result {
353        let Some(node) = self.graph.node(id) else {
354            return write_tree_line(f, prefix, edge, &format!("#{} <missing node>", id.index()));
355        };
356        let repeated = !visited.insert(id);
357        let mut line = if repeated && self.options.repeated_subtrees == RepeatedSubtrees::Reference
358        {
359            format!("#{0} <reference to #{0}>", id.index())
360        } else {
361            self.graph.node_label(id, node)
362        };
363        if let Some(color) = self.options.resolve(self.graph, id, node).foreground {
364            line = format!("{}{line}\x1b[0m", color.ansi_foreground());
365        }
366        write_tree_line(f, prefix, edge, &line)?;
367        if repeated && self.options.repeated_subtrees == RepeatedSubtrees::Reference {
368            return Ok(());
369        }
370
371        let children = node_children(node);
372        let child_prefix = match edge {
373            Some((_, true)) => format!("{prefix}   "),
374            Some((_, false)) => format!("{prefix}┃  "),
375            None => prefix.to_owned(),
376        };
377        for (index, (label, child)) in children.iter().enumerate() {
378            self.fmt_node(
379                f,
380                *child,
381                &child_prefix,
382                Some((label, index + 1 == children.len())),
383                visited,
384            )?;
385        }
386        Ok(())
387    }
388}
389
390impl fmt::Display for ExprGraphTreeDisplay<'_> {
391    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
392        writeln!(f, "ExprGraph(root=#{})", self.graph.root().index())?;
393        self.fmt_node(f, self.graph.root(), "", None, &mut HashSet::new())
394    }
395}
396
397/// Configurable Graphviz DOT display for an [`ExprGraph`].
398pub struct ExprGraphDotDisplay<'a> {
399    graph: &'a ExprGraph,
400    options: DisplayOptions,
401}
402
403impl<'a> ExprGraphDotDisplay<'a> {
404    pub(crate) fn new(graph: &'a ExprGraph) -> Self {
405        Self {
406            graph,
407            options: DisplayOptions::default(),
408        }
409    }
410
411    display_builder_methods!();
412
413    #[cfg(feature = "svg")]
414    /// Renders the generated Graphviz graph as an SVG document.
415    ///
416    /// # Errors
417    ///
418    /// Returns [`GraphRenderError::Dot`] when the generated Graphviz DOT
419    /// source cannot be parsed.
420    pub fn render_svg(&self) -> Result<String, GraphRenderError> {
421        use layout::{backends::svg::SVGWriter, gv};
422
423        let dot = self.to_string();
424        let mut parser = gv::DotParser::new(&dot);
425        let graph = parser.process().map_err(GraphRenderError::Dot)?;
426        let mut builder = gv::GraphBuilder::new();
427        builder.visit_graph(&graph);
428        let mut graph = builder.get();
429        let mut svg = SVGWriter::new();
430        graph.do_it(false, false, false, &mut svg);
431        Ok(svg.finalize())
432    }
433
434    fn node_attributes(&self, id: ExprId, node: &ExprNode) -> String {
435        let mut attributes = vec![format!(
436            "label=\"{}\"",
437            escape_dot(&self.graph.node_label(id, node))
438        )];
439        let style = self.options.resolve(self.graph, id, node);
440        if let Some(color) = style.foreground {
441            attributes.push(format!("fontcolor=\"{}\"", color.dot()));
442        }
443        if let Some(color) = style.border {
444            attributes.push(format!("color=\"{}\"", color.dot()));
445        }
446        if let Some(color) = style.fill {
447            attributes.push(format!("fillcolor=\"{}\"", color.dot()));
448            attributes.push("style=filled".to_owned());
449        }
450        attributes.join(", ")
451    }
452
453    fn write_expanded(
454        &self,
455        f: &mut fmt::Formatter<'_>,
456        id: ExprId,
457        occurrence: &mut usize,
458    ) -> fmt::Result {
459        let current = *occurrence;
460        *occurrence += 1;
461        let Some(node) = self.graph.node(id) else {
462            return Ok(());
463        };
464        writeln!(f, "  n{current} [{}];", self.node_attributes(id, node))?;
465        for (label, child) in node_children(node) {
466            let child_occurrence = *occurrence;
467            self.write_expanded(f, child, occurrence)?;
468            writeln!(
469                f,
470                "  n{current} -> n{child_occurrence} [label=\"{}\"];",
471                escape_dot(&label)
472            )?;
473        }
474        Ok(())
475    }
476
477    fn write_shared(
478        &self,
479        f: &mut fmt::Formatter<'_>,
480        id: ExprId,
481        visited: &mut HashSet<ExprId>,
482    ) -> fmt::Result {
483        if !visited.insert(id) {
484            return Ok(());
485        }
486        let Some(node) = self.graph.node(id) else {
487            return Ok(());
488        };
489        writeln!(f, "  n{} [{}];", id.index(), self.node_attributes(id, node))?;
490        for (label, child) in node_children(node) {
491            self.write_shared(f, child, visited)?;
492            writeln!(
493                f,
494                "  n{} -> n{} [label=\"{}\"];",
495                id.index(),
496                child.index(),
497                escape_dot(&label)
498            )?;
499        }
500        Ok(())
501    }
502}
503
504#[cfg(feature = "svg")]
505/// Errors produced while rendering an expression graph.
506#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
507pub enum GraphRenderError {
508    /// The generated Graphviz DOT source could not be parsed.
509    #[error("failed to parse generated DOT: {0}")]
510    Dot(String),
511}
512
513impl fmt::Display for ExprGraphDotDisplay<'_> {
514    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
515        writeln!(f, "digraph ExprGraph {{")?;
516        match self.options.repeated_subtrees {
517            RepeatedSubtrees::Expand => self.write_expanded(f, self.graph.root(), &mut 0)?,
518            RepeatedSubtrees::Reference => {
519                self.write_shared(f, self.graph.root(), &mut HashSet::new())?
520            }
521        }
522        writeln!(f, "}}")
523    }
524}
525
526fn write_tree_line(
527    f: &mut fmt::Formatter<'_>,
528    prefix: &str,
529    edge: Option<(&str, bool)>,
530    text: &str,
531) -> fmt::Result {
532    if let Some((label, is_last)) = edge {
533        let connector = if is_last { "┗" } else { "┣" };
534        writeln!(f, "{prefix}{connector} {label}: {text}")
535    } else {
536        writeln!(f, "{text}")
537    }
538}
539
540fn escape_dot(value: &str) -> String {
541    value
542        .replace('\\', "\\\\")
543        .replace('"', "\\\"")
544        .replace('\n', "\\n")
545        .replace('\r', "\\r")
546}
547
548#[cfg(test)]
549mod tests {
550    use super::*;
551    use crate::event_scalar;
552
553    fn shared_graph() -> ExprGraph {
554        let shared = event_scalar("x").named("shared").tagged("data");
555        ((shared.clone() + 1.0) * (shared + 2.0)).to_graph()
556    }
557
558    #[test]
559    fn tree_and_dot_expand_repeated_subtrees_without_color_by_default() {
560        let graph = shared_graph();
561        let shared_id = graph
562            .nodes()
563            .iter()
564            .position(|node| matches!(node, ExprNode::EventScalar(name) if name.as_ref() == "x"))
565            .unwrap();
566        let needle = format!("#{shared_id} EventScalar(x)");
567        let tree = graph.display_tree().to_string();
568        let dot = graph.display_dot().to_string();
569
570        assert_eq!(tree.matches(&needle).count(), 2);
571        assert_eq!(dot.matches(&needle).count(), 2);
572        assert!(!tree.contains("\x1b["));
573        assert!(!dot.contains("fontcolor="));
574        assert!(!dot.contains("fillcolor="));
575    }
576
577    #[test]
578    fn reference_mode_suppresses_repeated_tree_expansion_and_emits_a_shared_dag() {
579        let graph = shared_graph();
580        let tree = graph
581            .display_tree()
582            .repeated_subtrees(RepeatedSubtrees::Reference)
583            .to_string();
584        let dot = graph.display_dot().expand_repeated(false).to_string();
585
586        assert_eq!(tree.matches("EventScalar(x)").count(), 1);
587        assert_eq!(tree.matches("<reference to #").count(), 1);
588        assert_eq!(dot.matches("EventScalar(x)").count(), 1);
589        assert_eq!(dot.matches(" -> ").count(), 6);
590    }
591
592    #[test]
593    fn later_style_rules_override_matching_preset_fields() {
594        let graph = shared_graph();
595        let override_color = DisplayColor::rgb(1, 2, 3);
596        let rule = NodeStyleRule::new(
597            NodeSelector::Tag("data".to_owned()),
598            NodeStyle::new().with_foreground(override_color),
599        );
600        let tree = graph
601            .display_tree()
602            .with_preset(ColorPreset::Light)
603            .with_style_rule(rule.clone())
604            .to_string();
605        let dot = graph
606            .display_dot()
607            .with_preset(ColorPreset::Light)
608            .with_style_rule(rule)
609            .to_string();
610
611        assert!(tree.contains("\x1b[38;2;1;2;3m"));
612        assert!(dot.contains("fontcolor=\"#010203\""));
613    }
614
615    #[test]
616    fn dot_escapes_metadata_and_event_labels() {
617        let graph = event_scalar("x\\\"y").named("quoted\"name").to_graph();
618        let dot = graph.display_dot().to_string();
619
620        assert!(dot.contains("x\\\\\\\"y"));
621        assert!(dot.contains("quoted\\\"name"));
622    }
623
624    #[cfg(feature = "svg")]
625    #[test]
626    fn dot_display_renders_svg_in_process() {
627        let svg = shared_graph()
628            .display_dot()
629            .with_preset(ColorPreset::Light)
630            .render_svg()
631            .unwrap();
632
633        assert!(svg.contains("<svg"));
634        assert!(svg.contains("</svg>"));
635    }
636}
laddu_expr/visualization.rs

laddu_expr/
visualization.rs
