Skip to main content

powdr_expression/
display.rs

1use std::fmt::{self, Display, Formatter};
2
3use crate::{
4    AlgebraicBinaryOperation, AlgebraicBinaryOperator, AlgebraicExpression,
5    AlgebraicUnaryOperation, AlgebraicUnaryOperator,
6};
7
8type ExpressionPrecedence = u64;
9trait Precedence {
10    fn precedence(&self) -> Option<ExpressionPrecedence>;
11}
12
13impl Precedence for AlgebraicUnaryOperator {
14    fn precedence(&self) -> Option<ExpressionPrecedence> {
15        Some(match self {
16            AlgebraicUnaryOperator::Minus => 1,
17        })
18    }
19}
20
21impl Precedence for AlgebraicBinaryOperator {
22    fn precedence(&self) -> Option<ExpressionPrecedence> {
23        Some(match self {
24            Self::Mul => 3,
25            Self::Add | Self::Sub => 4,
26        })
27    }
28}
29
30impl<T, R> Precedence for AlgebraicExpression<T, R> {
31    fn precedence(&self) -> Option<ExpressionPrecedence> {
32        match self {
33            AlgebraicExpression::UnaryOperation(operation) => operation.op.precedence(),
34            AlgebraicExpression::BinaryOperation(operation) => operation.op.precedence(),
35            AlgebraicExpression::Number(..) | AlgebraicExpression::Reference(..) => None,
36        }
37    }
38}
39
40impl<T: Display, R: Display> Display for AlgebraicBinaryOperation<T, R> {
41    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
42        let op_precedence = self.op.precedence().unwrap();
43        let use_left_parentheses = match self.left.precedence() {
44            Some(left_precedence) => left_precedence > op_precedence,
45            None => false,
46        };
47
48        let use_right_parentheses = match self.right.precedence() {
49            Some(right_precedence) => right_precedence >= op_precedence,
50            None => false,
51        };
52
53        let left_string = if use_left_parentheses {
54            format!("({})", self.left)
55        } else {
56            format!("{}", self.left)
57        };
58        let right_string = if use_right_parentheses {
59            format!("({})", self.right)
60        } else {
61            format!("{}", self.right)
62        };
63
64        write!(f, "{left_string} {} {right_string}", self.op)
65    }
66}
67
68impl<T: Display, R: Display> Display for AlgebraicUnaryOperation<T, R> {
69    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
70        let exp_string = match (self.op.precedence(), self.expr.precedence()) {
71            (Some(precedence), Some(inner_precedence)) if precedence < inner_precedence => {
72                format!("({})", self.expr)
73            }
74            _ => {
75                format!("{}", self.expr)
76            }
77        };
78
79        write!(f, "{}{exp_string}", self.op)
80    }
81}
82
83impl Display for AlgebraicUnaryOperator {
84    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
85        match self {
86            AlgebraicUnaryOperator::Minus => write!(f, "-"),
87        }
88    }
89}
90
91impl Display for AlgebraicBinaryOperator {
92    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
93        match self {
94            AlgebraicBinaryOperator::Add => write!(f, "+"),
95            AlgebraicBinaryOperator::Sub => write!(f, "-"),
96            AlgebraicBinaryOperator::Mul => write!(f, "*"),
97        }
98    }
99}
100
101#[cfg(test)]
102mod test {
103    use powdr_number::GoldilocksField;
104    use pretty_assertions::assert_eq;
105    use test_log::test;
106
107    use super::AlgebraicExpression;
108
109    fn test_display(expr: AlgebraicExpression<GoldilocksField, &str>, expected: &str) {
110        assert_eq!(expr.to_string(), expected);
111    }
112
113    #[test]
114    fn binary_op() {
115        let x = AlgebraicExpression::Reference("x");
116        let y = AlgebraicExpression::Reference("y");
117        let z = AlgebraicExpression::Reference("z");
118        // Don't add extra
119        test_display(x.clone() + y.clone() + z.clone(), "x + y + z");
120        test_display(x.clone() * y.clone() * z.clone(), "x * y * z");
121        // Remove unneeded
122        test_display(-x.clone() + y.clone() * z.clone(), "-x + y * z");
123        test_display((x.clone() * y.clone()) * z.clone(), "x * y * z");
124        test_display(x.clone() - (y.clone() + z.clone()), "x - (y + z)");
125        test_display((x.clone() * y.clone()) + z.clone(), "x * y + z");
126        // Observe associativity
127        test_display(x.clone() * (y.clone() * z.clone()), "x * (y * z)");
128        test_display(x.clone() + (y.clone() + z.clone()), "x + (y + z)");
129        // Don't remove needed
130        test_display((x.clone() + y.clone()) * z.clone(), "(x + y) * z");
131        test_display(-(x.clone() + y.clone()), "-(x + y)");
132    }
133}