laddu_expr/
visualization.rs

1use std::{collections::HashSet, fmt};
2
3use num::complex::Complex64;
4
5use crate::{
6    BinaryOp, ExprGraph, ExprId, ExprMetadata, ExprNode, UnaryOp, expression::node_children,
7};
8
9/// Node categories available to visualization style selectors.
10#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
11pub enum ExprNodeKind {
12    /// Real constant node.
13    RealConst,
14    /// Complex constant node.
15    ComplexConst,
16    /// Scalar parameter node.
17    ScalarParam,
18    /// Scalar event-data node.
19    EventScalar,
20    /// Four-momentum component event-data node.
21    EventP4Component,
22    /// Unary-operation node.
23    Unary,
24    /// Binary-operation node.
25    Binary,
26    /// N-ary addition node.
27    NaryAdd,
28    /// N-ary multiplication node.
29    NaryMul,
30    /// Complex-construction node.
31    Complex,
32    /// Vector-construction node.
33    Vector,
34    /// Matrix-construction node.
35    Matrix,
36    /// Vector-component node.
37    Component,
38    /// Matrix-element node.
39    MatrixElement,
40    /// Matrix-matrix multiplication node.
41    MatMul,
42    /// Matrix-vector multiplication node.
43    MatVec,
44    /// Dot-product node.
45    Dot,
46    /// Linear-system solution node.
47    Solve,
48}
49
50impl ExprNodeKind {
51    /// Returns the category corresponding to `node`.
52    pub fn of(node: &ExprNode) -> Self {
53        match node {
54            ExprNode::RealConst(_) => Self::RealConst,
55            ExprNode::ComplexConst(_) => Self::ComplexConst,
56            ExprNode::ScalarParam(_) => Self::ScalarParam,
57            ExprNode::EventScalar(_) => Self::EventScalar,
58            ExprNode::EventP4Component { .. } => Self::EventP4Component,
59            ExprNode::Unary { .. } => Self::Unary,
60            ExprNode::Binary { .. } => Self::Binary,
61            ExprNode::NaryAdd { .. } => Self::NaryAdd,
62            ExprNode::NaryMul { .. } => Self::NaryMul,
63            ExprNode::Complex { .. } => Self::Complex,
64            ExprNode::Vector { .. } => Self::Vector,
65            ExprNode::Matrix { .. } => Self::Matrix,
66            ExprNode::Component { .. } => Self::Component,
67            ExprNode::MatrixElement { .. } => Self::MatrixElement,
68            ExprNode::MatMul { .. } => Self::MatMul,
69            ExprNode::MatVec { .. } => Self::MatVec,
70            ExprNode::Dot { .. } => Self::Dot,
71            ExprNode::Solve { .. } => Self::Solve,
72        }
73    }
74}
75
76/// An RGB color used by tree and Graphviz displays.
77#[derive(Copy, Clone, Debug, PartialEq, Eq)]
78pub struct DisplayColor {
79    red: u8,
80    green: u8,
81    blue: u8,
82}
83
84impl DisplayColor {
85    /// Creates a color from red, green, and blue channels.
86    pub const fn rgb(red: u8, green: u8, blue: u8) -> Self {
87        Self { red, green, blue }
88    }
89
90    fn dot(self) -> String {
91        format!("#{:02x}{:02x}{:02x}", self.red, self.green, self.blue)
92    }
93
94    fn ansi_foreground(self) -> String {
95        format!("\x1b[38;2;{};{};{}m", self.red, self.green, self.blue)
96    }
97
98    fn ansi_background(self) -> String {
99        format!("\x1b[48;2;{};{};{}m", self.red, self.green, self.blue)
100    }
101}
102
103/// Optional foreground, fill, and border colors for a displayed node.
104#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
105pub struct NodeStyle {
106    /// Text color.
107    pub foreground: Option<DisplayColor>,
108    /// Background or fill color.
109    pub fill: Option<DisplayColor>,
110    /// Outline color.
111    pub border: Option<DisplayColor>,
112}
113
114impl NodeStyle {
115    /// Creates a style with no color overrides.
116    pub const fn new() -> Self {
117        Self {
118            foreground: None,
119            fill: None,
120            border: None,
121        }
122    }
123
124    /// Sets the text color.
125    pub const fn with_foreground(mut self, color: DisplayColor) -> Self {
126        self.foreground = Some(color);
127        self
128    }
129
130    /// Sets the background or fill color.
131    pub const fn with_fill(mut self, color: DisplayColor) -> Self {
132        self.fill = Some(color);
133        self
134    }
135
136    /// Sets the outline color.
137    pub const fn with_border(mut self, color: DisplayColor) -> Self {
138        self.border = Some(color);
139        self
140    }
141
142    fn overlay(&mut self, other: Self) {
143        if other.foreground.is_some() {
144            self.foreground = other.foreground;
145        }
146        if other.fill.is_some() {
147            self.fill = other.fill;
148        }
149        if other.border.is_some() {
150            self.border = other.border;
151        }
152    }
153
154    fn ansi(self, text: String) -> String {
155        if self.foreground.is_none() && self.fill.is_none() {
156            return text;
157        }
158        let mut prefix = String::new();
159        if let Some(color) = self.foreground {
160            prefix.push_str(&color.ansi_foreground());
161        }
162        if let Some(color) = self.fill {
163            prefix.push_str(&color.ansi_background());
164        }
165        format!("{prefix}{text}\x1b[0m")
166    }
167
168    fn latex(self, text: String) -> String {
169        self.foreground.map_or(text.clone(), |color| {
170            format!(
171                "{{\\color[RGB]{{{},{},{}}}{text}}}",
172                color.red, color.green, color.blue
173            )
174        })
175    }
176}
177
178/// Predicate selecting expression nodes for a [`NodeStyleRule`].
179#[derive(Clone, Debug, PartialEq, Eq)]
180pub enum NodeSelector {
181    /// Select every node.
182    Any,
183    /// Select nodes in a category.
184    Kind(ExprNodeKind),
185    /// Select nodes with a matching metadata or source name.
186    Name(String),
187    /// Select nodes carrying a metadata tag.
188    Tag(String),
189}
190
191impl NodeSelector {
192    fn matches(&self, node: &ExprNode, metadata: Option<&ExprMetadata>) -> bool {
193        match self {
194            Self::Any => true,
195            Self::Kind(kind) => *kind == ExprNodeKind::of(node),
196            Self::Name(name) => {
197                metadata.and_then(ExprMetadata::name) == Some(name.as_str())
198                    || match node {
199                        ExprNode::ScalarParam(parameter) => parameter.name() == name,
200                        ExprNode::EventScalar(node_name)
201                        | ExprNode::EventP4Component {
202                            name: node_name, ..
203                        } => node_name.as_ref() == name,
204                        _ => false,
205                    }
206            }
207            Self::Tag(tag) => metadata.is_some_and(|metadata| metadata.has_tag(tag)),
208        }
209    }
210}
211
212/// A selector and the style to overlay on matching nodes.
213#[derive(Clone, Debug, PartialEq, Eq)]
214pub struct NodeStyleRule {
215    /// Predicate used to select nodes.
216    pub selector: NodeSelector,
217    /// Style overlaid on selected nodes.
218    pub style: NodeStyle,
219}
220
221impl NodeStyleRule {
222    /// Creates a style rule from a selector and style.
223    pub fn new(selector: NodeSelector, style: NodeStyle) -> Self {
224        Self { selector, style }
225    }
226}
227
228/// Built-in color palette for expression graphs.
229#[derive(Copy, Clone, Debug, PartialEq, Eq)]
230pub enum ColorPreset {
231    /// Colors selected for light backgrounds.
232    Light,
233    /// Colors selected for dark backgrounds.
234    Dark,
235}
236
237#[derive(Clone, Debug)]
238struct DisplayOptions {
239    expand_repeated: bool,
240    rules: Vec<NodeStyleRule>,
241}
242
243impl Default for DisplayOptions {
244    fn default() -> Self {
245        Self {
246            expand_repeated: true,
247            rules: Vec::new(),
248        }
249    }
250}
251
252impl DisplayOptions {
253    fn with_preset(&mut self, preset: ColorPreset) {
254        let (constant, parameter, event, operation, linear_algebra) = match preset {
255            ColorPreset::Light => (
256                DisplayColor::rgb(88, 96, 105),
257                DisplayColor::rgb(0, 92, 197),
258                DisplayColor::rgb(3, 102, 214),
259                DisplayColor::rgb(130, 80, 223),
260                DisplayColor::rgb(207, 34, 46),
261            ),
262            ColorPreset::Dark => (
263                DisplayColor::rgb(139, 148, 158),
264                DisplayColor::rgb(88, 166, 255),
265                DisplayColor::rgb(121, 192, 255),
266                DisplayColor::rgb(210, 168, 255),
267                DisplayColor::rgb(255, 123, 114),
268            ),
269        };
270        let style = |color| NodeStyle::new().with_foreground(color).with_border(color);
271        for kind in [ExprNodeKind::RealConst, ExprNodeKind::ComplexConst] {
272            self.rules.push(NodeStyleRule::new(
273                NodeSelector::Kind(kind),
274                style(constant),
275            ));
276        }
277        self.rules.push(NodeStyleRule::new(
278            NodeSelector::Kind(ExprNodeKind::ScalarParam),
279            style(parameter),
280        ));
281        for kind in [ExprNodeKind::EventScalar, ExprNodeKind::EventP4Component] {
282            self.rules
283                .push(NodeStyleRule::new(NodeSelector::Kind(kind), style(event)));
284        }
285        for kind in [
286            ExprNodeKind::Unary,
287            ExprNodeKind::Binary,
288            ExprNodeKind::NaryAdd,
289            ExprNodeKind::NaryMul,
290            ExprNodeKind::Complex,
291            ExprNodeKind::Vector,
292            ExprNodeKind::Matrix,
293            ExprNodeKind::Component,
294            ExprNodeKind::MatrixElement,
295        ] {
296            self.rules.push(NodeStyleRule::new(
297                NodeSelector::Kind(kind),
298                style(operation),
299            ));
300        }
301        for kind in [
302            ExprNodeKind::MatMul,
303            ExprNodeKind::MatVec,
304            ExprNodeKind::Dot,
305            ExprNodeKind::Solve,
306        ] {
307            self.rules.push(NodeStyleRule::new(
308                NodeSelector::Kind(kind),
309                style(linear_algebra),
310            ));
311        }
312    }
313
314    fn resolve(&self, graph: &ExprGraph, id: ExprId, node: &ExprNode) -> NodeStyle {
315        let mut style = NodeStyle::default();
316        let metadata = graph.metadata(id);
317        for rule in &self.rules {
318            if rule.selector.matches(node, metadata) {
319                style.overlay(rule.style);
320            }
321        }
322        style
323    }
324}
325
326impl ExprGraph {
327    /// Creates a configurable indented-tree display.
328    pub fn display_tree(&self) -> crate::ExprGraphTreeDisplay<'_> {
329        crate::ExprGraphTreeDisplay::new(self)
330    }
331
332    /// Creates a configurable compact-equation display.
333    pub fn display_equation(&self) -> crate::ExprGraphEquationDisplay<'_> {
334        crate::ExprGraphEquationDisplay::new(self)
335    }
336
337    /// Creates a configurable LaTeX equation display.
338    pub fn display_latex(&self) -> crate::ExprGraphLatexDisplay<'_> {
339        crate::ExprGraphLatexDisplay::new(self)
340    }
341
342    /// Creates a configurable Graphviz DOT display.
343    pub fn display_dot(&self) -> crate::ExprGraphDotDisplay<'_> {
344        crate::ExprGraphDotDisplay::new(self)
345    }
346
347    fn format_expression(&self, id: ExprId) -> String {
348        self.format_expression_with(id, &|_, _, text| text)
349    }
350
351    pub(crate) fn format_expression_with(
352        &self,
353        id: ExprId,
354        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
355    ) -> String {
356        self.format_child(id, ExprPrecedence::Lowest, false, decorate)
357    }
358
359    fn format_child(
360        &self,
361        id: ExprId,
362        parent_precedence: ExprPrecedence,
363        parenthesize_equal: bool,
364        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
365    ) -> String {
366        let (text, precedence) = self.format_node_expression(id, decorate);
367        if precedence < parent_precedence || (parenthesize_equal && precedence == parent_precedence)
368        {
369            format!("({text})")
370        } else {
371            text
372        }
373    }
374
375    fn format_node_expression(
376        &self,
377        id: ExprId,
378        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
379    ) -> (String, ExprPrecedence) {
380        let Some(node) = self.node(id) else {
381            return (format!("<missing #{}>", id.index()), ExprPrecedence::Atom);
382        };
383
384        let (text, precedence) = match node {
385            ExprNode::RealConst(value) => (Self::format_real_number(*value), ExprPrecedence::Atom),
386            ExprNode::ComplexConst(value) => self.format_complex_const(*value),
387            ExprNode::ScalarParam(parameter) => (parameter.name().to_owned(), ExprPrecedence::Atom),
388            ExprNode::EventScalar(name) => (name.to_string(), ExprPrecedence::Atom),
389            ExprNode::EventP4Component { name, component } => (
390                format!("{name}.{}", component.label()),
391                ExprPrecedence::Atom,
392            ),
393            ExprNode::Unary { op, input } => self.format_unary_expression(*op, *input, decorate),
394            ExprNode::Binary { op, lhs, rhs } => {
395                self.format_binary_expression(*op, *lhs, *rhs, decorate)
396            }
397            ExprNode::NaryAdd { terms } => self.format_sum_expression(terms, decorate),
398            ExprNode::NaryMul { factors } => self.format_product_expression(factors, decorate),
399            ExprNode::Complex { re, im } => (
400                format!(
401                    "complex({}, {})",
402                    self.format_expression_with(*re, decorate),
403                    self.format_expression_with(*im, decorate)
404                ),
405                ExprPrecedence::Atom,
406            ),
407            ExprNode::Vector { elements } => (
408                format!(
409                    "[{}]",
410                    elements
411                        .iter()
412                        .map(|id| self.format_expression_with(*id, decorate))
413                        .collect::<Vec<_>>()
414                        .join(", ")
415                ),
416                ExprPrecedence::Atom,
417            ),
418            ExprNode::Matrix {
419                rows,
420                cols,
421                elements,
422            } => {
423                let rows = (0..*rows)
424                    .map(|row| {
425                        let start = row * *cols;
426                        let end = start + *cols;
427                        format!(
428                            "[{}]",
429                            elements[start..end]
430                                .iter()
431                                .map(|id| self.format_expression_with(*id, decorate))
432                                .collect::<Vec<_>>()
433                                .join(", ")
434                        )
435                    })
436                    .collect::<Vec<_>>()
437                    .join(", ");
438                (format!("[{rows}]"), ExprPrecedence::Atom)
439            }
440            ExprNode::Component { input, index } => (
441                format!(
442                    "{}[{index}]",
443                    self.format_child(*input, ExprPrecedence::Postfix, false, decorate)
444                ),
445                ExprPrecedence::Postfix,
446            ),
447            ExprNode::MatrixElement { input, row, col } => (
448                format!(
449                    "{}[{row}, {col}]",
450                    self.format_child(*input, ExprPrecedence::Postfix, false, decorate)
451                ),
452                ExprPrecedence::Postfix,
453            ),
454            ExprNode::MatMul { lhs, rhs } => (
455                format!(
456                    "matmul({}, {})",
457                    self.format_expression_with(*lhs, decorate),
458                    self.format_expression_with(*rhs, decorate)
459                ),
460                ExprPrecedence::Atom,
461            ),
462            ExprNode::MatVec { matrix, vector } => (
463                format!(
464                    "matvec({}, {})",
465                    self.format_expression_with(*matrix, decorate),
466                    self.format_expression_with(*vector, decorate)
467                ),
468                ExprPrecedence::Atom,
469            ),
470            ExprNode::Dot { lhs, rhs } => (
471                format!(
472                    "dot({}, {})",
473                    self.format_expression_with(*lhs, decorate),
474                    self.format_expression_with(*rhs, decorate)
475                ),
476                ExprPrecedence::Atom,
477            ),
478            ExprNode::Solve { matrix, rhs } => (
479                format!(
480                    "solve({}, {})",
481                    self.format_expression_with(*matrix, decorate),
482                    self.format_expression_with(*rhs, decorate)
483                ),
484                ExprPrecedence::Atom,
485            ),
486        };
487        (decorate(id, node, text), precedence)
488    }
489
490    fn format_unary_expression(
491        &self,
492        op: UnaryOp,
493        input: ExprId,
494        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
495    ) -> (String, ExprPrecedence) {
496        match op {
497            UnaryOp::Neg => (
498                format!(
499                    "-{}",
500                    self.format_child(input, ExprPrecedence::Unary, true, decorate)
501                ),
502                ExprPrecedence::Unary,
503            ),
504            UnaryOp::Real => (
505                self.format_call_expression("real", input, decorate),
506                ExprPrecedence::Atom,
507            ),
508            UnaryOp::Imag => (
509                self.format_call_expression("imag", input, decorate),
510                ExprPrecedence::Atom,
511            ),
512            UnaryOp::Conj => (
513                self.format_call_expression("conj", input, decorate),
514                ExprPrecedence::Atom,
515            ),
516            UnaryOp::NormSqr => (
517                format!("|{}|^2", self.format_expression_with(input, decorate)),
518                ExprPrecedence::Pow,
519            ),
520            UnaryOp::Sqrt => (
521                self.format_call_expression("sqrt", input, decorate),
522                ExprPrecedence::Atom,
523            ),
524            UnaryOp::Exp => (
525                self.format_call_expression("exp", input, decorate),
526                ExprPrecedence::Atom,
527            ),
528            UnaryOp::Sin => (
529                self.format_call_expression("sin", input, decorate),
530                ExprPrecedence::Atom,
531            ),
532            UnaryOp::Cos => (
533                self.format_call_expression("cos", input, decorate),
534                ExprPrecedence::Atom,
535            ),
536            UnaryOp::Log => (
537                self.format_call_expression("log", input, decorate),
538                ExprPrecedence::Atom,
539            ),
540            UnaryOp::PowI(power) => {
541                let exponent = if power < 0 {
542                    format!("({power})")
543                } else {
544                    power.to_string()
545                };
546                (
547                    format!(
548                        "{}^{exponent}",
549                        self.format_child(input, ExprPrecedence::Pow, true, decorate)
550                    ),
551                    ExprPrecedence::Pow,
552                )
553            }
554        }
555    }
556
557    fn format_call_expression(
558        &self,
559        name: &str,
560        input: ExprId,
561        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
562    ) -> String {
563        format!("{name}({})", self.format_expression_with(input, decorate))
564    }
565
566    fn format_binary_expression(
567        &self,
568        op: BinaryOp,
569        lhs: ExprId,
570        rhs: ExprId,
571        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
572    ) -> (String, ExprPrecedence) {
573        match op {
574            BinaryOp::Add => self.format_sum_expression(&[lhs, rhs], decorate),
575            BinaryOp::Sub => {
576                let lhs = self.format_child(lhs, ExprPrecedence::Add, false, decorate);
577                let rhs = self.format_child(rhs, ExprPrecedence::Add, true, decorate);
578                (format!("{lhs} - {rhs}"), ExprPrecedence::Add)
579            }
580            BinaryOp::Mul => self.format_product_expression(&[lhs, rhs], decorate),
581            BinaryOp::Div => {
582                let lhs = self.format_child(lhs, ExprPrecedence::Mul, false, decorate);
583                let rhs = self.format_child(rhs, ExprPrecedence::Mul, true, decorate);
584                (format!("{lhs} / {rhs}"), ExprPrecedence::Mul)
585            }
586            BinaryOp::Atan2 => (
587                format!(
588                    "atan2({}, {})",
589                    self.format_expression_with(lhs, decorate),
590                    self.format_expression_with(rhs, decorate)
591                ),
592                ExprPrecedence::Atom,
593            ),
594        }
595    }
596
597    fn format_sum_expression(
598        &self,
599        terms: &[ExprId],
600        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
601    ) -> (String, ExprPrecedence) {
602        let mut formatted = String::new();
603        for term in terms {
604            let (negative, term) = self.format_signed_term(*term, decorate);
605            if formatted.is_empty() {
606                if negative {
607                    formatted.push('-');
608                }
609                formatted.push_str(&term);
610            } else if negative {
611                formatted.push_str(" - ");
612                formatted.push_str(&term);
613            } else {
614                formatted.push_str(" + ");
615                formatted.push_str(&term);
616            }
617        }
618        if formatted.is_empty() {
619            formatted.push('0');
620        }
621        (formatted, ExprPrecedence::Add)
622    }
623
624    fn format_signed_term(
625        &self,
626        id: ExprId,
627        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
628    ) -> (bool, String) {
629        match self.node(id) {
630            Some(ExprNode::RealConst(value)) if *value < 0.0 => (
631                true,
632                decorate(
633                    id,
634                    self.node(id).unwrap(),
635                    Self::format_real_number(-*value),
636                ),
637            ),
638            Some(ExprNode::Unary {
639                op: UnaryOp::Neg,
640                input,
641            }) => (
642                true,
643                self.format_child(*input, ExprPrecedence::Add, false, decorate),
644            ),
645            Some(ExprNode::NaryMul { factors }) => {
646                let (negative, product) = self.format_product_parts(factors, decorate);
647                (negative, product)
648            }
649            _ => (
650                false,
651                self.format_child(id, ExprPrecedence::Add, false, decorate),
652            ),
653        }
654    }
655
656    fn format_product_expression(
657        &self,
658        factors: &[ExprId],
659        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
660    ) -> (String, ExprPrecedence) {
661        let (negative, product) = self.format_product_parts(factors, decorate);
662        if negative {
663            (format!("-{product}"), ExprPrecedence::Unary)
664        } else {
665            (product, ExprPrecedence::Mul)
666        }
667    }
668
669    fn format_product_parts(
670        &self,
671        factors: &[ExprId],
672        decorate: &dyn Fn(ExprId, &ExprNode, String) -> String,
673    ) -> (bool, String) {
674        let mut negative = false;
675        let mut pieces = Vec::new();
676
677        for factor in factors {
678            match self.node(*factor) {
679                Some(ExprNode::RealConst(value)) if *value < 0.0 => {
680                    negative = !negative;
681                    if *value != -1.0 || factors.len() == 1 {
682                        pieces.push(decorate(
683                            *factor,
684                            self.node(*factor).unwrap(),
685                            Self::format_real_number(-*value),
686                        ));
687                    }
688                }
689                Some(ExprNode::Unary {
690                    op: UnaryOp::Neg,
691                    input,
692                }) => {
693                    negative = !negative;
694                    pieces.push(self.format_child(*input, ExprPrecedence::Mul, false, decorate));
695                }
696                _ => pieces.push(self.format_child(*factor, ExprPrecedence::Mul, false, decorate)),
697            }
698        }
699
700        if pieces.is_empty() {
701            pieces.push("1".to_owned());
702        }
703
704        (negative, pieces.join(" * "))
705    }
706
707    fn format_complex_const(&self, value: Complex64) -> (String, ExprPrecedence) {
708        match (value.re, value.im) {
709            (re, 0.0) => (Self::format_real_number(re), ExprPrecedence::Atom),
710            (0.0, im) => (Self::format_imaginary_unit(im), ExprPrecedence::Atom),
711            (re, im) if im < 0.0 => (
712                format!(
713                    "{} - {}",
714                    Self::format_real_number(re),
715                    Self::format_imaginary_unit(-im)
716                ),
717                ExprPrecedence::Add,
718            ),
719            (re, im) => (
720                format!(
721                    "{} + {}",
722                    Self::format_real_number(re),
723                    Self::format_imaginary_unit(im)
724                ),
725                ExprPrecedence::Add,
726            ),
727        }
728    }
729
730    fn format_real_number(value: f64) -> String {
731        let Some((value, decimals)) = Self::nearby_simple_decimal(value) else {
732            return value.to_string();
733        };
734
735        if decimals == 0 {
736            return value.to_string();
737        }
738
739        let mut formatted = format!("{value:.decimals$}");
740        while formatted.contains('.') && formatted.ends_with('0') {
741            formatted.pop();
742        }
743        if formatted.ends_with('.') {
744            formatted.pop();
745        }
746        formatted
747    }
748
749    fn nearby_simple_decimal(value: f64) -> Option<(f64, usize)> {
750        if !value.is_finite() {
751            return None;
752        }
753
754        for decimals in 0..=12 {
755            let scale = 10_f64.powi(decimals as i32);
756            let rounded = (value * scale).round() / scale;
757            if Self::nearly_equal(value, rounded) {
758                return Some((rounded, decimals));
759            }
760        }
761
762        None
763    }
764
765    fn nearly_equal(lhs: f64, rhs: f64) -> bool {
766        (lhs - rhs).abs() <= f64::EPSILON * lhs.abs().max(rhs.abs()).max(1.0) * 16.0
767    }
768
769    fn format_imaginary_unit(value: f64) -> String {
770        match value {
771            1.0 => "i".to_owned(),
772            -1.0 => "-i".to_owned(),
773            value => format!("{}i", Self::format_real_number(value)),
774        }
775    }
776
777    pub(crate) fn node_label(&self, id: ExprId, node: &ExprNode) -> String {
778        let mut label = match node {
779            ExprNode::RealConst(value) => {
780                format!(
781                    "#{} RealConst({})",
782                    id.index(),
783                    Self::format_real_number(*value)
784                )
785            }
786            ExprNode::ComplexConst(value) => {
787                let (value, _) = self.format_complex_const(*value);
788                format!("#{} ComplexConst({value})", id.index())
789            }
790            ExprNode::ScalarParam(parameter) => {
791                format!("#{} ScalarParam({})", id.index(), parameter.name())
792            }
793            ExprNode::EventScalar(name) => format!("#{} EventScalar({name})", id.index()),
794            ExprNode::EventP4Component { name, component } => {
795                format!(
796                    "#{} EventP4Component({name}.{})",
797                    id.index(),
798                    component.label()
799                )
800            }
801            ExprNode::Unary { op, .. } => format!("#{} Unary({op:?})", id.index()),
802            ExprNode::Binary { op, .. } => format!("#{} Binary({op:?})", id.index()),
803            ExprNode::NaryAdd { terms } => {
804                format!("#{} NaryAdd(len={})", id.index(), terms.len())
805            }
806            ExprNode::NaryMul { factors } => {
807                format!("#{} NaryMul(len={})", id.index(), factors.len())
808            }
809            ExprNode::Complex { .. } => format!("#{} Complex", id.index()),
810            ExprNode::Vector { elements } => {
811                format!("#{} Vector(len={})", id.index(), elements.len())
812            }
813            ExprNode::Matrix { rows, cols, .. } => {
814                format!("#{} Matrix({rows}x{cols})", id.index())
815            }
816            ExprNode::Component { index, .. } => {
817                format!("#{} Component(index={index})", id.index())
818            }
819            ExprNode::MatrixElement { row, col, .. } => {
820                format!("#{} MatrixElement(row={row}, col={col})", id.index())
821            }
822            ExprNode::MatMul { .. } => format!("#{} MatMul", id.index()),
823            ExprNode::MatVec { .. } => format!("#{} MatVec", id.index()),
824            ExprNode::Dot { .. } => format!("#{} Dot", id.index()),
825            ExprNode::Solve { .. } => format!("#{} Solve", id.index()),
826        };
827
828        if let Some(metadata) = self.metadata(id) {
829            if let Some(name) = metadata.name() {
830                label.push_str(&format!(" name=\"{name}\""));
831            }
832            if !metadata.tags().is_empty() {
833                label.push_str(" tags=[");
834                for (index, tag) in metadata.tags().iter().enumerate() {
835                    if index != 0 {
836                        label.push_str(", ");
837                    }
838                    label.push_str(tag);
839                }
840                label.push(']');
841            }
842        }
843
844        label
845    }
846}
847
848impl fmt::Display for ExprGraph {
849    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
850        f.write_str(&self.format_expression(self.root()))
851    }
852}
853
854#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
855enum ExprPrecedence {
856    Lowest,
857    Add,
858    Mul,
859    Unary,
860    Pow,
861    Postfix,
862    Atom,
863}
864
865/// Configurable compact-equation display for an [`ExprGraph`].
866pub struct ExprGraphEquationDisplay<'a> {
867    graph: &'a ExprGraph,
868    options: DisplayOptions,
869}
870
871impl<'a> ExprGraphEquationDisplay<'a> {
872    pub(crate) fn new(graph: &'a ExprGraph) -> Self {
873        Self {
874            graph,
875            options: DisplayOptions::default(),
876        }
877    }
878
879    /// Adds the style rules from a built-in color palette.
880    pub fn with_preset(mut self, preset: ColorPreset) -> Self {
881        self.options.with_preset(preset);
882        self
883    }
884
885    /// Appends a node style rule.
886    ///
887    /// Later matching rules override fields set by earlier rules.
888    pub fn with_style_rule(mut self, rule: NodeStyleRule) -> Self {
889        self.options.rules.push(rule);
890        self
891    }
892}
893
894impl fmt::Display for ExprGraphEquationDisplay<'_> {
895    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
896        f.write_str(
897            &self
898                .graph
899                .format_expression_with(self.graph.root(), &|id, node, text| {
900                    self.options.resolve(self.graph, id, node).ansi(text)
901                }),
902        )
903    }
904}
905
906/// Configurable LaTeX-equation display for an [`ExprGraph`].
907///
908/// The output is a math-mode fragment. It does not include `$` delimiters or
909/// a document preamble. Vector and matrix nodes use `bmatrix` from `amsmath`.
910/// Color rules emit `\\color[RGB]` declarations and require `xcolor`.
911pub struct ExprGraphLatexDisplay<'a> {
912    graph: &'a ExprGraph,
913    options: DisplayOptions,
914}
915
916impl<'a> ExprGraphLatexDisplay<'a> {
917    pub(crate) fn new(graph: &'a ExprGraph) -> Self {
918        Self {
919            graph,
920            options: DisplayOptions::default(),
921        }
922    }
923
924    /// Adds the style rules from a built-in color palette.
925    pub fn with_preset(mut self, preset: ColorPreset) -> Self {
926        self.options.with_preset(preset);
927        self
928    }
929
930    /// Appends a node style rule.
931    ///
932    /// Later matching rules override fields set by earlier rules. Only the
933    /// foreground color is meaningful for LaTeX output.
934    pub fn with_style_rule(mut self, rule: NodeStyleRule) -> Self {
935        self.options.rules.push(rule);
936        self
937    }
938}
939
940impl fmt::Display for ExprGraphLatexDisplay<'_> {
941    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
942        let decorate = |id: ExprId, node: &ExprNode, text: String| {
943            self.options.resolve(self.graph, id, node).latex(text)
944        };
945        f.write_str(&LatexFormatter::new(self.graph, &decorate).format())
946    }
947}
948
949struct LatexFormatter<'a> {
950    graph: &'a ExprGraph,
951    decorate: &'a dyn Fn(ExprId, &ExprNode, String) -> String,
952}
953
954impl<'a> LatexFormatter<'a> {
955    fn new(
956        graph: &'a ExprGraph,
957        decorate: &'a dyn Fn(ExprId, &ExprNode, String) -> String,
958    ) -> Self {
959        Self { graph, decorate }
960    }
961
962    fn format(&self) -> String {
963        self.format_expression(self.graph.root())
964    }
965
966    fn format_expression(&self, id: ExprId) -> String {
967        self.format_child(id, ExprPrecedence::Lowest, false)
968    }
969
970    fn format_child(
971        &self,
972        id: ExprId,
973        parent_precedence: ExprPrecedence,
974        parenthesize_equal: bool,
975    ) -> String {
976        let (text, precedence) = self.format_node(id);
977        if precedence < parent_precedence || (parenthesize_equal && precedence == parent_precedence)
978        {
979            format!("\\left({text}\\right)")
980        } else {
981            text
982        }
983    }
984
985    fn format_node(&self, id: ExprId) -> (String, ExprPrecedence) {
986        let Some(node) = self.graph.node(id) else {
987            return (
988                format!("\\text{{missing node \\#{}}}", id.index()),
989                ExprPrecedence::Atom,
990            );
991        };
992
993        let (text, precedence) = match node {
994            ExprNode::RealConst(value) => (format_latex_number(*value), ExprPrecedence::Atom),
995            ExprNode::ComplexConst(value) => format_latex_complex(*value),
996            ExprNode::ScalarParam(parameter) => (
997                parameter
998                    .latex_label()
999                    .map(str::to_owned)
1000                    .unwrap_or_else(|| escape_latex(parameter.name())),
1001                ExprPrecedence::Atom,
1002            ),
1003            ExprNode::EventScalar(name) => (escape_latex(name), ExprPrecedence::Atom),
1004            ExprNode::EventP4Component { name, component } => (
1005                format!(
1006                    "{}_{{\\mathrm{{{}}}}}",
1007                    escape_latex(name),
1008                    escape_latex(component.label())
1009                ),
1010                ExprPrecedence::Atom,
1011            ),
1012            ExprNode::Unary { op, input } => self.format_unary(*op, *input),
1013            ExprNode::Binary { op, lhs, rhs } => self.format_binary(*op, *lhs, *rhs),
1014            ExprNode::NaryAdd { terms } => self.format_sum(terms),
1015            ExprNode::NaryMul { factors } => self.format_product(factors),
1016            ExprNode::Complex { re, im } => (
1017                self.format_operator("complex", &[*re, *im]),
1018                ExprPrecedence::Atom,
1019            ),
1020            ExprNode::Vector { elements } => (
1021                format!(
1022                    "\\begin{{bmatrix}}{}\\end{{bmatrix}}",
1023                    elements
1024                        .iter()
1025                        .map(|id| self.format_expression(*id))
1026                        .collect::<Vec<_>>()
1027                        .join(" \\\\ ")
1028                ),
1029                ExprPrecedence::Atom,
1030            ),
1031            ExprNode::Matrix {
1032                rows,
1033                cols,
1034                elements,
1035            } => {
1036                let rows = (0..*rows)
1037                    .map(|row| {
1038                        let start = row * *cols;
1039                        let end = start + *cols;
1040                        elements[start..end]
1041                            .iter()
1042                            .map(|id| self.format_expression(*id))
1043                            .collect::<Vec<_>>()
1044                            .join(" & ")
1045                    })
1046                    .collect::<Vec<_>>()
1047                    .join(" \\\\ ");
1048                (
1049                    format!("\\begin{{bmatrix}}{rows}\\end{{bmatrix}}"),
1050                    ExprPrecedence::Atom,
1051                )
1052            }
1053            ExprNode::Component { input, index } => (
1054                format!(
1055                    "{}_{{{index}}}",
1056                    self.format_child(*input, ExprPrecedence::Postfix, false)
1057                ),
1058                ExprPrecedence::Postfix,
1059            ),
1060            ExprNode::MatrixElement { input, row, col } => (
1061                format!(
1062                    "{}_{{{row},{col}}}",
1063                    self.format_child(*input, ExprPrecedence::Postfix, false)
1064                ),
1065                ExprPrecedence::Postfix,
1066            ),
1067            ExprNode::MatMul { lhs, rhs } => (
1068                self.format_operator("matmul", &[*lhs, *rhs]),
1069                ExprPrecedence::Atom,
1070            ),
1071            ExprNode::MatVec { matrix, vector } => (
1072                self.format_operator("matvec", &[*matrix, *vector]),
1073                ExprPrecedence::Atom,
1074            ),
1075            ExprNode::Dot { lhs, rhs } => (
1076                self.format_operator("dot", &[*lhs, *rhs]),
1077                ExprPrecedence::Atom,
1078            ),
1079            ExprNode::Solve { matrix, rhs } => (
1080                self.format_operator("solve", &[*matrix, *rhs]),
1081                ExprPrecedence::Atom,
1082            ),
1083        };
1084        ((self.decorate)(id, node, text), precedence)
1085    }
1086
1087    fn format_unary(&self, op: UnaryOp, input: ExprId) -> (String, ExprPrecedence) {
1088        match op {
1089            UnaryOp::Neg => (
1090                format!("-{}", self.format_child(input, ExprPrecedence::Unary, true)),
1091                ExprPrecedence::Unary,
1092            ),
1093            UnaryOp::Real => (self.format_operator("Re", &[input]), ExprPrecedence::Atom),
1094            UnaryOp::Imag => (self.format_operator("Im", &[input]), ExprPrecedence::Atom),
1095            UnaryOp::Conj => (
1096                format!("\\overline{{{}}}", self.format_expression(input)),
1097                ExprPrecedence::Atom,
1098            ),
1099            UnaryOp::NormSqr => (
1100                format!("\\left|{}\\right|^{{2}}", self.format_expression(input)),
1101                ExprPrecedence::Pow,
1102            ),
1103            UnaryOp::Sqrt => (
1104                format!("\\sqrt{{{}}}", self.format_expression(input)),
1105                ExprPrecedence::Atom,
1106            ),
1107            UnaryOp::Exp => (self.format_function("exp", input), ExprPrecedence::Atom),
1108            UnaryOp::Sin => (self.format_function("sin", input), ExprPrecedence::Atom),
1109            UnaryOp::Cos => (self.format_function("cos", input), ExprPrecedence::Atom),
1110            UnaryOp::Log => (self.format_function("log", input), ExprPrecedence::Atom),
1111            UnaryOp::PowI(power) => (
1112                format!(
1113                    "{}^{{{power}}}",
1114                    self.format_child(input, ExprPrecedence::Pow, true)
1115                ),
1116                ExprPrecedence::Pow,
1117            ),
1118        }
1119    }
1120
1121    fn format_function(&self, name: &str, input: ExprId) -> String {
1122        format!("\\{name}\\left({}\\right)", self.format_expression(input))
1123    }
1124
1125    fn format_operator(&self, name: &str, inputs: &[ExprId]) -> String {
1126        format!(
1127            "\\operatorname{{{name}}}\\left({}\\right)",
1128            inputs
1129                .iter()
1130                .map(|id| self.format_expression(*id))
1131                .collect::<Vec<_>>()
1132                .join(", ")
1133        )
1134    }
1135
1136    fn format_binary(&self, op: BinaryOp, lhs: ExprId, rhs: ExprId) -> (String, ExprPrecedence) {
1137        match op {
1138            BinaryOp::Add => self.format_sum(&[lhs, rhs]),
1139            BinaryOp::Sub => {
1140                let lhs = self.format_child(lhs, ExprPrecedence::Add, false);
1141                let rhs = self.format_child(rhs, ExprPrecedence::Add, true);
1142                (format!("{lhs} - {rhs}"), ExprPrecedence::Add)
1143            }
1144            BinaryOp::Mul => self.format_product(&[lhs, rhs]),
1145            BinaryOp::Div => (
1146                format!(
1147                    "\\frac{{{}}}{{{}}}",
1148                    self.format_expression(lhs),
1149                    self.format_expression(rhs)
1150                ),
1151                ExprPrecedence::Atom,
1152            ),
1153            BinaryOp::Atan2 => (
1154                self.format_operator("atan2", &[lhs, rhs]),
1155                ExprPrecedence::Atom,
1156            ),
1157        }
1158    }
1159
1160    fn format_sum(&self, terms: &[ExprId]) -> (String, ExprPrecedence) {
1161        let mut formatted = String::new();
1162        for term in terms {
1163            let (negative, term) = self.format_signed_term(*term);
1164            if formatted.is_empty() {
1165                if negative {
1166                    formatted.push('-');
1167                }
1168                formatted.push_str(&term);
1169            } else if negative {
1170                formatted.push_str(" - ");
1171                formatted.push_str(&term);
1172            } else {
1173                formatted.push_str(" + ");
1174                formatted.push_str(&term);
1175            }
1176        }
1177        if formatted.is_empty() {
1178            formatted.push('0');
1179        }
1180        (formatted, ExprPrecedence::Add)
1181    }
1182
1183    fn format_signed_term(&self, id: ExprId) -> (bool, String) {
1184        match self.graph.node(id) {
1185            Some(node @ ExprNode::RealConst(value)) if *value < 0.0 => (
1186                true,
1187                (self.decorate)(id, node, format_latex_number(-*value)),
1188            ),
1189            Some(ExprNode::Unary {
1190                op: UnaryOp::Neg,
1191                input,
1192            }) => (true, self.format_child(*input, ExprPrecedence::Add, false)),
1193            Some(ExprNode::NaryMul { factors }) => self.format_product_parts(factors),
1194            _ => (false, self.format_child(id, ExprPrecedence::Add, false)),
1195        }
1196    }
1197
1198    fn format_product(&self, factors: &[ExprId]) -> (String, ExprPrecedence) {
1199        let (negative, product) = self.format_product_parts(factors);
1200        if negative {
1201            (format!("-{product}"), ExprPrecedence::Unary)
1202        } else {
1203            (product, ExprPrecedence::Mul)
1204        }
1205    }
1206
1207    fn format_product_parts(&self, factors: &[ExprId]) -> (bool, String) {
1208        let mut negative = false;
1209        let mut pieces = Vec::new();
1210        for factor in factors {
1211            match self.graph.node(*factor) {
1212                Some(node @ ExprNode::RealConst(value)) if *value < 0.0 => {
1213                    negative = !negative;
1214                    if *value != -1.0 || factors.len() == 1 {
1215                        pieces.push((self.decorate)(*factor, node, format_latex_number(-*value)));
1216                    }
1217                }
1218                Some(ExprNode::Unary {
1219                    op: UnaryOp::Neg,
1220                    input,
1221                }) => {
1222                    negative = !negative;
1223                    pieces.push(self.format_child(*input, ExprPrecedence::Mul, false));
1224                }
1225                _ => pieces.push(self.format_child(*factor, ExprPrecedence::Mul, false)),
1226            }
1227        }
1228        if pieces.is_empty() {
1229            pieces.push("1".to_owned());
1230        }
1231        (negative, pieces.join(" \\cdot "))
1232    }
1233}
1234
1235fn format_latex_complex(value: Complex64) -> (String, ExprPrecedence) {
1236    match (value.re, value.im) {
1237        (re, 0.0) => (format_latex_number(re), ExprPrecedence::Atom),
1238        (0.0, im) => (format_latex_imaginary(im), ExprPrecedence::Atom),
1239        (re, im) if im < 0.0 => (
1240            format!(
1241                "{} - {}",
1242                format_latex_number(re),
1243                format_latex_imaginary(-im)
1244            ),
1245            ExprPrecedence::Add,
1246        ),
1247        (re, im) => (
1248            format!(
1249                "{} + {}",
1250                format_latex_number(re),
1251                format_latex_imaginary(im)
1252            ),
1253            ExprPrecedence::Add,
1254        ),
1255    }
1256}
1257
1258fn format_latex_number(value: f64) -> String {
1259    if value == f64::INFINITY {
1260        "\\infty".to_owned()
1261    } else if value == f64::NEG_INFINITY {
1262        "-\\infty".to_owned()
1263    } else if value.is_nan() {
1264        "\\mathrm{NaN}".to_owned()
1265    } else {
1266        ExprGraph::format_real_number(value)
1267    }
1268}
1269
1270fn format_latex_imaginary(value: f64) -> String {
1271    match value {
1272        1.0 => "\\mathrm{i}".to_owned(),
1273        -1.0 => "-\\mathrm{i}".to_owned(),
1274        value => format!("{}\\mathrm{{i}}", format_latex_number(value)),
1275    }
1276}
1277
1278fn escape_latex(value: &str) -> String {
1279    let mut escaped = String::with_capacity(value.len());
1280    for character in value.chars() {
1281        match character {
1282            '\\' => escaped.push_str("\\backslash "),
1283            '{' => escaped.push_str("\\{"),
1284            '}' => escaped.push_str("\\}"),
1285            '_' => escaped.push_str("\\_"),
1286            '^' => escaped.push_str("\\^{}"),
1287            '#' => escaped.push_str("\\#"),
1288            '$' => escaped.push_str("\\$"),
1289            '%' => escaped.push_str("\\%"),
1290            '&' => escaped.push_str("\\&"),
1291            '~' => escaped.push_str("\\~{}"),
1292            _ => escaped.push(character),
1293        }
1294    }
1295    escaped
1296}
1297
1298/// Configurable indented-tree display for an [`ExprGraph`].
1299pub struct ExprGraphTreeDisplay<'a> {
1300    graph: &'a ExprGraph,
1301    options: DisplayOptions,
1302}
1303
1304impl<'a> ExprGraphTreeDisplay<'a> {
1305    pub(crate) fn new(graph: &'a ExprGraph) -> Self {
1306        Self {
1307            graph,
1308            options: DisplayOptions::default(),
1309        }
1310    }
1311
1312    /// Sets whether nodes reached through multiple paths are fully expanded.
1313    pub fn expand_repeated(mut self, expand: bool) -> Self {
1314        self.options.expand_repeated = expand;
1315        self
1316    }
1317
1318    /// Adds the style rules from a built-in color palette.
1319    pub fn with_preset(mut self, preset: ColorPreset) -> Self {
1320        self.options.with_preset(preset);
1321        self
1322    }
1323
1324    /// Appends a node style rule.
1325    ///
1326    /// Later matching rules override fields set by earlier rules.
1327    pub fn with_style_rule(mut self, rule: NodeStyleRule) -> Self {
1328        self.options.rules.push(rule);
1329        self
1330    }
1331}
1332
1333impl fmt::Display for ExprGraphTreeDisplay<'_> {
1334    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1335        writeln!(f, "ExprGraph(root=#{})", self.graph.root().index())?;
1336        let mut visited = HashSet::new();
1337        let mut stack = vec![(self.graph.root(), String::new(), None::<(String, bool)>)];
1338        while let Some((id, prefix, edge)) = stack.pop() {
1339            let Some(node) = self.graph.node(id) else {
1340                write_tree_line(
1341                    f,
1342                    &prefix,
1343                    edge.as_ref().map(|(label, last)| (label.as_str(), *last)),
1344                    &format!("#{} <missing node>", id.index()),
1345                )?;
1346                continue;
1347            };
1348            let repeated = !visited.insert(id);
1349            let mut line = if repeated && !self.options.expand_repeated {
1350                format!("#{0} <reference to #{0}>", id.index())
1351            } else {
1352                self.graph.node_label(id, node)
1353            };
1354            line = self.options.resolve(self.graph, id, node).ansi(line);
1355            write_tree_line(
1356                f,
1357                &prefix,
1358                edge.as_ref().map(|(label, last)| (label.as_str(), *last)),
1359                &line,
1360            )?;
1361            if repeated && !self.options.expand_repeated {
1362                continue;
1363            }
1364
1365            let children = node_children(node);
1366            let child_prefix = match edge {
1367                Some((_, true)) => format!("{prefix}   "),
1368                Some((_, false)) => format!("{prefix}┃  "),
1369                None => prefix,
1370            };
1371            let child_count = children.len();
1372            for (index, (label, child)) in children.into_iter().enumerate().rev() {
1373                stack.push((
1374                    child,
1375                    child_prefix.clone(),
1376                    Some((label, index + 1 == child_count)),
1377                ));
1378            }
1379        }
1380        Ok(())
1381    }
1382}
1383
1384/// Configurable Graphviz DOT display for an [`ExprGraph`].
1385pub struct ExprGraphDotDisplay<'a> {
1386    graph: &'a ExprGraph,
1387    options: DisplayOptions,
1388}
1389
1390impl<'a> ExprGraphDotDisplay<'a> {
1391    pub(crate) fn new(graph: &'a ExprGraph) -> Self {
1392        Self {
1393            graph,
1394            options: DisplayOptions::default(),
1395        }
1396    }
1397
1398    /// Sets whether nodes reached through multiple paths are fully expanded.
1399    pub fn expand_repeated(mut self, expand: bool) -> Self {
1400        self.options.expand_repeated = expand;
1401        self
1402    }
1403
1404    /// Adds the style rules from a built-in color palette.
1405    pub fn with_preset(mut self, preset: ColorPreset) -> Self {
1406        self.options.with_preset(preset);
1407        self
1408    }
1409
1410    /// Appends a node style rule.
1411    ///
1412    /// Later matching rules override fields set by earlier rules.
1413    pub fn with_style_rule(mut self, rule: NodeStyleRule) -> Self {
1414        self.options.rules.push(rule);
1415        self
1416    }
1417
1418    #[cfg(feature = "svg")]
1419    /// Renders the generated Graphviz graph as an SVG document.
1420    ///
1421    /// # Errors
1422    ///
1423    /// Returns [`GraphRenderError::Dot`] when the generated Graphviz DOT
1424    /// source cannot be parsed.
1425    pub fn render_svg(&self) -> Result<String, GraphRenderError> {
1426        use layout::{backends::svg::SVGWriter, gv};
1427
1428        let dot = self.to_string();
1429        let mut parser = gv::DotParser::new(&dot);
1430        let graph = parser.process().map_err(GraphRenderError::Dot)?;
1431        let mut builder = gv::GraphBuilder::new();
1432        builder.visit_graph(&graph);
1433        let mut graph = builder.get();
1434        let mut svg = SVGWriter::new();
1435        graph.do_it(false, false, false, &mut svg);
1436        Ok(svg.finalize())
1437    }
1438
1439    fn node_attributes(&self, id: ExprId, node: &ExprNode) -> String {
1440        let mut attributes = vec![format!(
1441            "label=\"{}\"",
1442            escape_dot(&self.graph.node_label(id, node))
1443        )];
1444        let style = self.options.resolve(self.graph, id, node);
1445        if let Some(color) = style.foreground {
1446            attributes.push(format!("fontcolor=\"{}\"", color.dot()));
1447        }
1448        if let Some(color) = style.border {
1449            attributes.push(format!("color=\"{}\"", color.dot()));
1450        }
1451        if let Some(color) = style.fill {
1452            attributes.push(format!("fillcolor=\"{}\"", color.dot()));
1453            attributes.push("style=filled".to_owned());
1454        }
1455        attributes.join(", ")
1456    }
1457
1458    fn write_expanded(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1459        enum Frame {
1460            Visit(ExprId),
1461            Children {
1462                parent: usize,
1463                children: Vec<(String, ExprId)>,
1464                index: usize,
1465            },
1466            Edge {
1467                parent: usize,
1468                child: usize,
1469                label: String,
1470            },
1471        }
1472
1473        let mut occurrence = 0;
1474        let mut stack = vec![Frame::Visit(self.graph.root())];
1475        while let Some(frame) = stack.pop() {
1476            match frame {
1477                Frame::Visit(id) => {
1478                    let current = occurrence;
1479                    occurrence += 1;
1480                    let Some(node) = self.graph.node(id) else {
1481                        continue;
1482                    };
1483                    writeln!(f, "  n{current} [{}];", self.node_attributes(id, node))?;
1484                    stack.push(Frame::Children {
1485                        parent: current,
1486                        children: node_children(node),
1487                        index: 0,
1488                    });
1489                }
1490                Frame::Children {
1491                    parent,
1492                    children,
1493                    index,
1494                } => {
1495                    if let Some((label, child)) = children.get(index).cloned() {
1496                        let child_occurrence = occurrence;
1497                        stack.push(Frame::Children {
1498                            parent,
1499                            children,
1500                            index: index + 1,
1501                        });
1502                        stack.push(Frame::Edge {
1503                            parent,
1504                            child: child_occurrence,
1505                            label,
1506                        });
1507                        stack.push(Frame::Visit(child));
1508                    }
1509                }
1510                Frame::Edge {
1511                    parent,
1512                    child,
1513                    label,
1514                } => writeln!(
1515                    f,
1516                    "  n{parent} -> n{child} [label=\"{}\"];",
1517                    escape_dot(&label)
1518                )?,
1519            }
1520        }
1521        Ok(())
1522    }
1523
1524    fn write_shared(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1525        enum Frame {
1526            Visit(ExprId),
1527            Edge(ExprId, ExprId, String),
1528        }
1529
1530        let mut visited = HashSet::new();
1531        let mut stack = vec![Frame::Visit(self.graph.root())];
1532        while let Some(frame) = stack.pop() {
1533            match frame {
1534                Frame::Visit(id) => {
1535                    if !visited.insert(id) {
1536                        continue;
1537                    }
1538                    let Some(node) = self.graph.node(id) else {
1539                        continue;
1540                    };
1541                    writeln!(f, "  n{} [{}];", id.index(), self.node_attributes(id, node))?;
1542                    for (label, child) in node_children(node).into_iter().rev() {
1543                        stack.push(Frame::Edge(id, child, label));
1544                        stack.push(Frame::Visit(child));
1545                    }
1546                }
1547                Frame::Edge(parent, child, label) => writeln!(
1548                    f,
1549                    "  n{} -> n{} [label=\"{}\"];",
1550                    parent.index(),
1551                    child.index(),
1552                    escape_dot(&label)
1553                )?,
1554            }
1555        }
1556        Ok(())
1557    }
1558}
1559
1560#[cfg(feature = "svg")]
1561/// Errors produced while rendering an expression graph.
1562#[derive(Clone, Debug, thiserror::Error, PartialEq, Eq)]
1563pub enum GraphRenderError {
1564    /// The generated Graphviz DOT source could not be parsed.
1565    #[error("failed to parse generated DOT: {0}")]
1566    Dot(String),
1567}
1568
1569impl fmt::Display for ExprGraphDotDisplay<'_> {
1570    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1571        writeln!(f, "digraph ExprGraph {{")?;
1572        if self.options.expand_repeated {
1573            self.write_expanded(f)?;
1574        } else {
1575            self.write_shared(f)?;
1576        }
1577        writeln!(f, "}}")
1578    }
1579}
1580
1581fn write_tree_line(
1582    f: &mut fmt::Formatter<'_>,
1583    prefix: &str,
1584    edge: Option<(&str, bool)>,
1585    text: &str,
1586) -> fmt::Result {
1587    if let Some((label, is_last)) = edge {
1588        let connector = if is_last { "┗" } else { "┣" };
1589        writeln!(f, "{prefix}{connector} {label}: {text}")
1590    } else {
1591        writeln!(f, "{text}")
1592    }
1593}
1594
1595fn escape_dot(value: &str) -> String {
1596    value
1597        .replace('\\', "\\\\")
1598        .replace('"', "\\\"")
1599        .replace('\n', "\\n")
1600        .replace('\r', "\\r")
1601}
1602
1603#[cfg(test)]
1604mod tests {
1605    use std::sync::Arc;
1606
1607    use super::*;
1608    use crate::{ExprSourceKind, event_scalar};
1609
1610    fn shared_graph() -> ExprGraph {
1611        let shared = event_scalar("x").named("shared").tagged("data");
1612        ((shared.clone() + 1.0) * (shared + 2.0)).to_graph()
1613    }
1614
1615    #[test]
1616    fn equation_latex_tree_and_dot_have_no_color_by_default() {
1617        let graph = shared_graph();
1618        let shared_id = graph
1619            .nodes()
1620            .iter()
1621            .position(|node| matches!(node, ExprNode::EventScalar(name) if name.as_ref() == "x"))
1622            .unwrap();
1623        let needle = format!("#{shared_id} EventScalar(x)");
1624        let tree = graph.display_tree().to_string();
1625        let dot = graph.display_dot().to_string();
1626        let equation = graph.display_equation().to_string();
1627        let latex = graph.display_latex().to_string();
1628
1629        assert_eq!(tree.matches(&needle).count(), 2);
1630        assert_eq!(dot.matches(&needle).count(), 2);
1631        assert!(!tree.contains("\x1b["));
1632        assert!(!dot.contains("fontcolor="));
1633        assert!(!dot.contains("fillcolor="));
1634        assert!(!equation.contains("\x1b["));
1635        assert_eq!(equation, graph.to_string());
1636        assert!(!latex.contains("\\color"));
1637        assert!(latex.contains("\\cdot"));
1638    }
1639
1640    #[test]
1641    fn reference_mode_suppresses_repeated_tree_expansion_and_emits_a_shared_dag() {
1642        let graph = shared_graph();
1643        let tree = graph.display_tree().expand_repeated(false).to_string();
1644        let dot = graph.display_dot().expand_repeated(false).to_string();
1645
1646        assert_eq!(tree.matches("EventScalar(x)").count(), 1);
1647        assert_eq!(tree.matches("<reference to #").count(), 1);
1648        assert_eq!(dot.matches("EventScalar(x)").count(), 1);
1649        assert_eq!(dot.matches(" -> ").count(), 6);
1650    }
1651
1652    #[test]
1653    fn later_style_rules_override_matching_preset_fields() {
1654        let graph = shared_graph();
1655        let override_color = DisplayColor::rgb(1, 2, 3);
1656        let rule = NodeStyleRule::new(
1657            NodeSelector::Tag("data".to_owned()),
1658            NodeStyle::new().with_foreground(override_color),
1659        );
1660        let tree = graph
1661            .display_tree()
1662            .with_preset(ColorPreset::Light)
1663            .with_style_rule(rule.clone())
1664            .to_string();
1665        let dot = graph
1666            .display_dot()
1667            .with_preset(ColorPreset::Light)
1668            .with_style_rule(rule)
1669            .to_string();
1670        let equation = graph
1671            .display_equation()
1672            .with_preset(ColorPreset::Light)
1673            .with_style_rule(NodeStyleRule::new(
1674                NodeSelector::Tag("data".to_owned()),
1675                NodeStyle::new().with_foreground(override_color),
1676            ))
1677            .to_string();
1678        let latex = graph
1679            .display_latex()
1680            .with_preset(ColorPreset::Light)
1681            .with_style_rule(NodeStyleRule::new(
1682                NodeSelector::Tag("data".to_owned()),
1683                NodeStyle::new().with_foreground(override_color),
1684            ))
1685            .to_string();
1686
1687        assert!(tree.contains("\x1b[38;2;1;2;3m"));
1688        assert!(dot.contains("fontcolor=\"#010203\""));
1689        assert!(equation.contains("\x1b[38;2;1;2;3m"));
1690        assert!(latex.contains("\\color[RGB]{1,2,3}"));
1691    }
1692
1693    #[test]
1694    fn latex_uses_math_constructs_and_escapes_names() {
1695        let numerator = crate::event_scalar("x_value");
1696        let denominator = crate::event_scalar("y").sqrt();
1697        let latex = (numerator / denominator)
1698            .powi(-2)
1699            .to_graph()
1700            .display_latex()
1701            .to_string();
1702
1703        assert!(latex.contains("\\frac{"));
1704        assert!(latex.contains("x\\_value"));
1705        assert!(latex.contains("\\sqrt{y}"));
1706        assert!(latex.contains("^{-2}"));
1707    }
1708
1709    #[test]
1710    fn latex_uses_parameter_labels_with_name_fallback() {
1711        let labeled = crate::parameter!("alpha_internal", latex: r"\alpha");
1712        let unlabeled = crate::parameter!("beta_internal");
1713        let latex = (labeled + unlabeled).to_graph().display_latex().to_string();
1714
1715        assert!(latex.contains(r"\alpha"));
1716        assert!(!latex.contains("alpha_internal"));
1717        assert!(latex.contains(r"beta\_internal"));
1718    }
1719
1720    #[test]
1721    fn dot_escapes_metadata_and_event_labels() {
1722        let graph = event_scalar("x\\\"y").named("quoted\"name").to_graph();
1723        let dot = graph.display_dot().to_string();
1724
1725        assert!(dot.contains("x\\\\\\\"y"));
1726        assert!(dot.contains("quoted\\\"name"));
1727    }
1728
1729    #[test]
1730    fn iterative_dot_display_handles_deep_graphs() {
1731        let metadata = ExprMetadata::new(ExprSourceKind::Unary);
1732        let mut nodes = vec![ExprNode::EventScalar(Arc::from("x"))];
1733        let mut metadata_nodes = vec![ExprMetadata::new(ExprSourceKind::Event)];
1734        for index in 1..20_000 {
1735            nodes.push(ExprNode::Unary {
1736                op: UnaryOp::Sin,
1737                input: ExprId::from_index(index - 1),
1738            });
1739            metadata_nodes.push(metadata.clone());
1740        }
1741        let graph = ExprGraph::from_parts(ExprId::from_index(19_999), nodes, metadata_nodes)
1742            .expect("deep graph is valid");
1743
1744        let dot = graph.display_dot().expand_repeated(false).to_string();
1745        assert!(dot.starts_with("digraph ExprGraph {\n"));
1746        assert_eq!(dot.matches(" -> ").count(), 19_999);
1747    }
1748
1749    #[cfg(feature = "svg")]
1750    #[test]
1751    fn dot_display_renders_svg_in_process() {
1752        let svg = shared_graph()
1753            .display_dot()
1754            .with_preset(ColorPreset::Light)
1755            .render_svg()
1756            .unwrap();
1757
1758        assert!(svg.contains("<svg"));
1759        assert!(svg.contains("</svg>"));
1760    }
1761}
laddu_expr/visualization.rs

laddu_expr/
visualization.rs
