laddu_expr/
expression.rs

1use std::{
2    collections::HashMap,
3    fmt,
4    sync::{Arc, OnceLock},
5};
6
7use num::complex::Complex64;
8use serde::{Deserialize, Serialize};
9
10use crate::{ExprGraphError, ExprShapeError, ParamError, ParamResult, parameters::Parameter};
11
12/// Stable identifier for a node in a serialized [`ExprGraph`].
13#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
14pub struct ExprId(u64);
15
16impl ExprId {
17    /// Creates an identifier from a zero-based node index.
18    pub fn from_index(index: usize) -> Self {
19        Self(index as u64)
20    }
21
22    /// Returns the zero-based node index.
23    pub fn index(self) -> usize {
24        self.0 as usize
25    }
26}
27
28/// Runtime value category produced by an expression node.
29#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
30pub enum ValueKind {
31    /// A real scalar.
32    Real,
33    /// A complex scalar.
34    Complex,
35    /// A vector with a fixed number of elements.
36    Vector {
37        /// Number of vector elements.
38        len: usize,
39    },
40    /// A matrix with fixed dimensions.
41    Matrix {
42        /// Number of rows.
43        rows: usize,
44        /// Number of columns.
45        cols: usize,
46    },
47}
48
49/// Structural shape of an expression.
50#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
51pub enum ExprShape {
52    /// A scalar expression.
53    Scalar,
54    /// A vector expression.
55    Vector {
56        /// Number of vector elements.
57        len: usize,
58    },
59    /// A matrix expression.
60    Matrix {
61        /// Number of rows.
62        rows: usize,
63        /// Number of columns.
64        cols: usize,
65    },
66}
67
68impl fmt::Display for ExprShape {
69    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
70        match self {
71            Self::Scalar => write!(f, "scalar"),
72            Self::Vector { len } => write!(f, "vector[{len}]"),
73            Self::Matrix { rows, cols } => write!(f, "matrix[{rows}x{cols}]"),
74        }
75    }
76}
77
78/// Converts a component selector into a zero-based index.
79pub trait ComponentIndex {
80    /// Returns the selected zero-based component index.
81    fn component_index(self) -> usize;
82}
83
84impl ComponentIndex for usize {
85    fn component_index(self) -> usize {
86        self
87    }
88}
89
90impl ComponentIndex for i32 {
91    fn component_index(self) -> usize {
92        usize::try_from(self).expect("component index must be nonnegative")
93    }
94}
95
96#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
97/// A named component of a four-momentum in `(E, px, py, pz)` order.
98pub enum P4Component {
99    /// Energy.
100    E,
101    /// Momentum in the x direction.
102    Px,
103    /// Momentum in the y direction.
104    Py,
105    /// Momentum in the z direction.
106    Pz,
107}
108
109impl P4Component {
110    /// Return the lowercase event-column suffix for this component.
111    pub fn label(self) -> &'static str {
112        match self {
113            Self::E => "e",
114            Self::Px => "px",
115            Self::Py => "py",
116            Self::Pz => "pz",
117        }
118    }
119
120    /// Return the component's position in `(E, px, py, pz)` order.
121    pub fn index(self) -> usize {
122        match self {
123            Self::E => 0,
124            Self::Px => 1,
125            Self::Py => 2,
126            Self::Pz => 3,
127        }
128    }
129}
130
131/// Unary operation in an expression graph.
132#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
133pub enum UnaryOp {
134    /// Arithmetic negation.
135    Neg,
136    /// Real part.
137    Real,
138    /// Imaginary part.
139    Imag,
140    /// Complex conjugate.
141    Conj,
142    /// Squared complex norm.
143    NormSqr,
144    /// Principal square root.
145    Sqrt,
146    /// Exponential.
147    Exp,
148    /// Sine.
149    Sin,
150    /// Cosine.
151    Cos,
152    /// Natural logarithm.
153    Log,
154    /// Integer power.
155    PowI(i32),
156}
157
158impl UnaryOp {
159    /// Applies this operation to a scalar complex value.
160    pub fn evaluate(&self, value: Complex64) -> Complex64 {
161        match self {
162            Self::Neg => -value,
163            Self::Real => Complex64::from(value.re),
164            Self::Imag => Complex64::from(value.im),
165            Self::Conj => value.conj(),
166            Self::NormSqr => Complex64::from(value.norm_sqr()),
167            Self::Sqrt => value.sqrt(),
168            Self::Exp => value.exp(),
169            Self::Sin => value.sin(),
170            Self::Cos => value.cos(),
171            Self::Log => value.ln(),
172            Self::PowI(power) => value.powi(*power),
173        }
174    }
175}
176
177/// Binary operation in an expression graph.
178#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
179pub enum BinaryOp {
180    /// Addition.
181    Add,
182    /// Subtraction.
183    Sub,
184    /// Multiplication.
185    Mul,
186    /// Division.
187    Div,
188    /// Two-argument arctangent of the real parts.
189    Atan2,
190}
191
192impl BinaryOp {
193    /// Applies this operation to two scalar complex values.
194    pub fn evaluate(&self, a: Complex64, b: Complex64) -> Complex64 {
195        match self {
196            Self::Add => a + b,
197            Self::Sub => a - b,
198            Self::Mul => a * b,
199            Self::Div => a / b,
200            Self::Atan2 => Complex64::from(a.re.atan2(b.re)),
201        }
202    }
203}
204
205/// Serialized node in a topologically ordered [`ExprGraph`].
206#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
207pub enum ExprNode {
208    /// A real constant.
209    RealConst(f64),
210    /// A complex constant.
211    ComplexConst(Complex64),
212    /// A scalar fit parameter.
213    ScalarParam(Parameter),
214    /// A named scalar event column.
215    EventScalar(Arc<str>),
216    /// One component of a named event four-momentum.
217    EventP4Component {
218        /// Event column base name.
219        name: Arc<str>,
220        /// Requested four-momentum component.
221        component: P4Component,
222    },
223    /// A unary operation.
224    Unary {
225        /// Operation to apply.
226        op: UnaryOp,
227        /// Input node.
228        input: ExprId,
229    },
230    /// A binary operation.
231    Binary {
232        /// Operation to apply.
233        op: BinaryOp,
234        /// Left operand.
235        lhs: ExprId,
236        /// Right operand.
237        rhs: ExprId,
238    },
239    /// A sum of zero or more terms.
240    NaryAdd {
241        /// Term nodes.
242        terms: Vec<ExprId>,
243    },
244    /// A product of zero or more factors.
245    NaryMul {
246        /// Factor nodes.
247        factors: Vec<ExprId>,
248    },
249    /// A complex scalar assembled from real and imaginary expressions.
250    Complex {
251        /// Real component.
252        re: ExprId,
253        /// Imaginary component.
254        im: ExprId,
255    },
256    /// A vector assembled from scalar elements.
257    Vector {
258        /// Scalar element nodes.
259        elements: Vec<ExprId>,
260    },
261    /// A row-major matrix assembled from scalar elements.
262    Matrix {
263        /// Number of rows.
264        rows: usize,
265        /// Number of columns.
266        cols: usize,
267        /// Row-major scalar elements.
268        elements: Vec<ExprId>,
269    },
270    /// A vector component selection.
271    Component {
272        /// Vector input.
273        input: ExprId,
274        /// Zero-based component index.
275        index: usize,
276    },
277    /// A matrix element selection.
278    MatrixElement {
279        /// Matrix input.
280        input: ExprId,
281        /// Zero-based row index.
282        row: usize,
283        /// Zero-based column index.
284        col: usize,
285    },
286    /// Matrix-matrix multiplication.
287    MatMul {
288        /// Left matrix.
289        lhs: ExprId,
290        /// Right matrix.
291        rhs: ExprId,
292    },
293    /// Matrix-vector multiplication.
294    MatVec {
295        /// Matrix operand.
296        matrix: ExprId,
297        /// Vector operand.
298        vector: ExprId,
299    },
300    /// Vector dot product.
301    Dot {
302        /// Left vector.
303        lhs: ExprId,
304        /// Right vector.
305        rhs: ExprId,
306    },
307    /// Solution of a linear system.
308    Solve {
309        /// Coefficient matrix.
310        matrix: ExprId,
311        /// Right-hand-side vector or matrix.
312        rhs: ExprId,
313    },
314}
315
316impl From<Complex64> for ExprNode {
317    fn from(value: Complex64) -> Self {
318        if value.im == 0.0 {
319            Self::RealConst(value.re)
320        } else {
321            Self::ComplexConst(value)
322        }
323    }
324}
325
326impl ExprNode {
327    /// Creates the most compact constant-node representation for `value`.
328    pub fn from_folded_const(value: Complex64) -> Self {
329        if value.im == 0.0 && value.im.is_sign_positive() {
330            Self::RealConst(value.re)
331        } else {
332            Self::ComplexConst(value)
333        }
334    }
335
336    /// Returns the node's scalar constant value, if it is a constant.
337    pub fn const_value(&self) -> Option<Complex64> {
338        match self {
339            ExprNode::RealConst(value) => Some(Complex64::from(*value)),
340            ExprNode::ComplexConst(value) => Some(*value),
341            _ => None,
342        }
343    }
344
345    /// Returns whether `node` is the scalar constant zero.
346    pub fn is_zero(node: &ExprNode) -> bool {
347        node.const_value()
348            .is_some_and(|value| value == Complex64::ZERO)
349    }
350
351    /// Returns whether `node` is the scalar constant one.
352    pub fn is_one(node: &ExprNode) -> bool {
353        node.const_value()
354            .is_some_and(|value| value == Complex64::ONE)
355    }
356
357    /// Returns the identifiers of this node's direct dependencies.
358    pub fn child_ids(&self) -> Vec<ExprId> {
359        match self {
360            Self::RealConst(_)
361            | Self::ComplexConst(_)
362            | Self::ScalarParam(_)
363            | Self::EventScalar(_)
364            | Self::EventP4Component { .. } => Vec::new(),
365            Self::Unary { input, .. }
366            | Self::Component { input, .. }
367            | Self::MatrixElement { input, .. } => vec![*input],
368            Self::Binary { lhs, rhs, .. }
369            | Self::Complex { re: lhs, im: rhs }
370            | Self::MatMul { lhs, rhs }
371            | Self::Dot { lhs, rhs } => vec![*lhs, *rhs],
372            Self::NaryAdd { terms } => terms.clone(),
373            Self::NaryMul { factors } => factors.clone(),
374            Self::Vector { elements } | Self::Matrix { elements, .. } => elements.clone(),
375            Self::MatVec { matrix, vector } => vec![*matrix, *vector],
376            Self::Solve { matrix, rhs } => vec![*matrix, *rhs],
377        }
378    }
379}
380
381/// Broad origin category recorded in [`ExprMetadata`].
382#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
383pub enum ExprSourceKind {
384    /// Constant literal.
385    Const,
386    /// Fit parameter.
387    Param,
388    /// Event data.
389    Event,
390    /// Unary operation.
391    Unary,
392    /// Binary or n-ary operation.
393    Binary,
394    /// Complex-number construction.
395    Complex,
396    /// Vector construction or selection.
397    Vector,
398    /// Matrix construction or selection.
399    Matrix,
400    /// Linear-algebra operation.
401    LinearAlgebra,
402}
403
404/// User-facing annotations and origin information for an expression node.
405#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
406pub struct ExprMetadata {
407    source: ExprSourceKind,
408    name: Option<Arc<str>>,
409    tags: Vec<Arc<str>>,
410}
411
412impl ExprMetadata {
413    /// Creates metadata for the given source category.
414    pub fn new(source: ExprSourceKind) -> Self {
415        Self {
416            source,
417            name: None,
418            tags: Vec::new(),
419        }
420    }
421
422    /// Returns the node's source category.
423    pub fn source(&self) -> ExprSourceKind {
424        self.source
425    }
426
427    /// Returns the optional user-assigned name.
428    pub fn name(&self) -> Option<&str> {
429        self.name.as_deref()
430    }
431
432    /// Returns the user-assigned tags.
433    pub fn tags(&self) -> &[Arc<str>] {
434        &self.tags
435    }
436
437    /// Returns whether the metadata contains `tag`.
438    pub fn has_tag(&self, tag: &str) -> bool {
439        self.tags.iter().any(|candidate| candidate.as_ref() == tag)
440    }
441}
442
443/// Shareable symbolic expression represented internally as a directed acyclic graph.
444#[derive(Clone, Debug)]
445pub struct Expr {
446    node: Arc<DagNode>,
447}
448
449#[derive(Clone, Debug)]
450struct DagNode {
451    kind: DagNodeKind,
452    metadata: ExprMetadata,
453    shape: OnceLock<Result<ExprShape, ExprShapeError>>,
454}
455
456#[derive(Clone, Debug)]
457enum DagNodeKind {
458    RealConst(f64),
459    ComplexConst(Complex64),
460    ScalarParam(Parameter),
461    EventScalar(Arc<str>),
462    EventP4Component {
463        name: Arc<str>,
464        component: P4Component,
465    },
466    Unary {
467        op: UnaryOp,
468        input: Expr,
469    },
470    Binary {
471        op: BinaryOp,
472        lhs: Expr,
473        rhs: Expr,
474    },
475    Complex {
476        re: Expr,
477        im: Expr,
478    },
479    Vector {
480        elements: Vec<Expr>,
481    },
482    Matrix {
483        rows: usize,
484        cols: usize,
485        elements: Vec<Expr>,
486    },
487    Component {
488        input: Expr,
489        index: usize,
490    },
491    MatrixElement {
492        input: Expr,
493        row: usize,
494        col: usize,
495    },
496    MatMul {
497        lhs: Expr,
498        rhs: Expr,
499    },
500    MatVec {
501        matrix: Expr,
502        vector: Expr,
503    },
504    Dot {
505        lhs: Expr,
506        rhs: Expr,
507    },
508    Solve {
509        matrix: Expr,
510        rhs: Expr,
511    },
512}
513
514impl Expr {
515    fn new(kind: DagNodeKind) -> Self {
516        let source = source_kind(&kind);
517        Self {
518            node: Arc::new(DagNode {
519                kind,
520                metadata: ExprMetadata::new(source),
521                shape: OnceLock::new(),
522            }),
523        }
524    }
525
526    /// Assigns a display name to the expression root.
527    pub fn named(self, name: impl Into<Arc<str>>) -> Self {
528        self.with_metadata(|metadata| metadata.name = Some(name.into()))
529    }
530
531    /// Adds a tag to the expression root.
532    pub fn tagged(self, tag: impl Into<Arc<str>>) -> Self {
533        let tag = tag.into();
534        self.with_metadata(|metadata| {
535            if !metadata.tags.iter().any(|existing| existing == &tag) {
536                metadata.tags.push(tag);
537            }
538        })
539    }
540
541    /// Adds each supplied tag to the expression root.
542    pub fn tagged_with(self, tags: impl IntoIterator<Item = impl Into<Arc<str>>>) -> Self {
543        tags.into_iter().fold(self, Self::tagged)
544    }
545
546    /// Replace tagged components that do not match any requested tag with zero.
547    ///
548    /// Untagged nodes remain active, while a matching tagged node retains its complete subtree.
549    pub fn project_tags<'a>(&self, tags: impl IntoIterator<Item = &'a str>) -> Self {
550        let tags: Vec<_> = tags.into_iter().collect();
551        self.project_tags_inner(&tags)
552    }
553
554    fn project_tags_inner(&self, tags: &[&str]) -> Self {
555        if !self.node.metadata.tags.is_empty() {
556            return if self
557                .node
558                .metadata
559                .tags
560                .iter()
561                .any(|candidate| tags.contains(&candidate.as_ref()))
562            {
563                self.clone()
564            } else {
565                self.zero_like()
566            };
567        }
568
569        let kind = match &self.node.kind {
570            DagNodeKind::RealConst(_)
571            | DagNodeKind::ComplexConst(_)
572            | DagNodeKind::ScalarParam(_)
573            | DagNodeKind::EventScalar(_)
574            | DagNodeKind::EventP4Component { .. } => return self.clone(),
575            DagNodeKind::Unary { op, input } => DagNodeKind::Unary {
576                op: *op,
577                input: input.project_tags_inner(tags),
578            },
579            DagNodeKind::Binary { op, lhs, rhs } => DagNodeKind::Binary {
580                op: *op,
581                lhs: lhs.project_tags_inner(tags),
582                rhs: rhs.project_tags_inner(tags),
583            },
584            DagNodeKind::Complex { re, im } => DagNodeKind::Complex {
585                re: re.project_tags_inner(tags),
586                im: im.project_tags_inner(tags),
587            },
588            DagNodeKind::Vector { elements } => DagNodeKind::Vector {
589                elements: elements
590                    .iter()
591                    .map(|value| value.project_tags_inner(tags))
592                    .collect(),
593            },
594            DagNodeKind::Matrix {
595                rows,
596                cols,
597                elements,
598            } => DagNodeKind::Matrix {
599                rows: *rows,
600                cols: *cols,
601                elements: elements
602                    .iter()
603                    .map(|value| value.project_tags_inner(tags))
604                    .collect(),
605            },
606            DagNodeKind::Component { input, index } => DagNodeKind::Component {
607                input: input.project_tags_inner(tags),
608                index: *index,
609            },
610            DagNodeKind::MatrixElement { input, row, col } => DagNodeKind::MatrixElement {
611                input: input.project_tags_inner(tags),
612                row: *row,
613                col: *col,
614            },
615            DagNodeKind::MatMul { lhs, rhs } => DagNodeKind::MatMul {
616                lhs: lhs.project_tags_inner(tags),
617                rhs: rhs.project_tags_inner(tags),
618            },
619            DagNodeKind::MatVec { matrix, vector } => DagNodeKind::MatVec {
620                matrix: matrix.project_tags_inner(tags),
621                vector: vector.project_tags_inner(tags),
622            },
623            DagNodeKind::Dot { lhs, rhs } => DagNodeKind::Dot {
624                lhs: lhs.project_tags_inner(tags),
625                rhs: rhs.project_tags_inner(tags),
626            },
627            DagNodeKind::Solve { matrix, rhs } => DagNodeKind::Solve {
628                matrix: matrix.project_tags_inner(tags),
629                rhs: rhs.project_tags_inner(tags),
630            },
631        };
632        Expr::new(kind).with_metadata(|metadata| *metadata = self.node.metadata.clone())
633    }
634
635    fn zero_like(&self) -> Self {
636        match self
637            .shape()
638            .expect("valid expression shapes are cached eagerly")
639        {
640            ExprShape::Scalar => Expr::from(0.0),
641            ExprShape::Vector { len } => vector((0..len).map(|_| Expr::from(0.0))),
642            ExprShape::Matrix { rows, cols } => {
643                matrix_from_flat(rows, cols, (0..rows * cols).map(|_| Expr::from(0.0)))
644                    .expect("zero matrix dimensions match")
645            }
646        }
647    }
648
649    /// Returns an expression for the real part.
650    pub fn real(&self) -> Self {
651        unary(UnaryOp::Real, self)
652    }
653
654    /// Returns an expression for the imaginary part.
655    pub fn imag(&self) -> Self {
656        unary(UnaryOp::Imag, self)
657    }
658
659    /// Returns an expression for the complex conjugate.
660    pub fn conj(&self) -> Self {
661        unary(UnaryOp::Conj, self)
662    }
663
664    /// Returns an expression for the squared complex norm.
665    pub fn norm_sqr(&self) -> Self {
666        unary(UnaryOp::NormSqr, self)
667    }
668
669    /// Returns an expression for the principal square root.
670    pub fn sqrt(&self) -> Self {
671        unary(UnaryOp::Sqrt, self)
672    }
673
674    /// Returns an expression for the exponential.
675    pub fn exp(&self) -> Self {
676        unary(UnaryOp::Exp, self)
677    }
678
679    /// Returns an expression for the sine.
680    pub fn sin(&self) -> Self {
681        unary(UnaryOp::Sin, self)
682    }
683
684    /// Returns an expression for the cosine.
685    pub fn cos(&self) -> Self {
686        unary(UnaryOp::Cos, self)
687    }
688
689    /// Returns an expression for the principal arccosine.
690    pub fn acos(&self) -> Self {
691        atan2((Expr::from(1.0) - self.powi(2)).sqrt(), self)
692    }
693
694    /// Returns an expression for the natural logarithm.
695    pub fn log(&self) -> Self {
696        unary(UnaryOp::Log, self)
697    }
698
699    /// Returns an expression raised to an integer power.
700    pub fn powi(&self, power: i32) -> Self {
701        unary(UnaryOp::PowI(power), self)
702    }
703
704    /// Selects a component from a vector-valued expression.
705    pub fn component(&self, index: impl ComponentIndex) -> Self {
706        Expr::new(DagNodeKind::Component {
707            input: self.clone(),
708            index: index.component_index(),
709        })
710    }
711
712    /// Selects an element from a matrix-valued expression.
713    pub fn matrix_element(&self, row: usize, col: usize) -> Self {
714        Expr::new(DagNodeKind::MatrixElement {
715            input: self.clone(),
716            row,
717            col,
718        })
719    }
720
721    /// Serializes the shareable expression DAG into a topologically ordered graph.
722    pub fn to_graph(&self) -> ExprGraph {
723        GraphBuilder::new().build(self)
724    }
725
726    /// Rebuilds a shareable expression DAG from its serialized graph form.
727    ///
728    /// # Errors
729    ///
730    /// Returns [`ExprGraphError`] when the graph is empty, its root or a child
731    /// identifier is invalid, its metadata length does not match its node
732    /// count, or its nodes are not topologically ordered.
733    pub fn from_graph(graph: ExprGraph) -> Result<Self, ExprGraphError> {
734        let ExprGraph {
735            root,
736            nodes,
737            metadata,
738        } = graph;
739        let graph = ExprGraph::from_parts(root, nodes, metadata)?;
740        let mut expressions: Vec<Expr> = Vec::with_capacity(graph.nodes.len());
741        for (index, node) in graph.nodes.iter().enumerate() {
742            let child = |id: ExprId| expressions[id.index()].clone();
743            let expression = match node {
744                ExprNode::RealConst(value) => Expr::new(DagNodeKind::RealConst(*value)),
745                ExprNode::ComplexConst(value) => Expr::new(DagNodeKind::ComplexConst(*value)),
746                ExprNode::ScalarParam(parameter) => {
747                    Expr::new(DagNodeKind::ScalarParam(parameter.clone()))
748                }
749                ExprNode::EventScalar(name) => {
750                    Expr::new(DagNodeKind::EventScalar(Arc::clone(name)))
751                }
752                ExprNode::EventP4Component { name, component } => {
753                    Expr::new(DagNodeKind::EventP4Component {
754                        name: Arc::clone(name),
755                        component: *component,
756                    })
757                }
758                ExprNode::Unary { op, input } => Expr::new(DagNodeKind::Unary {
759                    op: *op,
760                    input: child(*input),
761                }),
762                ExprNode::Binary { op, lhs, rhs } => Expr::new(DagNodeKind::Binary {
763                    op: *op,
764                    lhs: child(*lhs),
765                    rhs: child(*rhs),
766                }),
767                ExprNode::NaryAdd { terms } => terms
768                    .iter()
769                    .map(|id| child(*id))
770                    .reduce(|lhs, rhs| binary(BinaryOp::Add, &lhs, &rhs))
771                    .unwrap_or_else(|| Expr::from(0.0)),
772                ExprNode::NaryMul { factors } => factors
773                    .iter()
774                    .map(|id| child(*id))
775                    .reduce(|lhs, rhs| binary(BinaryOp::Mul, &lhs, &rhs))
776                    .unwrap_or_else(|| Expr::from(1.0)),
777                ExprNode::Complex { re, im } => Expr::new(DagNodeKind::Complex {
778                    re: child(*re),
779                    im: child(*im),
780                }),
781                ExprNode::Vector { elements } => Expr::new(DagNodeKind::Vector {
782                    elements: elements.iter().map(|id| child(*id)).collect(),
783                }),
784                ExprNode::Matrix {
785                    rows,
786                    cols,
787                    elements,
788                } => Expr::new(DagNodeKind::Matrix {
789                    rows: *rows,
790                    cols: *cols,
791                    elements: elements.iter().map(|id| child(*id)).collect(),
792                }),
793                ExprNode::Component { input, index } => Expr::new(DagNodeKind::Component {
794                    input: child(*input),
795                    index: *index,
796                }),
797                ExprNode::MatrixElement { input, row, col } => {
798                    Expr::new(DagNodeKind::MatrixElement {
799                        input: child(*input),
800                        row: *row,
801                        col: *col,
802                    })
803                }
804                ExprNode::MatMul { lhs, rhs } => Expr::new(DagNodeKind::MatMul {
805                    lhs: child(*lhs),
806                    rhs: child(*rhs),
807                }),
808                ExprNode::MatVec { matrix, vector } => Expr::new(DagNodeKind::MatVec {
809                    matrix: child(*matrix),
810                    vector: child(*vector),
811                }),
812                ExprNode::Dot { lhs, rhs } => Expr::new(DagNodeKind::Dot {
813                    lhs: child(*lhs),
814                    rhs: child(*rhs),
815                }),
816                ExprNode::Solve { matrix, rhs } => Expr::new(DagNodeKind::Solve {
817                    matrix: child(*matrix),
818                    rhs: child(*rhs),
819                }),
820            };
821            let mut dag = (*expression.node).clone();
822            dag.metadata = graph.metadata[index].clone();
823            expressions.push(Expr {
824                node: Arc::new(dag),
825            });
826        }
827        Ok(expressions[graph.root.index()].clone())
828    }
829
830    /// Determines and validates the expression's structural shape.
831    ///
832    /// # Errors
833    ///
834    /// Returns [`ExprShapeError`] when this expression contains an operation
835    /// whose operand shapes are incompatible.
836    pub fn shape(&self) -> Result<ExprShape, ExprShapeError> {
837        self.node
838            .shape
839            .get_or_init(|| self.node.kind.shape())
840            .clone()
841    }
842
843    fn with_metadata(self, f: impl FnOnce(&mut ExprMetadata)) -> Self {
844        let mut node = (*self.node).clone();
845        f(&mut node.metadata);
846        Self {
847            node: Arc::new(node),
848        }
849    }
850}
851
852impl Serialize for Expr {
853    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
854    where
855        S: serde::Serializer,
856    {
857        self.to_graph().serialize(serializer)
858    }
859}
860
861impl<'de> Deserialize<'de> for Expr {
862    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
863    where
864        D: serde::Deserializer<'de>,
865    {
866        Expr::from_graph(ExprGraph::deserialize(deserializer)?).map_err(serde::de::Error::custom)
867    }
868}
869
870impl DagNodeKind {
871    fn shape(&self) -> Result<ExprShape, ExprShapeError> {
872        match self {
873            Self::RealConst(_)
874            | Self::ComplexConst(_)
875            | Self::ScalarParam(_)
876            | Self::EventScalar(_)
877            | Self::EventP4Component { .. } => Ok(ExprShape::Scalar),
878            Self::Unary { input, .. } => {
879                input.expect_shape("unary operation", ExprShape::Scalar)?;
880                Ok(ExprShape::Scalar)
881            }
882            Self::Binary { lhs, rhs, .. } => {
883                lhs.expect_shape("binary operation", ExprShape::Scalar)?;
884                rhs.expect_shape("binary operation", ExprShape::Scalar)?;
885                Ok(ExprShape::Scalar)
886            }
887            Self::Complex { re, im } => {
888                re.expect_shape("complex constructor", ExprShape::Scalar)?;
889                im.expect_shape("complex constructor", ExprShape::Scalar)?;
890                Ok(ExprShape::Scalar)
891            }
892            Self::Vector { elements } => {
893                for element in elements {
894                    element.expect_shape("vector constructor", ExprShape::Scalar)?;
895                }
896                Ok(ExprShape::Vector {
897                    len: elements.len(),
898                })
899            }
900            Self::Matrix {
901                rows,
902                cols,
903                elements,
904            } => {
905                let expected = rows.checked_mul(*cols).ok_or_else(|| {
906                    ExprShapeError::new("matrix constructor", "row/column product overflowed")
907                })?;
908                if elements.len() != expected {
909                    return Err(ExprShapeError::new(
910                        "matrix constructor",
911                        format!(
912                            "shape {rows}x{cols} requires {expected} elements, got {}",
913                            elements.len()
914                        ),
915                    ));
916                }
917                for element in elements {
918                    element.expect_shape("matrix constructor", ExprShape::Scalar)?;
919                }
920                Ok(ExprShape::Matrix {
921                    rows: *rows,
922                    cols: *cols,
923                })
924            }
925            Self::Component { input, index } => {
926                let ExprShape::Vector { len } = input.shape()? else {
927                    return Err(ExprShapeError::new(
928                        "component",
929                        format!("expected vector, got {}", input.shape()?),
930                    ));
931                };
932                if *index >= len {
933                    return Err(ExprShapeError::new(
934                        "component",
935                        format!("index {index} is out of bounds for vector[{len}]"),
936                    ));
937                }
938                Ok(ExprShape::Scalar)
939            }
940            Self::MatrixElement { input, row, col } => {
941                let ExprShape::Matrix { rows, cols } = input.shape()? else {
942                    return Err(ExprShapeError::new(
943                        "matrix element",
944                        format!("expected matrix, got {}", input.shape()?),
945                    ));
946                };
947                if *row >= rows || *col >= cols {
948                    return Err(ExprShapeError::new(
949                        "matrix element",
950                        format!("index ({row}, {col}) is out of bounds for matrix[{rows}x{cols}]"),
951                    ));
952                }
953                Ok(ExprShape::Scalar)
954            }
955            Self::MatMul { lhs, rhs } => {
956                let ExprShape::Matrix {
957                    rows: lhs_rows,
958                    cols: lhs_cols,
959                } = lhs.shape()?
960                else {
961                    return Err(ExprShapeError::new(
962                        "matrix multiplication",
963                        format!("left input must be a matrix, got {}", lhs.shape()?),
964                    ));
965                };
966                let ExprShape::Matrix {
967                    rows: rhs_rows,
968                    cols: rhs_cols,
969                } = rhs.shape()?
970                else {
971                    return Err(ExprShapeError::new(
972                        "matrix multiplication",
973                        format!("right input must be a matrix, got {}", rhs.shape()?),
974                    ));
975                };
976                if lhs_cols != rhs_rows {
977                    return Err(ExprShapeError::new(
978                        "matrix multiplication",
979                        format!("cannot multiply {lhs_rows}x{lhs_cols} by {rhs_rows}x{rhs_cols}"),
980                    ));
981                }
982                Ok(ExprShape::Matrix {
983                    rows: lhs_rows,
984                    cols: rhs_cols,
985                })
986            }
987            Self::MatVec { matrix, vector } => {
988                let ExprShape::Matrix { rows, cols } = matrix.shape()? else {
989                    return Err(ExprShapeError::new(
990                        "matrix-vector multiplication",
991                        format!("left input must be a matrix, got {}", matrix.shape()?),
992                    ));
993                };
994                let ExprShape::Vector { len } = vector.shape()? else {
995                    return Err(ExprShapeError::new(
996                        "matrix-vector multiplication",
997                        format!("right input must be a vector, got {}", vector.shape()?),
998                    ));
999                };
1000                if cols != len {
1001                    return Err(ExprShapeError::new(
1002                        "matrix-vector multiplication",
1003                        format!("cannot multiply {rows}x{cols} matrix by vector[{len}]"),
1004                    ));
1005                }
1006                Ok(ExprShape::Vector { len: rows })
1007            }
1008            Self::Dot { lhs, rhs } => {
1009                let ExprShape::Vector { len: lhs_len } = lhs.shape()? else {
1010                    return Err(ExprShapeError::new(
1011                        "dot product",
1012                        format!("left input must be a vector, got {}", lhs.shape()?),
1013                    ));
1014                };
1015                let ExprShape::Vector { len: rhs_len } = rhs.shape()? else {
1016                    return Err(ExprShapeError::new(
1017                        "dot product",
1018                        format!("right input must be a vector, got {}", rhs.shape()?),
1019                    ));
1020                };
1021                if lhs_len != rhs_len {
1022                    return Err(ExprShapeError::new(
1023                        "dot product",
1024                        format!("vector lengths differ: {lhs_len} and {rhs_len}"),
1025                    ));
1026                }
1027                Ok(ExprShape::Scalar)
1028            }
1029            Self::Solve { matrix, rhs } => {
1030                let ExprShape::Matrix { rows, cols } = matrix.shape()? else {
1031                    return Err(ExprShapeError::new(
1032                        "linear solve",
1033                        format!("left input must be a matrix, got {}", matrix.shape()?),
1034                    ));
1035                };
1036                let ExprShape::Vector { len } = rhs.shape()? else {
1037                    return Err(ExprShapeError::new(
1038                        "linear solve",
1039                        format!("right input must be a vector, got {}", rhs.shape()?),
1040                    ));
1041                };
1042                if rows != cols || rows != len {
1043                    return Err(ExprShapeError::new(
1044                        "linear solve",
1045                        format!("cannot solve matrix[{rows}x{cols}] against vector[{len}]"),
1046                    ));
1047                }
1048                Ok(ExprShape::Vector { len })
1049            }
1050        }
1051    }
1052}
1053
1054impl Expr {
1055    fn expect_shape(
1056        &self,
1057        operation: &'static str,
1058        expected: ExprShape,
1059    ) -> Result<(), ExprShapeError> {
1060        let actual = self.shape()?;
1061        if actual != expected {
1062            return Err(ExprShapeError::new(
1063                operation,
1064                format!("expected {expected}, got {actual}"),
1065            ));
1066        }
1067        Ok(())
1068    }
1069}
1070
1071auto_ops::impl_op_ex!(+ |a: &Expr, b: &Expr| -> Expr { binary(BinaryOp::Add, a, b) });
1072
1073auto_ops::impl_op_ex!(+ |a: &Expr, b: &f64| -> Expr { binary(BinaryOp::Add, a, b) });
1074auto_ops::impl_op_ex!(+ |a: &f64, b: &Expr| -> Expr { binary(BinaryOp::Add, a, b) });
1075
1076auto_ops::impl_op_ex!(+ |a: &Expr, b: &Complex64| -> Expr { binary(BinaryOp::Add, a, b) });
1077auto_ops::impl_op_ex!(+ |a: &Complex64, b: &Expr| -> Expr { binary(BinaryOp::Add, a, b) });
1078
1079auto_ops::impl_op_ex!(+ |a: &Expr, b: &Parameter| -> Expr { binary(BinaryOp::Add, a, b) });
1080auto_ops::impl_op_ex!(+ |a: &Parameter, b: &Expr| -> Expr { binary(BinaryOp::Add, a, b) });
1081
1082auto_ops::impl_op_ex!(+ |a: &Parameter, b: &f64| -> Expr { binary(BinaryOp::Add, a, b) });
1083auto_ops::impl_op_ex!(+ |a: &f64, b: &Parameter| -> Expr { binary(BinaryOp::Add, a, b) });
1084
1085auto_ops::impl_op_ex!(+ |a: &Parameter, b: &Complex64| -> Expr { binary(BinaryOp::Add, a, b) });
1086auto_ops::impl_op_ex!(+ |a: &Complex64, b: &Parameter| -> Expr { binary(BinaryOp::Add, a, b) });
1087
1088auto_ops::impl_op_ex!(+ |a: &Parameter, b: &Parameter| -> Expr { binary(BinaryOp::Add, a, b) });
1089
1090auto_ops::impl_op_ex!(-|a: &Expr, b: &Expr| -> Expr { binary(BinaryOp::Sub, a, b) });
1091
1092auto_ops::impl_op_ex!(-|a: &Expr, b: &f64| -> Expr { binary(BinaryOp::Sub, a, b) });
1093
1094auto_ops::impl_op_ex!(-|a: &f64, b: &Expr| -> Expr { binary(BinaryOp::Sub, a, b) });
1095
1096auto_ops::impl_op_ex!(-|a: &Expr, b: &Complex64| -> Expr { binary(BinaryOp::Sub, a, b) });
1097
1098auto_ops::impl_op_ex!(-|a: &Complex64, b: &Expr| -> Expr { binary(BinaryOp::Sub, a, b) });
1099
1100auto_ops::impl_op_ex!(-|a: &Expr, b: &Parameter| -> Expr { binary(BinaryOp::Sub, a, b) });
1101
1102auto_ops::impl_op_ex!(-|a: &Parameter, b: &Expr| -> Expr { binary(BinaryOp::Sub, a, b) });
1103
1104auto_ops::impl_op_ex!(-|a: &f64, b: &Parameter| -> Expr { binary(BinaryOp::Sub, a, b) });
1105
1106auto_ops::impl_op_ex!(-|a: &Parameter, b: &f64| -> Expr { binary(BinaryOp::Sub, a, b) });
1107
1108auto_ops::impl_op_ex!(-|a: &Complex64, b: &Parameter| -> Expr { binary(BinaryOp::Sub, a, b) });
1109
1110auto_ops::impl_op_ex!(-|a: &Parameter, b: &Complex64| -> Expr { binary(BinaryOp::Sub, a, b) });
1111
1112auto_ops::impl_op_ex!(-|a: &Parameter, b: &Parameter| -> Expr { binary(BinaryOp::Sub, a, b) });
1113
1114auto_ops::impl_op_ex!(*|a: &Expr, b: &Expr| -> Expr { binary(BinaryOp::Mul, a, b) });
1115
1116auto_ops::impl_op_ex!(*|a: &Expr, b: &f64| -> Expr { binary(BinaryOp::Mul, a, b) });
1117auto_ops::impl_op_ex!(*|a: &f64, b: &Expr| -> Expr { binary(BinaryOp::Mul, a, b) });
1118
1119auto_ops::impl_op_ex!(*|a: &Expr, b: &Complex64| -> Expr { binary(BinaryOp::Mul, a, b) });
1120auto_ops::impl_op_ex!(*|a: &Complex64, b: &Expr| -> Expr { binary(BinaryOp::Mul, a, b) });
1121
1122auto_ops::impl_op_ex!(*|a: &Expr, b: &Parameter| -> Expr { binary(BinaryOp::Mul, a, b) });
1123auto_ops::impl_op_ex!(*|a: &Parameter, b: &Expr| -> Expr { binary(BinaryOp::Mul, a, b) });
1124
1125auto_ops::impl_op_ex!(*|a: &f64, b: &Parameter| -> Expr { binary(BinaryOp::Mul, a, b) });
1126auto_ops::impl_op_ex!(*|a: &Parameter, b: &f64| -> Expr { binary(BinaryOp::Mul, a, b) });
1127
1128auto_ops::impl_op_ex!(*|a: &Complex64, b: &Parameter| -> Expr { binary(BinaryOp::Mul, a, b) });
1129auto_ops::impl_op_ex!(*|a: &Parameter, b: &Complex64| -> Expr { binary(BinaryOp::Mul, a, b) });
1130
1131auto_ops::impl_op_ex!(*|a: &Parameter, b: &Parameter| -> Expr { binary(BinaryOp::Mul, a, b) });
1132
1133auto_ops::impl_op_ex!(/ |a: &Expr, b: &Expr| -> Expr {
1134    binary(BinaryOp::Div, a, b)
1135});
1136
1137auto_ops::impl_op_ex!(/ |a: &Expr, b: &Complex64| -> Expr { binary(BinaryOp::Div, a, b) });
1138auto_ops::impl_op_ex!(/ |a: &Complex64, b: &Expr| -> Expr { binary(BinaryOp::Div, a, b) });
1139
1140auto_ops::impl_op_ex!(/ |a: &Expr, b: &f64| -> Expr { binary(BinaryOp::Div, a, b) });
1141auto_ops::impl_op_ex!(/ |a: &f64, b: &Expr| -> Expr { binary(BinaryOp::Div, a, b) });
1142
1143auto_ops::impl_op_ex!(/|a: &Expr, b: &Parameter| -> Expr {
1144    binary(BinaryOp::Div, a, b)
1145});
1146auto_ops::impl_op_ex!(/|a: &Parameter, b: &Expr| -> Expr {
1147    binary(BinaryOp::Div, a, b)
1148});
1149
1150auto_ops::impl_op_ex!(/|a: &f64, b: &Parameter| -> Expr {
1151    binary(BinaryOp::Div, a, b)
1152});
1153auto_ops::impl_op_ex!(/|a: &Parameter, b: &f64| -> Expr {
1154    binary(BinaryOp::Div, a, b)
1155});
1156
1157auto_ops::impl_op_ex!(/|a: &Complex64, b: &Parameter| -> Expr {
1158    binary(BinaryOp::Div, a, b)
1159});
1160auto_ops::impl_op_ex!(/|a: &Parameter, b: &Complex64| -> Expr {
1161    binary(BinaryOp::Div, a, b)
1162});
1163
1164auto_ops::impl_op_ex!(/|a: &Parameter, b: &Parameter| -> Expr {
1165    binary(BinaryOp::Div, a, b)
1166});
1167
1168auto_ops::impl_op_ex!(-|a: &Expr| -> Expr { unary(UnaryOp::Neg, a) });
1169auto_ops::impl_op_ex!(-|a: &Parameter| -> Expr { unary(UnaryOp::Neg, a) });
1170
1171auto_ops::impl_op_ex!(+= |a: &mut Expr, b: &Expr| {
1172    *a = binary(BinaryOp::Add, &*a, b);
1173});
1174auto_ops::impl_op_ex!(+= |a: &mut Expr, b: &f64| {
1175    *a = binary(BinaryOp::Add, &*a, b);
1176});
1177auto_ops::impl_op_ex!(+= |a: &mut Expr, b: &Complex64| {
1178    *a = binary(BinaryOp::Add, &*a, b);
1179});
1180auto_ops::impl_op_ex!(+= |a: &mut Expr, b: &Parameter| {
1181    *a = binary(BinaryOp::Add, &*a, b);
1182});
1183
1184auto_ops::impl_op_ex!(-= |a: &mut Expr, b: &Expr| {
1185    *a = binary(BinaryOp::Sub, &*a, b);
1186});
1187auto_ops::impl_op_ex!(-= |a: &mut Expr, b: &f64| {
1188    *a = binary(BinaryOp::Sub, &*a, b);
1189});
1190auto_ops::impl_op_ex!(-= |a: &mut Expr, b: &Complex64| {
1191    *a = binary(BinaryOp::Sub, &*a, b);
1192});
1193auto_ops::impl_op_ex!(-= |a: &mut Expr, b: &Parameter| {
1194    *a = binary(BinaryOp::Sub, &*a, b);
1195});
1196
1197auto_ops::impl_op_ex!(*= |a: &mut Expr, b: &Expr| {
1198    *a = binary(BinaryOp::Mul, &*a, b);
1199});
1200auto_ops::impl_op_ex!(*= |a: &mut Expr, b: &f64| {
1201    *a = binary(BinaryOp::Mul, &*a, b);
1202});
1203auto_ops::impl_op_ex!(*= |a: &mut Expr, b: &Complex64| {
1204    *a = binary(BinaryOp::Mul, &*a, b);
1205});
1206auto_ops::impl_op_ex!(*= |a: &mut Expr, b: &Parameter| {
1207    *a = binary(BinaryOp::Mul, &*a, b);
1208});
1209
1210auto_ops::impl_op_ex!(/= |a: &mut Expr, b: &Expr| {
1211    *a = binary(BinaryOp::Div, &*a, b);
1212});
1213auto_ops::impl_op_ex!(/= |a: &mut Expr, b: &f64| {
1214    *a = binary(BinaryOp::Div, &*a, b);
1215});
1216auto_ops::impl_op_ex!(/= |a: &mut Expr, b: &Complex64| {
1217    *a = binary(BinaryOp::Div, &*a, b);
1218});
1219auto_ops::impl_op_ex!(/= |a: &mut Expr, b: &Parameter| {
1220    *a = binary(BinaryOp::Div, &*a, b);
1221});
1222
1223impl From<f64> for Expr {
1224    fn from(value: f64) -> Self {
1225        Self::new(DagNodeKind::RealConst(value))
1226    }
1227}
1228
1229impl From<&f64> for Expr {
1230    fn from(value: &f64) -> Self {
1231        Self::new(DagNodeKind::RealConst(*value))
1232    }
1233}
1234
1235impl From<Complex64> for Expr {
1236    fn from(value: Complex64) -> Self {
1237        Self::new(DagNodeKind::ComplexConst(value))
1238    }
1239}
1240
1241impl From<&Complex64> for Expr {
1242    fn from(value: &Complex64) -> Self {
1243        Self::new(DagNodeKind::ComplexConst(*value))
1244    }
1245}
1246
1247impl From<&Expr> for Expr {
1248    fn from(value: &Expr) -> Self {
1249        value.clone()
1250    }
1251}
1252
1253impl From<Parameter> for Expr {
1254    fn from(parameter: Parameter) -> Self {
1255        Expr::new(DagNodeKind::ScalarParam(parameter))
1256    }
1257}
1258
1259impl From<&Parameter> for Expr {
1260    fn from(parameter: &Parameter) -> Self {
1261        parameter.clone().into()
1262    }
1263}
1264
1265/// Constructs `cos(phase) + i sin(phase)`.
1266pub fn cis(phase: Expr) -> Expr {
1267    phase.cos() + Complex64::I * phase.sin()
1268}
1269
1270/// Constructs a complex scalar from real and imaginary expressions.
1271pub fn complex(re: impl Into<Expr>, im: impl Into<Expr>) -> Expr {
1272    Expr::new(DagNodeKind::Complex {
1273        re: re.into(),
1274        im: im.into(),
1275    })
1276}
1277
1278/// Constructs a complex scalar from magnitude and phase expressions.
1279pub fn polar_complex(mag: impl Into<Expr>, phase: impl Into<Expr>) -> Expr {
1280    mag.into() * (Complex64::I * phase.into()).exp()
1281}
1282
1283/// References a named scalar column in each event.
1284pub fn event_scalar(name: impl Into<Arc<str>>) -> Expr {
1285    Expr::new(DagNodeKind::EventScalar(name.into()))
1286}
1287
1288/// References one component of a named event four-momentum.
1289pub fn event_p4_component(name: impl Into<Arc<str>>, component: P4Component) -> Expr {
1290    Expr::new(DagNodeKind::EventP4Component {
1291        name: name.into(),
1292        component,
1293    })
1294}
1295
1296/// Constructs the two-argument arctangent `atan2(y, x)`.
1297pub fn atan2(y: impl Into<Expr>, x: impl Into<Expr>) -> Expr {
1298    binary(BinaryOp::Atan2, y, x)
1299}
1300
1301/// Constructs the principal arccosine of `value`.
1302pub fn acos(value: impl Into<Expr>) -> Expr {
1303    value.into().acos()
1304}
1305
1306/// Constructs a vector expression from scalar elements.
1307pub fn vector<E>(elements: impl IntoIterator<Item = E>) -> Expr
1308where
1309    E: Into<Expr>,
1310    Expr: From<E>,
1311{
1312    Expr::new(DagNodeKind::Vector {
1313        elements: elements.into_iter().map(Expr::from).collect(),
1314    })
1315}
1316
1317/// Constructs a row-major matrix expression from a nested array.
1318pub fn matrix<const R: usize, const C: usize, E>(elements: [[E; C]; R]) -> Expr
1319where
1320    E: Into<Expr>,
1321    Expr: From<E>,
1322{
1323    Expr::new(DagNodeKind::Matrix {
1324        rows: R,
1325        cols: C,
1326        elements: elements.into_iter().flatten().map(Expr::from).collect(),
1327    })
1328}
1329
1330/// Constructs a row-major matrix from a flat sequence.
1331///
1332/// # Errors
1333///
1334/// Returns [`ExprShapeError`] when either dimension is zero, the dimension
1335/// product overflows, the element count differs from `rows * cols`, or an
1336/// element is not scalar-valued.
1337pub fn matrix_from_flat<E>(
1338    rows: usize,
1339    cols: usize,
1340    elements: impl IntoIterator<Item = E>,
1341) -> Result<Expr, ExprShapeError>
1342where
1343    E: Into<Expr>,
1344    Expr: From<E>,
1345{
1346    if rows == 0 || cols == 0 {
1347        return Err(ExprShapeError::new(
1348            "matrix constructor",
1349            format!("matrix dimensions must be nonzero, got {rows}x{cols}"),
1350        ));
1351    }
1352    let expected = rows.checked_mul(cols).ok_or_else(|| {
1353        ExprShapeError::new("matrix constructor", "row/column product overflowed")
1354    })?;
1355    let elements = elements.into_iter().map(Expr::from).collect::<Vec<_>>();
1356    if elements.len() != expected {
1357        return Err(ExprShapeError::new(
1358            "matrix constructor",
1359            format!(
1360                "shape {rows}x{cols} requires {expected} elements, got {}",
1361                elements.len()
1362            ),
1363        ));
1364    }
1365    for element in &elements {
1366        element.expect_shape("matrix constructor", ExprShape::Scalar)?;
1367    }
1368    Ok(Expr::new(DagNodeKind::Matrix {
1369        rows,
1370        cols,
1371        elements,
1372    }))
1373}
1374
1375/// Constructs a matrix-matrix multiplication expression.
1376pub fn matmul(lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Expr {
1377    Expr::new(DagNodeKind::MatMul {
1378        lhs: lhs.into(),
1379        rhs: rhs.into(),
1380    })
1381}
1382
1383/// Constructs a matrix-vector multiplication expression.
1384pub fn matvec(matrix: impl Into<Expr>, vector: impl Into<Expr>) -> Expr {
1385    Expr::new(DagNodeKind::MatVec {
1386        matrix: matrix.into(),
1387        vector: vector.into(),
1388    })
1389}
1390
1391/// Constructs a vector dot-product expression.
1392pub fn dot(lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Expr {
1393    Expr::new(DagNodeKind::Dot {
1394        lhs: lhs.into(),
1395        rhs: rhs.into(),
1396    })
1397}
1398
1399/// Constructs an expression that solves a linear system.
1400pub fn solve(matrix: impl Into<Expr>, rhs: impl Into<Expr>) -> Expr {
1401    Expr::new(DagNodeKind::Solve {
1402        matrix: matrix.into(),
1403        rhs: rhs.into(),
1404    })
1405}
1406
1407fn unary(op: UnaryOp, expr: impl Into<Expr>) -> Expr {
1408    Expr::new(DagNodeKind::Unary {
1409        op,
1410        input: expr.into(),
1411    })
1412}
1413
1414fn binary(op: BinaryOp, lhs: impl Into<Expr>, rhs: impl Into<Expr>) -> Expr {
1415    Expr::new(DagNodeKind::Binary {
1416        op,
1417        lhs: lhs.into(),
1418        rhs: rhs.into(),
1419    })
1420}
1421
1422/// Topologically ordered, serializable representation of an [`Expr`] DAG.
1423#[derive(Clone, Debug, Serialize, Deserialize)]
1424pub struct ExprGraph {
1425    root: ExprId,
1426    nodes: Vec<ExprNode>,
1427    metadata: Vec<ExprMetadata>,
1428}
1429
1430impl ExprGraph {
1431    /// Return a copy of this graph with the named scalar parameter fixed.
1432    ///
1433    /// # Errors
1434    ///
1435    /// Returns [`ParamError::UnknownName`] when the graph has no scalar
1436    /// parameter named `name`, or [`ParamError::FixedValueOutOfBounds`] when
1437    /// `value` is outside that parameter's bounds.
1438    pub fn fix_parameter(&self, name: &str, value: f64) -> ParamResult<Self> {
1439        self.map_parameter(name, |parameter| {
1440            if !parameter.bounds_spec().contains(value) {
1441                return Err(ParamError::FixedValueOutOfBounds {
1442                    name: name.to_owned(),
1443                    value,
1444                });
1445            }
1446            Ok(parameter.clone().with_fixed_value(value))
1447        })
1448    }
1449
1450    /// Return a copy of this graph with the named scalar parameter free.
1451    ///
1452    /// # Errors
1453    ///
1454    /// Returns [`ParamError::UnknownName`] when the graph has no scalar
1455    /// parameter named `name`.
1456    pub fn free_parameter(&self, name: &str) -> ParamResult<Self> {
1457        self.map_parameter(name, |parameter| Ok(parameter.clone().with_free()))
1458    }
1459
1460    fn map_parameter(
1461        &self,
1462        name: &str,
1463        mut map: impl FnMut(&Parameter) -> ParamResult<Parameter>,
1464    ) -> ParamResult<Self> {
1465        let mut found = false;
1466        let mut graph = self.clone();
1467        for node in &mut graph.nodes {
1468            if let ExprNode::ScalarParam(parameter) = node
1469                && parameter.name() == name
1470            {
1471                *parameter = map(parameter)?;
1472                found = true;
1473            }
1474        }
1475        if !found {
1476            return Err(ParamError::UnknownName(name.to_owned()));
1477        }
1478        Ok(graph)
1479    }
1480
1481    /// Replaces tagged components that match none of `tags` with zero.
1482    ///
1483    /// Untagged nodes remain active, and a matching tagged node retains its
1484    /// entire subtree.
1485    pub fn project_tags<'a>(&self, tags: impl IntoIterator<Item = &'a str>) -> Self {
1486        let tags: Vec<_> = tags.into_iter().collect();
1487        let mut nodes = Vec::new();
1488        let mut metadata = Vec::new();
1489        let mut remapped = HashMap::new();
1490        let root = self.project_node(
1491            self.root,
1492            &tags,
1493            false,
1494            &mut nodes,
1495            &mut metadata,
1496            &mut remapped,
1497        );
1498        Self {
1499            root,
1500            nodes,
1501            metadata,
1502        }
1503    }
1504
1505    fn project_node(
1506        &self,
1507        old: ExprId,
1508        tags: &[&str],
1509        retain_all: bool,
1510        nodes: &mut Vec<ExprNode>,
1511        metadata: &mut Vec<ExprMetadata>,
1512        remapped: &mut HashMap<(ExprId, bool), ExprId>,
1513    ) -> ExprId {
1514        if let Some(id) = remapped.get(&(old, retain_all)) {
1515            return *id;
1516        }
1517        let old_metadata = &self.metadata[old.index()];
1518        let matches = old_metadata
1519            .tags
1520            .iter()
1521            .any(|tag| tags.contains(&tag.as_ref()));
1522        let node = if !retain_all && !old_metadata.tags.is_empty() && !matches {
1523            ExprNode::RealConst(0.0)
1524        } else {
1525            let retain_children = retain_all || matches;
1526            let map =
1527                |id, nodes: &mut Vec<_>, metadata: &mut Vec<_>, remapped: &mut HashMap<_, _>| {
1528                    self.project_node(id, tags, retain_children, nodes, metadata, remapped)
1529                };
1530            match &self.nodes[old.index()] {
1531                ExprNode::RealConst(value) => ExprNode::RealConst(*value),
1532                ExprNode::ComplexConst(value) => ExprNode::ComplexConst(*value),
1533                ExprNode::ScalarParam(value) => ExprNode::ScalarParam(value.clone()),
1534                ExprNode::EventScalar(value) => ExprNode::EventScalar(Arc::clone(value)),
1535                ExprNode::EventP4Component { name, component } => ExprNode::EventP4Component {
1536                    name: Arc::clone(name),
1537                    component: *component,
1538                },
1539                ExprNode::Unary { op, input } => ExprNode::Unary {
1540                    op: *op,
1541                    input: map(*input, nodes, metadata, remapped),
1542                },
1543                ExprNode::Binary { op, lhs, rhs } => ExprNode::Binary {
1544                    op: *op,
1545                    lhs: map(*lhs, nodes, metadata, remapped),
1546                    rhs: map(*rhs, nodes, metadata, remapped),
1547                },
1548                ExprNode::NaryAdd { terms } => ExprNode::NaryAdd {
1549                    terms: terms
1550                        .iter()
1551                        .map(|id| map(*id, nodes, metadata, remapped))
1552                        .collect(),
1553                },
1554                ExprNode::NaryMul { factors } => ExprNode::NaryMul {
1555                    factors: factors
1556                        .iter()
1557                        .map(|id| map(*id, nodes, metadata, remapped))
1558                        .collect(),
1559                },
1560                ExprNode::Complex { re, im } => ExprNode::Complex {
1561                    re: map(*re, nodes, metadata, remapped),
1562                    im: map(*im, nodes, metadata, remapped),
1563                },
1564                ExprNode::Vector { elements } => ExprNode::Vector {
1565                    elements: elements
1566                        .iter()
1567                        .map(|id| map(*id, nodes, metadata, remapped))
1568                        .collect(),
1569                },
1570                ExprNode::Matrix {
1571                    rows,
1572                    cols,
1573                    elements,
1574                } => ExprNode::Matrix {
1575                    rows: *rows,
1576                    cols: *cols,
1577                    elements: elements
1578                        .iter()
1579                        .map(|id| map(*id, nodes, metadata, remapped))
1580                        .collect(),
1581                },
1582                ExprNode::Component { input, index } => ExprNode::Component {
1583                    input: map(*input, nodes, metadata, remapped),
1584                    index: *index,
1585                },
1586                ExprNode::MatrixElement { input, row, col } => ExprNode::MatrixElement {
1587                    input: map(*input, nodes, metadata, remapped),
1588                    row: *row,
1589                    col: *col,
1590                },
1591                ExprNode::MatMul { lhs, rhs } => ExprNode::MatMul {
1592                    lhs: map(*lhs, nodes, metadata, remapped),
1593                    rhs: map(*rhs, nodes, metadata, remapped),
1594                },
1595                ExprNode::MatVec { matrix, vector } => ExprNode::MatVec {
1596                    matrix: map(*matrix, nodes, metadata, remapped),
1597                    vector: map(*vector, nodes, metadata, remapped),
1598                },
1599                ExprNode::Dot { lhs, rhs } => ExprNode::Dot {
1600                    lhs: map(*lhs, nodes, metadata, remapped),
1601                    rhs: map(*rhs, nodes, metadata, remapped),
1602                },
1603                ExprNode::Solve { matrix, rhs } => ExprNode::Solve {
1604                    matrix: map(*matrix, nodes, metadata, remapped),
1605                    rhs: map(*rhs, nodes, metadata, remapped),
1606                },
1607            }
1608        };
1609        let id = ExprId::from_index(nodes.len());
1610        nodes.push(node);
1611        metadata.push(if matches || retain_all {
1612            old_metadata.clone()
1613        } else {
1614            ExprMetadata::new(old_metadata.source)
1615        });
1616        remapped.insert((old, retain_all), id);
1617        id
1618    }
1619
1620    /// Validates and constructs a graph from its serialized parts.
1621    ///
1622    /// Child nodes must precede their parents, metadata must have one entry per
1623    /// node, and `root` must identify an existing node.
1624    ///
1625    /// # Errors
1626    ///
1627    /// Returns [`ExprGraphError`] when the graph is empty, the metadata and
1628    /// node lengths differ, `root` is invalid, or a child identifier is
1629    /// invalid or does not precede its parent.
1630    pub fn from_parts(
1631        root: ExprId,
1632        nodes: Vec<ExprNode>,
1633        metadata: Vec<ExprMetadata>,
1634    ) -> Result<Self, ExprGraphError> {
1635        if nodes.is_empty() {
1636            return Err(ExprGraphError::Empty);
1637        }
1638        if nodes.len() != metadata.len() {
1639            return Err(ExprGraphError::MetadataLength {
1640                node_len: nodes.len(),
1641                metadata_len: metadata.len(),
1642            });
1643        }
1644        if root.index() >= nodes.len() {
1645            return Err(ExprGraphError::InvalidRoot {
1646                root: root.index(),
1647                node_len: nodes.len(),
1648            });
1649        }
1650        for (index, node) in nodes.iter().enumerate() {
1651            for child in node_child_ids(node) {
1652                if child.index() >= nodes.len() {
1653                    return Err(ExprGraphError::InvalidChild {
1654                        node: index,
1655                        child: child.index(),
1656                    });
1657                }
1658                if child.index() >= index {
1659                    return Err(ExprGraphError::InvalidChildOrder {
1660                        node: index,
1661                        child: child.index(),
1662                    });
1663                }
1664            }
1665        }
1666        Ok(Self {
1667            root,
1668            nodes,
1669            metadata,
1670        })
1671    }
1672
1673    /// Returns the root node identifier.
1674    pub fn root(&self) -> ExprId {
1675        self.root
1676    }
1677
1678    /// Returns the node identified by `id`, if it exists.
1679    pub fn node(&self, id: ExprId) -> Option<&ExprNode> {
1680        self.nodes.get(id.index())
1681    }
1682
1683    /// Returns all nodes in topological order.
1684    pub fn nodes(&self) -> &[ExprNode] {
1685        &self.nodes
1686    }
1687
1688    /// Returns the metadata associated with `id`, if it exists.
1689    pub fn metadata(&self, id: ExprId) -> Option<&ExprMetadata> {
1690        self.metadata.get(id.index())
1691    }
1692
1693    /// Creates a configurable indented-tree display.
1694    pub fn display_tree(&self) -> crate::ExprGraphTreeDisplay<'_> {
1695        crate::ExprGraphTreeDisplay::new(self)
1696    }
1697
1698    /// Creates a configurable Graphviz DOT display.
1699    pub fn display_dot(&self) -> crate::ExprGraphDotDisplay<'_> {
1700        crate::ExprGraphDotDisplay::new(self)
1701    }
1702
1703    fn format_expression(&self, id: ExprId) -> String {
1704        self.format_child(id, ExprPrecedence::Lowest, false)
1705    }
1706
1707    fn format_child(
1708        &self,
1709        id: ExprId,
1710        parent_precedence: ExprPrecedence,
1711        parenthesize_equal: bool,
1712    ) -> String {
1713        let (text, precedence) = self.format_node_expression(id);
1714        if precedence < parent_precedence || (parenthesize_equal && precedence == parent_precedence)
1715        {
1716            format!("({text})")
1717        } else {
1718            text
1719        }
1720    }
1721
1722    fn format_node_expression(&self, id: ExprId) -> (String, ExprPrecedence) {
1723        let Some(node) = self.node(id) else {
1724            return (format!("<missing #{}>", id.index()), ExprPrecedence::Atom);
1725        };
1726
1727        match node {
1728            ExprNode::RealConst(value) => (Self::format_real_number(*value), ExprPrecedence::Atom),
1729            ExprNode::ComplexConst(value) => self.format_complex_const(*value),
1730            ExprNode::ScalarParam(parameter) => (parameter.name().to_owned(), ExprPrecedence::Atom),
1731            ExprNode::EventScalar(name) => (name.to_string(), ExprPrecedence::Atom),
1732            ExprNode::EventP4Component { name, component } => (
1733                format!("{name}.{}", component.label()),
1734                ExprPrecedence::Atom,
1735            ),
1736            ExprNode::Unary { op, input } => self.format_unary_expression(*op, *input),
1737            ExprNode::Binary { op, lhs, rhs } => self.format_binary_expression(*op, *lhs, *rhs),
1738            ExprNode::NaryAdd { terms } => self.format_sum_expression(terms),
1739            ExprNode::NaryMul { factors } => self.format_product_expression(factors),
1740            ExprNode::Complex { re, im } => (
1741                format!(
1742                    "complex({}, {})",
1743                    self.format_expression(*re),
1744                    self.format_expression(*im)
1745                ),
1746                ExprPrecedence::Atom,
1747            ),
1748            ExprNode::Vector { elements } => (
1749                format!(
1750                    "[{}]",
1751                    elements
1752                        .iter()
1753                        .map(|id| self.format_expression(*id))
1754                        .collect::<Vec<_>>()
1755                        .join(", ")
1756                ),
1757                ExprPrecedence::Atom,
1758            ),
1759            ExprNode::Matrix {
1760                rows,
1761                cols,
1762                elements,
1763            } => {
1764                let rows = (0..*rows)
1765                    .map(|row| {
1766                        let start = row * *cols;
1767                        let end = start + *cols;
1768                        format!(
1769                            "[{}]",
1770                            elements[start..end]
1771                                .iter()
1772                                .map(|id| self.format_expression(*id))
1773                                .collect::<Vec<_>>()
1774                                .join(", ")
1775                        )
1776                    })
1777                    .collect::<Vec<_>>()
1778                    .join(", ");
1779                (format!("[{rows}]"), ExprPrecedence::Atom)
1780            }
1781            ExprNode::Component { input, index } => (
1782                format!(
1783                    "{}[{index}]",
1784                    self.format_child(*input, ExprPrecedence::Postfix, false)
1785                ),
1786                ExprPrecedence::Postfix,
1787            ),
1788            ExprNode::MatrixElement { input, row, col } => (
1789                format!(
1790                    "{}[{row}, {col}]",
1791                    self.format_child(*input, ExprPrecedence::Postfix, false)
1792                ),
1793                ExprPrecedence::Postfix,
1794            ),
1795            ExprNode::MatMul { lhs, rhs } => (
1796                format!(
1797                    "matmul({}, {})",
1798                    self.format_expression(*lhs),
1799                    self.format_expression(*rhs)
1800                ),
1801                ExprPrecedence::Atom,
1802            ),
1803            ExprNode::MatVec { matrix, vector } => (
1804                format!(
1805                    "matvec({}, {})",
1806                    self.format_expression(*matrix),
1807                    self.format_expression(*vector)
1808                ),
1809                ExprPrecedence::Atom,
1810            ),
1811            ExprNode::Dot { lhs, rhs } => (
1812                format!(
1813                    "dot({}, {})",
1814                    self.format_expression(*lhs),
1815                    self.format_expression(*rhs)
1816                ),
1817                ExprPrecedence::Atom,
1818            ),
1819            ExprNode::Solve { matrix, rhs } => (
1820                format!(
1821                    "solve({}, {})",
1822                    self.format_expression(*matrix),
1823                    self.format_expression(*rhs)
1824                ),
1825                ExprPrecedence::Atom,
1826            ),
1827        }
1828    }
1829
1830    fn format_unary_expression(&self, op: UnaryOp, input: ExprId) -> (String, ExprPrecedence) {
1831        match op {
1832            UnaryOp::Neg => (
1833                format!("-{}", self.format_child(input, ExprPrecedence::Unary, true)),
1834                ExprPrecedence::Unary,
1835            ),
1836            UnaryOp::Real => (
1837                self.format_call_expression("real", input),
1838                ExprPrecedence::Atom,
1839            ),
1840            UnaryOp::Imag => (
1841                self.format_call_expression("imag", input),
1842                ExprPrecedence::Atom,
1843            ),
1844            UnaryOp::Conj => (
1845                self.format_call_expression("conj", input),
1846                ExprPrecedence::Atom,
1847            ),
1848            UnaryOp::NormSqr => (
1849                format!("|{}|^2", self.format_expression(input)),
1850                ExprPrecedence::Pow,
1851            ),
1852            UnaryOp::Sqrt => (
1853                self.format_call_expression("sqrt", input),
1854                ExprPrecedence::Atom,
1855            ),
1856            UnaryOp::Exp => (
1857                self.format_call_expression("exp", input),
1858                ExprPrecedence::Atom,
1859            ),
1860            UnaryOp::Sin => (
1861                self.format_call_expression("sin", input),
1862                ExprPrecedence::Atom,
1863            ),
1864            UnaryOp::Cos => (
1865                self.format_call_expression("cos", input),
1866                ExprPrecedence::Atom,
1867            ),
1868            UnaryOp::Log => (
1869                self.format_call_expression("log", input),
1870                ExprPrecedence::Atom,
1871            ),
1872            UnaryOp::PowI(power) => {
1873                let exponent = if power < 0 {
1874                    format!("({power})")
1875                } else {
1876                    power.to_string()
1877                };
1878                (
1879                    format!(
1880                        "{}^{exponent}",
1881                        self.format_child(input, ExprPrecedence::Pow, true)
1882                    ),
1883                    ExprPrecedence::Pow,
1884                )
1885            }
1886        }
1887    }
1888
1889    fn format_call_expression(&self, name: &str, input: ExprId) -> String {
1890        format!("{name}({})", self.format_expression(input))
1891    }
1892
1893    fn format_binary_expression(
1894        &self,
1895        op: BinaryOp,
1896        lhs: ExprId,
1897        rhs: ExprId,
1898    ) -> (String, ExprPrecedence) {
1899        match op {
1900            BinaryOp::Add => self.format_sum_expression(&[lhs, rhs]),
1901            BinaryOp::Sub => {
1902                let lhs = self.format_child(lhs, ExprPrecedence::Add, false);
1903                let rhs = self.format_child(rhs, ExprPrecedence::Add, true);
1904                (format!("{lhs} - {rhs}"), ExprPrecedence::Add)
1905            }
1906            BinaryOp::Mul => self.format_product_expression(&[lhs, rhs]),
1907            BinaryOp::Div => {
1908                let lhs = self.format_child(lhs, ExprPrecedence::Mul, false);
1909                let rhs = self.format_child(rhs, ExprPrecedence::Mul, true);
1910                (format!("{lhs} / {rhs}"), ExprPrecedence::Mul)
1911            }
1912            BinaryOp::Atan2 => (
1913                format!(
1914                    "atan2({}, {})",
1915                    self.format_expression(lhs),
1916                    self.format_expression(rhs)
1917                ),
1918                ExprPrecedence::Atom,
1919            ),
1920        }
1921    }
1922
1923    fn format_sum_expression(&self, terms: &[ExprId]) -> (String, ExprPrecedence) {
1924        let mut formatted = String::new();
1925        for term in terms {
1926            let (negative, term) = self.format_signed_term(*term);
1927            if formatted.is_empty() {
1928                if negative {
1929                    formatted.push('-');
1930                }
1931                formatted.push_str(&term);
1932            } else if negative {
1933                formatted.push_str(" - ");
1934                formatted.push_str(&term);
1935            } else {
1936                formatted.push_str(" + ");
1937                formatted.push_str(&term);
1938            }
1939        }
1940        if formatted.is_empty() {
1941            formatted.push('0');
1942        }
1943        (formatted, ExprPrecedence::Add)
1944    }
1945
1946    fn format_signed_term(&self, id: ExprId) -> (bool, String) {
1947        match self.node(id) {
1948            Some(ExprNode::RealConst(value)) if *value < 0.0 => {
1949                (true, Self::format_real_number(-*value))
1950            }
1951            Some(ExprNode::Unary {
1952                op: UnaryOp::Neg,
1953                input,
1954            }) => (true, self.format_child(*input, ExprPrecedence::Add, false)),
1955            Some(ExprNode::NaryMul { factors }) => {
1956                let (negative, product) = self.format_product_parts(factors);
1957                (negative, product)
1958            }
1959            _ => (false, self.format_child(id, ExprPrecedence::Add, false)),
1960        }
1961    }
1962
1963    fn format_product_expression(&self, factors: &[ExprId]) -> (String, ExprPrecedence) {
1964        let (negative, product) = self.format_product_parts(factors);
1965        if negative {
1966            (format!("-{product}"), ExprPrecedence::Unary)
1967        } else {
1968            (product, ExprPrecedence::Mul)
1969        }
1970    }
1971
1972    fn format_product_parts(&self, factors: &[ExprId]) -> (bool, String) {
1973        let mut negative = false;
1974        let mut pieces = Vec::new();
1975
1976        for factor in factors {
1977            match self.node(*factor) {
1978                Some(ExprNode::RealConst(value)) if *value < 0.0 => {
1979                    negative = !negative;
1980                    if *value != -1.0 || factors.len() == 1 {
1981                        pieces.push(Self::format_real_number(-*value));
1982                    }
1983                }
1984                Some(ExprNode::Unary {
1985                    op: UnaryOp::Neg,
1986                    input,
1987                }) => {
1988                    negative = !negative;
1989                    pieces.push(self.format_child(*input, ExprPrecedence::Mul, false));
1990                }
1991                _ => pieces.push(self.format_child(*factor, ExprPrecedence::Mul, false)),
1992            }
1993        }
1994
1995        if pieces.is_empty() {
1996            pieces.push("1".to_owned());
1997        }
1998
1999        (negative, pieces.join(" * "))
2000    }
2001
2002    fn format_complex_const(&self, value: Complex64) -> (String, ExprPrecedence) {
2003        match (value.re, value.im) {
2004            (re, 0.0) => (Self::format_real_number(re), ExprPrecedence::Atom),
2005            (0.0, im) => (Self::format_imaginary_unit(im), ExprPrecedence::Atom),
2006            (re, im) if im < 0.0 => (
2007                format!(
2008                    "{} - {}",
2009                    Self::format_real_number(re),
2010                    Self::format_imaginary_unit(-im)
2011                ),
2012                ExprPrecedence::Add,
2013            ),
2014            (re, im) => (
2015                format!(
2016                    "{} + {}",
2017                    Self::format_real_number(re),
2018                    Self::format_imaginary_unit(im)
2019                ),
2020                ExprPrecedence::Add,
2021            ),
2022        }
2023    }
2024
2025    fn format_real_number(value: f64) -> String {
2026        let Some((value, decimals)) = Self::nearby_simple_decimal(value) else {
2027            return value.to_string();
2028        };
2029
2030        if decimals == 0 {
2031            return value.to_string();
2032        }
2033
2034        let mut formatted = format!("{value:.decimals$}");
2035        while formatted.contains('.') && formatted.ends_with('0') {
2036            formatted.pop();
2037        }
2038        if formatted.ends_with('.') {
2039            formatted.pop();
2040        }
2041        formatted
2042    }
2043
2044    fn nearby_simple_decimal(value: f64) -> Option<(f64, usize)> {
2045        if !value.is_finite() {
2046            return None;
2047        }
2048
2049        for decimals in 0..=12 {
2050            let scale = 10_f64.powi(decimals as i32);
2051            let rounded = (value * scale).round() / scale;
2052            if Self::nearly_equal(value, rounded) {
2053                return Some((rounded, decimals));
2054            }
2055        }
2056
2057        None
2058    }
2059
2060    fn nearly_equal(lhs: f64, rhs: f64) -> bool {
2061        (lhs - rhs).abs() <= f64::EPSILON * lhs.abs().max(rhs.abs()).max(1.0) * 16.0
2062    }
2063
2064    fn format_imaginary_unit(value: f64) -> String {
2065        match value {
2066            1.0 => "i".to_owned(),
2067            -1.0 => "-i".to_owned(),
2068            value => format!("{}i", Self::format_real_number(value)),
2069        }
2070    }
2071
2072    pub(crate) fn node_label(&self, id: ExprId, node: &ExprNode) -> String {
2073        let mut label = match node {
2074            ExprNode::RealConst(value) => {
2075                format!(
2076                    "#{} RealConst({})",
2077                    id.index(),
2078                    Self::format_real_number(*value)
2079                )
2080            }
2081            ExprNode::ComplexConst(value) => {
2082                let (value, _) = self.format_complex_const(*value);
2083                format!("#{} ComplexConst({value})", id.index())
2084            }
2085            ExprNode::ScalarParam(parameter) => {
2086                format!("#{} ScalarParam({})", id.index(), parameter.name())
2087            }
2088            ExprNode::EventScalar(name) => format!("#{} EventScalar({name})", id.index()),
2089            ExprNode::EventP4Component { name, component } => {
2090                format!(
2091                    "#{} EventP4Component({name}.{})",
2092                    id.index(),
2093                    component.label()
2094                )
2095            }
2096            ExprNode::Unary { op, .. } => format!("#{} Unary({op:?})", id.index()),
2097            ExprNode::Binary { op, .. } => format!("#{} Binary({op:?})", id.index()),
2098            ExprNode::NaryAdd { terms } => {
2099                format!("#{} NaryAdd(len={})", id.index(), terms.len())
2100            }
2101            ExprNode::NaryMul { factors } => {
2102                format!("#{} NaryMul(len={})", id.index(), factors.len())
2103            }
2104            ExprNode::Complex { .. } => format!("#{} Complex", id.index()),
2105            ExprNode::Vector { elements } => {
2106                format!("#{} Vector(len={})", id.index(), elements.len())
2107            }
2108            ExprNode::Matrix { rows, cols, .. } => {
2109                format!("#{} Matrix({rows}x{cols})", id.index())
2110            }
2111            ExprNode::Component { index, .. } => {
2112                format!("#{} Component(index={index})", id.index())
2113            }
2114            ExprNode::MatrixElement { row, col, .. } => {
2115                format!("#{} MatrixElement(row={row}, col={col})", id.index())
2116            }
2117            ExprNode::MatMul { .. } => format!("#{} MatMul", id.index()),
2118            ExprNode::MatVec { .. } => format!("#{} MatVec", id.index()),
2119            ExprNode::Dot { .. } => format!("#{} Dot", id.index()),
2120            ExprNode::Solve { .. } => format!("#{} Solve", id.index()),
2121        };
2122
2123        if let Some(metadata) = self.metadata(id) {
2124            if let Some(name) = metadata.name() {
2125                label.push_str(&format!(" name=\"{name}\""));
2126            }
2127            if !metadata.tags().is_empty() {
2128                label.push_str(" tags=[");
2129                for (index, tag) in metadata.tags().iter().enumerate() {
2130                    if index != 0 {
2131                        label.push_str(", ");
2132                    }
2133                    label.push_str(tag);
2134                }
2135                label.push(']');
2136            }
2137        }
2138
2139        label
2140    }
2141}
2142
2143impl fmt::Display for ExprGraph {
2144    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
2145        f.write_str(&self.format_expression(self.root))
2146    }
2147}
2148
2149#[derive(Copy, Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
2150enum ExprPrecedence {
2151    Lowest,
2152    Add,
2153    Mul,
2154    Unary,
2155    Pow,
2156    Postfix,
2157    Atom,
2158}
2159
2160pub(crate) fn node_children(node: &ExprNode) -> Vec<(String, ExprId)> {
2161    match node {
2162        ExprNode::RealConst(_)
2163        | ExprNode::ComplexConst(_)
2164        | ExprNode::ScalarParam(_)
2165        | ExprNode::EventScalar(_)
2166        | ExprNode::EventP4Component { .. } => Vec::new(),
2167        ExprNode::Unary { input, .. } => vec![("input".into(), *input)],
2168        ExprNode::Binary { lhs, rhs, .. } => vec![("lhs".into(), *lhs), ("rhs".into(), *rhs)],
2169        ExprNode::NaryAdd { terms } => terms
2170            .iter()
2171            .enumerate()
2172            .map(|(index, id)| (format!("term[{index}]"), *id))
2173            .collect(),
2174        ExprNode::NaryMul { factors } => factors
2175            .iter()
2176            .enumerate()
2177            .map(|(index, id)| (format!("factor[{index}]"), *id))
2178            .collect(),
2179        ExprNode::Complex { re, im } => vec![("re".into(), *re), ("im".into(), *im)],
2180        ExprNode::Vector { elements } => elements
2181            .iter()
2182            .enumerate()
2183            .map(|(index, id)| (format!("element[{index}]"), *id))
2184            .collect(),
2185        ExprNode::Matrix { cols, elements, .. } => elements
2186            .iter()
2187            .enumerate()
2188            .map(|(index, id)| (format!("element[{},{}]", index / cols, index % cols), *id))
2189            .collect(),
2190        ExprNode::Component { input, .. } | ExprNode::MatrixElement { input, .. } => {
2191            vec![("input".into(), *input)]
2192        }
2193        ExprNode::MatMul { lhs, rhs } | ExprNode::Dot { lhs, rhs } => {
2194            vec![("lhs".into(), *lhs), ("rhs".into(), *rhs)]
2195        }
2196        ExprNode::MatVec { matrix, vector } => {
2197            vec![("matrix".into(), *matrix), ("vector".into(), *vector)]
2198        }
2199        ExprNode::Solve { matrix, rhs } => vec![("matrix".into(), *matrix), ("rhs".into(), *rhs)],
2200    }
2201}
2202
2203fn node_child_ids(node: &ExprNode) -> Vec<ExprId> {
2204    node_children(node)
2205        .into_iter()
2206        .map(|(_, child)| child)
2207        .collect()
2208}
2209
2210#[derive(Default)]
2211struct GraphBuilder {
2212    nodes: Vec<ExprNode>,
2213    metadata: Vec<ExprMetadata>,
2214    ids: HashMap<usize, ExprId>,
2215}
2216
2217impl GraphBuilder {
2218    fn new() -> Self {
2219        Self::default()
2220    }
2221
2222    fn build(mut self, expr: &Expr) -> ExprGraph {
2223        let root = self.visit(expr);
2224        ExprGraph {
2225            root,
2226            nodes: self.nodes,
2227            metadata: self.metadata,
2228        }
2229    }
2230
2231    fn visit(&mut self, expr: &Expr) -> ExprId {
2232        let key = Arc::as_ptr(&expr.node) as usize;
2233        if let Some(id) = self.ids.get(&key) {
2234            return *id;
2235        }
2236        let node = match &expr.node.kind {
2237            DagNodeKind::RealConst(value) => ExprNode::RealConst(*value),
2238            DagNodeKind::ComplexConst(value) => ExprNode::ComplexConst(*value),
2239            DagNodeKind::ScalarParam(parameter) => ExprNode::ScalarParam(parameter.clone()),
2240            DagNodeKind::EventScalar(name) => ExprNode::EventScalar(Arc::clone(name)),
2241            DagNodeKind::EventP4Component { name, component } => ExprNode::EventP4Component {
2242                name: Arc::clone(name),
2243                component: *component,
2244            },
2245            DagNodeKind::Unary { op, input } => {
2246                let input = self.visit(input);
2247                ExprNode::Unary { op: *op, input }
2248            }
2249            DagNodeKind::Binary { op, lhs, rhs } => {
2250                let lhs = self.visit(lhs);
2251                let rhs = self.visit(rhs);
2252                ExprNode::Binary { op: *op, lhs, rhs }
2253            }
2254            DagNodeKind::Complex { re, im } => {
2255                let re = self.visit(re);
2256                let im = self.visit(im);
2257                ExprNode::Complex { re, im }
2258            }
2259            DagNodeKind::Vector { elements } => ExprNode::Vector {
2260                elements: elements.iter().map(|expr| self.visit(expr)).collect(),
2261            },
2262            DagNodeKind::Matrix {
2263                rows,
2264                cols,
2265                elements,
2266            } => ExprNode::Matrix {
2267                rows: *rows,
2268                cols: *cols,
2269                elements: elements.iter().map(|expr| self.visit(expr)).collect(),
2270            },
2271            DagNodeKind::Component { input, index } => {
2272                let input = self.visit(input);
2273                ExprNode::Component {
2274                    input,
2275                    index: *index,
2276                }
2277            }
2278            DagNodeKind::MatrixElement { input, row, col } => {
2279                let input = self.visit(input);
2280                ExprNode::MatrixElement {
2281                    input,
2282                    row: *row,
2283                    col: *col,
2284                }
2285            }
2286            DagNodeKind::MatMul { lhs, rhs } => {
2287                let lhs = self.visit(lhs);
2288                let rhs = self.visit(rhs);
2289                ExprNode::MatMul { lhs, rhs }
2290            }
2291            DagNodeKind::MatVec { matrix, vector } => {
2292                let matrix = self.visit(matrix);
2293                let vector = self.visit(vector);
2294                ExprNode::MatVec { matrix, vector }
2295            }
2296            DagNodeKind::Dot { lhs, rhs } => {
2297                let lhs = self.visit(lhs);
2298                let rhs = self.visit(rhs);
2299                ExprNode::Dot { lhs, rhs }
2300            }
2301            DagNodeKind::Solve { matrix, rhs } => {
2302                let matrix = self.visit(matrix);
2303                let rhs = self.visit(rhs);
2304                ExprNode::Solve { matrix, rhs }
2305            }
2306        };
2307
2308        let id = ExprId::from_index(self.nodes.len());
2309        self.nodes.push(node);
2310        self.metadata.push(expr.node.metadata.clone());
2311        self.ids.insert(key, id);
2312        id
2313    }
2314}
2315
2316fn source_kind(kind: &DagNodeKind) -> ExprSourceKind {
2317    match kind {
2318        DagNodeKind::RealConst(_) | DagNodeKind::ComplexConst(_) => ExprSourceKind::Const,
2319        DagNodeKind::ScalarParam(_) => ExprSourceKind::Param,
2320        DagNodeKind::EventScalar(_) | DagNodeKind::EventP4Component { .. } => ExprSourceKind::Event,
2321        DagNodeKind::Unary { .. } => ExprSourceKind::Unary,
2322        DagNodeKind::Binary { .. } => ExprSourceKind::Binary,
2323        DagNodeKind::Complex { .. } => ExprSourceKind::Complex,
2324        DagNodeKind::Vector { .. } | DagNodeKind::Component { .. } | DagNodeKind::Dot { .. } => {
2325            ExprSourceKind::Vector
2326        }
2327        DagNodeKind::Matrix { .. } | DagNodeKind::MatrixElement { .. } => ExprSourceKind::Matrix,
2328        DagNodeKind::MatMul { .. } | DagNodeKind::MatVec { .. } | DagNodeKind::Solve { .. } => {
2329            ExprSourceKind::LinearAlgebra
2330        }
2331    }
2332}
2333
2334#[cfg(test)]
2335mod tests {
2336    use super::*;
2337    use crate::parameter;
2338
2339    #[test]
2340    fn builds_target_syntax_without_layout_or_context() {
2341        let model = (Complex64::I * parameter!("y", initial : 1.0, bounds : (0.0, 2.0))
2342            + parameter!("x"))
2343        .norm_sqr();
2344
2345        let graph = model.to_graph();
2346        assert!(matches!(
2347            graph.node(graph.root()),
2348            Some(ExprNode::Unary {
2349                op: UnaryOp::NormSqr,
2350                ..
2351            })
2352        ));
2353    }
2354
2355    #[test]
2356    fn parameter_nodes_store_specs_but_do_not_make_layouts() {
2357        let graph = Expr::from(parameter!("x", initial: 1.0)).to_graph();
2358        assert!(matches!(
2359            graph.node(graph.root()),
2360            Some(ExprNode::ScalarParam(spec)) if spec.name() == "x"
2361        ));
2362    }
2363
2364    #[test]
2365    fn complex_constructor_builds_expression_node() {
2366        let graph = complex(parameter!("re"), parameter!("im")).to_graph();
2367
2368        assert!(matches!(
2369            graph.node(graph.root()),
2370            Some(ExprNode::Complex { .. })
2371        ));
2372    }
2373
2374    #[test]
2375    fn polar_complex_lowers_to_expression_graph() {
2376        let graph = polar_complex(parameter!("mag"), parameter!("phase")).to_graph();
2377
2378        assert!(graph.nodes().iter().any(|node| matches!(
2379            node,
2380            ExprNode::Unary {
2381                op: UnaryOp::Exp,
2382                ..
2383            }
2384        )));
2385    }
2386
2387    #[test]
2388    fn metadata_survives_graph_construction() {
2389        let graph = event_scalar("mass")
2390            .named("event mass")
2391            .tagged("data")
2392            .tagged("data")
2393            .to_graph();
2394        let metadata = graph.metadata(graph.root()).unwrap();
2395        assert_eq!(metadata.name(), Some("event mass"));
2396        assert_eq!(metadata.tags(), &[Arc::from("data")]);
2397        assert!(metadata.has_tag("data"));
2398    }
2399
2400    #[test]
2401    fn expressions_round_trip_through_serde_with_metadata() {
2402        let expression = ((parameter!("x", initial: 1.0) + 2.0).named("offset")
2403            * event_scalar("mass").tagged("data"))
2404        .tagged("model");
2405        let encoded = serde_json::to_string(&expression).unwrap();
2406        let decoded: Expr = serde_json::from_str(&encoded).unwrap();
2407
2408        assert_eq!(
2409            serde_json::to_value(expression.to_graph()).unwrap(),
2410            serde_json::to_value(decoded.to_graph()).unwrap()
2411        );
2412    }
2413
2414    #[test]
2415    fn display_formats_graph_as_labeled_tree() {
2416        let graph = ((parameter!("x") + 1.0).named("offset") * event_scalar("mass").tagged("data"))
2417            .to_graph();
2418        let display = graph.display_tree().to_string();
2419
2420        assert!(display.starts_with("ExprGraph(root=#"));
2421        assert!(display.contains("Binary(Mul)"));
2422        assert!(display.contains("┣ lhs:"));
2423        assert!(display.contains("┗ rhs:"));
2424        assert!(display.contains("Binary(Add) name=\"offset\""));
2425        assert!(display.contains("ScalarParam(x)"));
2426        assert!(display.contains("RealConst(1)"));
2427        assert!(display.contains("EventScalar(mass) tags=[data]"));
2428    }
2429
2430    #[test]
2431    fn display_formats_graph_as_expression() {
2432        let costheta = Expr::from(parameter!("costheta"));
2433        let phi = event_scalar("phi");
2434        let p = Expr::from(parameter!("p"));
2435        let phase = Expr::from(7.0) * Complex64::I;
2436        let graph =
2437            (((costheta.powi(2) * phi.sin()) - 5.2).norm_sqr() * p.conj() - phase.exp()).to_graph();
2438
2439        assert_eq!(
2440            graph.to_string(),
2441            "|costheta^2 * sin(phi) - 5.2|^2 * conj(p) - exp(7 * i)"
2442        );
2443    }
2444
2445    #[test]
2446    fn display_parenthesizes_when_precedence_requires_it() {
2447        let a = Expr::from(parameter!("a"));
2448        let b = Expr::from(parameter!("b"));
2449        let c = Expr::from(parameter!("c"));
2450
2451        assert_eq!(
2452            (a.clone() * (b.clone() + c.clone())).to_graph().to_string(),
2453            "a * (b + c)"
2454        );
2455        assert_eq!(
2456            (a.clone() - (b.clone() - c.clone())).to_graph().to_string(),
2457            "a - (b - c)"
2458        );
2459        assert_eq!(((a / b) / c).to_graph().to_string(), "a / b / c");
2460    }
2461
2462    #[test]
2463    fn display_rounds_tiny_float_representation_noise() {
2464        let metadata = ExprMetadata::new(ExprSourceKind::Const);
2465        let graph = ExprGraph::from_parts(
2466            ExprId::from_index(2),
2467            vec![
2468                ExprNode::RealConst(2.9999999999999996),
2469                ExprNode::ComplexConst(Complex64::new(0.30000000000000004, 1.9999999999999998)),
2470                ExprNode::Binary {
2471                    op: BinaryOp::Add,
2472                    lhs: ExprId::from_index(0),
2473                    rhs: ExprId::from_index(1),
2474                },
2475            ],
2476            vec![metadata.clone(), metadata.clone(), metadata],
2477        )
2478        .unwrap();
2479
2480        assert_eq!(graph.to_string(), "3 + 0.3 + 2i");
2481        assert!(graph.display_tree().to_string().contains("RealConst(3)"));
2482        assert!(
2483            graph
2484                .display_tree()
2485                .to_string()
2486                .contains("ComplexConst(0.3 + 2i)")
2487        );
2488    }
2489
2490    #[test]
2491    fn display_formats_p4_components_and_atan2() {
2492        let expr = atan2(
2493            event_p4_component("ks1", P4Component::Py),
2494            event_p4_component("ks1", P4Component::Px),
2495        );
2496
2497        assert_eq!(expr.to_graph().to_string(), "atan2(ks1.py, ks1.px)");
2498    }
2499
2500    #[test]
2501    fn graph_from_parts_validates_structure() {
2502        let metadata = ExprMetadata::new(ExprSourceKind::Const);
2503        let graph = ExprGraph::from_parts(
2504            ExprId::from_index(1),
2505            vec![
2506                ExprNode::RealConst(1.0),
2507                ExprNode::Unary {
2508                    op: UnaryOp::Neg,
2509                    input: ExprId::from_index(0),
2510                },
2511            ],
2512            vec![metadata.clone(), metadata.clone()],
2513        )
2514        .unwrap();
2515        assert!(matches!(
2516            graph.node(graph.root()),
2517            Some(ExprNode::Unary {
2518                op: UnaryOp::Neg,
2519                ..
2520            })
2521        ));
2522
2523        let err = ExprGraph::from_parts(
2524            ExprId::from_index(0),
2525            vec![ExprNode::RealConst(1.0)],
2526            Vec::new(),
2527        )
2528        .unwrap_err();
2529        assert!(matches!(err, ExprGraphError::MetadataLength { .. }));
2530
2531        let err = ExprGraph::from_parts(
2532            ExprId::from_index(0),
2533            vec![ExprNode::Unary {
2534                op: UnaryOp::Neg,
2535                input: ExprId::from_index(0),
2536            }],
2537            vec![metadata],
2538        )
2539        .unwrap_err();
2540        assert!(matches!(err, ExprGraphError::InvalidChildOrder { .. }));
2541    }
2542
2543    #[test]
2544    fn graph_preserves_unsimplified_expression_shape() {
2545        let graph = (parameter!("x") + 0.0).to_graph();
2546        assert!(matches!(
2547            graph.node(graph.root()),
2548            Some(ExprNode::Binary {
2549                op: BinaryOp::Add,
2550                ..
2551            })
2552        ));
2553    }
2554
2555    #[test]
2556    fn graph_preserves_written_operand_order_for_commutative_ops() {
2557        let left_param = (parameter!("x") + 1.0).to_graph();
2558        assert!(matches!(
2559            left_param.node(left_param.root()),
2560            Some(ExprNode::Binary {
2561                op: BinaryOp::Add,
2562                lhs,
2563                rhs
2564            }) if matches!(left_param.node(*lhs), Some(ExprNode::ScalarParam(parameter)) if parameter.name() == "x")
2565                && matches!(left_param.node(*rhs), Some(ExprNode::RealConst(1.0)))
2566        ));
2567
2568        let right_param = (1.0 + parameter!("x")).to_graph();
2569        assert!(matches!(
2570            right_param.node(right_param.root()),
2571            Some(ExprNode::Binary {
2572                op: BinaryOp::Add,
2573                lhs,
2574                rhs
2575            }) if matches!(right_param.node(*lhs), Some(ExprNode::RealConst(1.0)))
2576                && matches!(right_param.node(*rhs), Some(ExprNode::ScalarParam(parameter)) if parameter.name() == "x")
2577        ));
2578    }
2579
2580    #[test]
2581    fn represents_kmatrix_style_solve_graph() {
2582        let beta = vector([
2583            complex(parameter!("b0_re"), parameter!("b0_im")),
2584            complex(parameter!("b1_re"), parameter!("b1_im")),
2585        ]);
2586        let a = matrix([
2587            [Complex64::new(1.0, 0.0), Complex64::new(0.0, 1.0)],
2588            [Complex64::new(0.0, -1.0), Complex64::new(1.0, 0.0)],
2589        ]);
2590        let graph = solve(a, beta).component(0).to_graph();
2591
2592        assert!(
2593            graph
2594                .nodes()
2595                .iter()
2596                .any(|node| matches!(node, ExprNode::Solve { .. }))
2597        );
2598    }
2599
2600    #[test]
2601    fn graph_builder_preserves_shared_dag_nodes() {
2602        let shared = event_scalar("x").sin();
2603        let expression = vector((0..1_000).map(|_| shared.clone()));
2604        let graph = expression.to_graph();
2605
2606        assert_eq!(graph.nodes().len(), 3);
2607        let ExprNode::Vector { elements } = graph.node(graph.root()).unwrap() else {
2608            panic!("root should be a vector");
2609        };
2610        assert!(elements.windows(2).all(|pair| pair[0] == pair[1]));
2611    }
2612
2613    #[test]
2614    fn dynamic_matrices_and_shapes_are_checked_eagerly() {
2615        let dynamic = matrix_from_flat(2, 2, [1.0, 2.0, 3.0, 4.0]).unwrap();
2616        assert_eq!(
2617            dynamic.shape().unwrap(),
2618            ExprShape::Matrix { rows: 2, cols: 2 }
2619        );
2620        assert!(matrix_from_flat(2, 2, [1.0, 2.0, 3.0]).is_err());
2621        assert!(matmul(dynamic, matrix([[1.0, 2.0, 3.0]])).shape().is_err());
2622    }
2623
2624    #[test]
2625    fn assignment_operators_build_binary_expression_nodes() {
2626        let mut expr = Expr::from(parameter!("x"));
2627        expr += parameter!("y");
2628        expr -= 1.0;
2629        expr *= Complex64::I;
2630        expr /= Expr::from(parameter!("z"));
2631
2632        let graph = expr.to_graph();
2633        assert!(matches!(
2634            graph.node(graph.root()),
2635            Some(ExprNode::Binary {
2636                op: BinaryOp::Div,
2637                ..
2638            })
2639        ));
2640        assert_eq!(
2641            graph
2642                .nodes()
2643                .iter()
2644                .filter(|node| matches!(node, ExprNode::Binary { .. }))
2645                .count(),
2646            4
2647        );
2648    }
2649
2650    #[test]
2651    fn assignment_operators_accept_borrowed_rhs_values() {
2652        let y = parameter!("y");
2653        let one = 1.0;
2654        let i = Complex64::I;
2655        let z = Expr::from(parameter!("z"));
2656
2657        let mut expr = Expr::from(parameter!("x"));
2658        expr += &y;
2659        expr -= &one;
2660        expr *= &i;
2661        expr /= &z;
2662
2663        let graph = expr.to_graph();
2664        assert!(matches!(
2665            graph.node(graph.root()),
2666            Some(ExprNode::Binary {
2667                op: BinaryOp::Div,
2668                ..
2669            })
2670        ));
2671    }
2672}
laddu_expr/expression.rs

laddu_expr/
expression.rs
