UnaryExpression struct in Expressions (#1352)

This PR is part of issue #1345.
In particular, it adds the struct UnaryOperation to Expressions to
homogenise the structure before including source references.

asm-to-pil/src/vm_to_constrained.rs

@@ -14,7 +14,7 @@ use powdr_ast::{
         visitor::ExpressionVisitable,
         ArrayExpression, BinaryOperation, BinaryOperator, Expression, FunctionCall,
         FunctionDefinition, FunctionKind, LambdaExpression, MatchArm, Number, Pattern,
-        PilStatement, PolynomialName, SelectedExpressions, UnaryOperator,
+        PilStatement, PolynomialName, SelectedExpressions, UnaryOperation, UnaryOperator,
     },
     SourceRef,
 };
@@ -535,8 +535,11 @@ impl<T: FieldElement> VMConverter<T> {
                             .filter(|(_, reg)| reg.ty == RegisterTy::Assignment)
                             .map(|(name, _)| rhs_assignment_registers.insert(name.clone()));
                     }
-                    Expression::UnaryOperation(UnaryOperator::Next, e) => {
-                        if let Expression::Reference(poly) = e.as_ref() {
+                    Expression::UnaryOperation(UnaryOperation {
+                        op: UnaryOperator::Next,
+                        expr,
+                    }) => {
+                        if let Expression::Reference(poly) = expr.as_ref() {
                             poly.try_to_identifier()
                                 .and_then(|name| self.registers.get(name).map(|reg| (name, reg)))
                                 .filter(|(_, reg)| {
@@ -683,7 +686,7 @@ impl<T: FieldElement> VMConverter<T> {
                                 instruction_literal_arg.push(InstructionLiteralArg::Number(
                                     T::checked_from(value).unwrap(),
                                 ));
-                            } else if let Expression::UnaryOperation(UnaryOperator::Minus, expr) = a
+                            } else if let Expression::UnaryOperation(UnaryOperation { op: UnaryOperator::Minus, expr }) = a
                             {
                                 if let Expression::Number(Number {value, ..}) = *expr {
                                     instruction_literal_arg.push(InstructionLiteralArg::Number(
@@ -813,7 +816,7 @@ impl<T: FieldElement> VMConverter<T> {
                     panic!("Invalid operation in expression {left} {op} {right}")
                 }
             },
-            Expression::UnaryOperation(op, expr) => {
+            Expression::UnaryOperation(UnaryOperation { op, expr }) => {
                 assert!(op == UnaryOperator::Minus);
                 self.negate_assignment_value(self.process_assignment_value(*expr))
             }
@@ -1222,13 +1225,19 @@ fn extract_update(expr: Expression) -> (Option<String>, Expression) {
     };
     // TODO check that there are no other "next" references in the expression
     match *left {
-        Expression::UnaryOperation(UnaryOperator::Next, column) => match *column {
+        Expression::UnaryOperation(UnaryOperation {
+            op: UnaryOperator::Next,
+            expr,
+        }) => match *expr {
             Expression::Reference(column) => {
                 (Some(column.try_to_identifier().unwrap().clone()), *right)
             }
             _ => (
                 None,
-                Expression::UnaryOperation(UnaryOperator::Next, column) - *right,
+                Expression::UnaryOperation(UnaryOperation {
+                    op: UnaryOperator::Next,
+                    expr,
+                }) - *right,
             ),
         },
         _ => (None, *left - *right),

ast/src/parsed/build.rs

@@ -4,7 +4,8 @@ use crate::parsed::Expression;
 
 use super::{
     asm::{parse_absolute_path, Part, SymbolPath},
-    BinaryOperation, BinaryOperator, IndexAccess, NamespacedPolynomialReference, UnaryOperator,
+    BinaryOperation, BinaryOperator, IndexAccess, NamespacedPolynomialReference, UnaryOperation,
+    UnaryOperator,
 };
 
 pub fn absolute_reference(name: &str) -> Expression {
@@ -25,7 +26,11 @@ pub fn namespaced_reference<S: Into<String>>(namespace: String, name: S) -> Expr
 }
 
 pub fn next_reference<S: Into<String>>(name: S) -> Expression {
-    Expression::UnaryOperation(UnaryOperator::Next, Box::new(direct_reference(name)))
+    UnaryOperation {
+        op: UnaryOperator::Next,
+        expr: Box::new(direct_reference(name)),
+    }
+    .into()
 }
 
 /// Returns an index access operation to expr if the index is Some, otherwise returns expr itself.

ast/src/parsed/display.rs

@@ -595,30 +595,6 @@ fn format_list<L: IntoIterator<Item = I>, I: Display>(list: L) -> String {
     format!("{}", list.into_iter().format(", "))
 }
 
-impl<E: Display> Expression<E> {
-    pub fn format_unary_operation(
-        &self,
-        op: &UnaryOperator,
-        exp: &Expression<E>,
-        f: &mut Formatter<'_>,
-    ) -> Result {
-        let exp_string = match (self.precedence(), exp.precedence()) {
-            (Some(precedence), Some(inner_precedence)) if precedence < inner_precedence => {
-                format!("({exp})")
-            }
-            _ => {
-                format!("{exp}")
-            }
-        };
-
-        if op.is_prefix() {
-            write!(f, "{op}{exp_string}")
-        } else {
-            write!(f, "{exp_string}{op}")
-        }
-    }
-}
-
 impl<Ref: Display> Display for Expression<Ref> {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         match self {
@@ -632,7 +608,9 @@ impl<Ref: Display> Display for Expression<Ref> {
             Expression::BinaryOperation(binaryop) => {
                 write!(f, "{binaryop}")
             }
-            Expression::UnaryOperation(op, exp) => self.format_unary_operation(op, exp, f),
+            Expression::UnaryOperation(unaryop) => {
+                write!(f, "{unaryop}")
+            }
             Expression::IndexAccess(index_access) => write!(f, "{index_access}"),
             Expression::FunctionCall(fun_call) => write!(f, "{fun_call}"),
             Expression::FreeInput(input) => write!(f, "${{ {input} }}"),
@@ -785,6 +763,28 @@ impl Display for BinaryOperator {
     }
 }
 
+impl<E> Display for UnaryOperation<E>
+where
+    E: Display + Precedence,
+{
+    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
+        let exp_string = match (self.op.precedence(), self.expr.precedence()) {
+            (Some(precedence), Some(inner_precedence)) if precedence < inner_precedence => {
+                format!("({})", self.expr)
+            }
+            _ => {
+                format!("{}", self.expr)
+            }
+        };
+
+        if self.op.is_prefix() {
+            write!(f, "{}{exp_string}", self.op)
+        } else {
+            write!(f, "{exp_string}{}", self.op)
+        }
+    }
+}
+
 impl Display for UnaryOperator {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         write!(

ast/src/parsed/mod.rs

@@ -342,8 +342,8 @@ pub enum Expression<Ref = NamespacedPolynomialReference> {
     Tuple(Vec<Self>),
     LambdaExpression(LambdaExpression<Self>),
     ArrayLiteral(ArrayLiteral<Self>),
+    UnaryOperation(UnaryOperation<Self>),
     BinaryOperation(BinaryOperation<Self>),
-    UnaryOperation(UnaryOperator, Box<Self>),
     IndexAccess(IndexAccess<Self>),
     FunctionCall(FunctionCall<Self>),
     FreeInput(Box<Self>),
@@ -352,6 +352,18 @@ pub enum Expression<Ref = NamespacedPolynomialReference> {
     BlockExpression(Vec<StatementInsideBlock<Self>>, Box<Self>),
 }
 
+#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
+pub struct UnaryOperation<E = Expression<NamespacedPolynomialReference>> {
+    pub op: UnaryOperator,
+    pub expr: Box<E>,
+}
+
+impl<Ref> From<UnaryOperation<Expression<Ref>>> for Expression<Ref> {
+    fn from(operation: UnaryOperation<Expression<Ref>>) -> Self {
+        Expression::UnaryOperation(operation)
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Serialize, Deserialize, JsonSchema)]
 pub struct BinaryOperation<E = Expression<NamespacedPolynomialReference>> {
     pub left: Box<E>,
@@ -479,7 +491,7 @@ impl<R> Expression<R> {
             Expression::BinaryOperation(BinaryOperation { left, right, .. }) => {
                 [left.as_ref(), right.as_ref()].into_iter()
             }
-            Expression::UnaryOperation(_, e) => once(e.as_ref()),
+            Expression::UnaryOperation(UnaryOperation { expr, .. }) => once(expr.as_ref()),
             Expression::IndexAccess(IndexAccess { array, index }) => {
                 [array.as_ref(), index.as_ref()].into_iter()
             }
@@ -520,7 +532,7 @@ impl<R> Expression<R> {
             Expression::BinaryOperation(BinaryOperation { left, right, .. }) => {
                 [left.as_mut(), right.as_mut()].into_iter()
             }
-            Expression::UnaryOperation(_, e) => once(e.as_mut()),
+            Expression::UnaryOperation(UnaryOperation { expr, .. }) => once(expr.as_mut()),
             Expression::IndexAccess(IndexAccess { array, index }) => {
                 [array.as_mut(), index.as_mut()].into_iter()
             }
@@ -729,7 +741,7 @@ impl Precedence for BinaryOperator {
 impl<E> Precedence for Expression<E> {
     fn precedence(&self) -> Option<ExpressionPrecedence> {
         match self {
-            Expression::UnaryOperation(op, _) => op.precedence(),
+            Expression::UnaryOperation(operation) => operation.op.precedence(),
             Expression::BinaryOperation(operation) => operation.op.precedence(),
             _ => None,
         }

importer/src/path_canonicalizer.rs

@@ -16,7 +16,7 @@ use powdr_ast::parsed::{
     visitor::{Children, ExpressionVisitable},
     ArrayLiteral, BinaryOperation, EnumDeclaration, EnumVariant, Expression, FunctionCall,
     IndexAccess, LambdaExpression, LetStatementInsideBlock, MatchArm, Pattern, PilStatement,
-    StatementInsideBlock, TypedExpression,
+    StatementInsideBlock, TypedExpression, UnaryOperation,
 };
 
 /// Changes all symbol references (symbol paths) from relative paths
@@ -159,7 +159,7 @@ fn free_inputs_in_expression<'a>(
         Expression::BinaryOperation(BinaryOperation { left, right, .. }) => {
             Box::new(free_inputs_in_expression(left).chain(free_inputs_in_expression(right)))
         }
-        Expression::UnaryOperation(_, expr) => free_inputs_in_expression(expr),
+        Expression::UnaryOperation(UnaryOperation { expr, .. }) => free_inputs_in_expression(expr),
         Expression::FunctionCall(FunctionCall {
             function,
             arguments,
@@ -191,7 +191,9 @@ fn free_inputs_in_expression_mut<'a>(
         Expression::BinaryOperation(BinaryOperation { left, right, .. }) => Box::new(
             free_inputs_in_expression_mut(left).chain(free_inputs_in_expression_mut(right)),
         ),
-        Expression::UnaryOperation(_, expr) => free_inputs_in_expression_mut(expr),
+        Expression::UnaryOperation(UnaryOperation { expr, .. }) => {
+            free_inputs_in_expression_mut(expr)
+        }
         Expression::FunctionCall(FunctionCall {
             function,
             arguments,
@@ -667,8 +669,8 @@ fn check_expression(
             check_expression(location, a.as_ref(), state, local_variables)?;
             check_expression(location, b.as_ref(), state, local_variables)
         }
-        Expression::UnaryOperation(_, e) | Expression::FreeInput(e) => {
-            check_expression(location, e, state, local_variables)
+        Expression::UnaryOperation(UnaryOperation { expr, .. }) | Expression::FreeInput(expr) => {
+            check_expression(location, expr, state, local_variables)
         }
         Expression::FunctionCall(FunctionCall {
             function,

parser/src/powdr.lalrpop

@@ -554,7 +554,7 @@ PowOp: BinaryOperator = {
 }
 
 Unary: Box<Expression> = {
-    PrefixUnaryOp PostfixUnary => Box::new(Expression::UnaryOperation(<>)),
+    <op:PrefixUnaryOp> <expr:PostfixUnary> => Box::new(UnaryOperation{op, expr}.into()),
     PostfixUnary,
 }
 
@@ -564,7 +564,7 @@ PrefixUnaryOp: UnaryOperator = {
 }
 
 PostfixUnary: Box<Expression> = {
-    <t:Term> <o:PostfixUnaryOp> => Box::new(Expression::UnaryOperation(o, t)),
+    <t:Term> <op:PostfixUnaryOp> => Box::new(UnaryOperation{op, expr: t}.into()),
     Term,
 }
 

pil-analyzer/src/evaluator.rs

@@ -18,7 +18,7 @@ use powdr_ast::{
         types::{Type, TypeScheme},
         ArrayLiteral, BinaryOperation, BinaryOperator, FunctionCall, IfExpression, IndexAccess,
         LambdaExpression, LetStatementInsideBlock, MatchArm, Number, Pattern, StatementInsideBlock,
-        UnaryOperator,
+        UnaryOperation, UnaryOperator,
     },
     SourceRef,
 };
@@ -697,7 +697,7 @@ impl<'a, 'b, T: FieldElement, S: SymbolLookup<'a, T>> Evaluator<'a, 'b, T, S> {
                 self.op_stack.push(Operation::Expand(right));
                 self.expand(left)?;
             }
-            Expression::UnaryOperation(_, inner) => {
+            Expression::UnaryOperation(UnaryOperation { expr: inner, .. }) => {
                 self.op_stack.push(Operation::Combine(expr));
                 self.expand(inner)?;
             }
@@ -800,7 +800,7 @@ impl<'a, 'b, T: FieldElement, S: SymbolLookup<'a, T>> Evaluator<'a, 'b, T, S> {
                 let left = self.value_stack.pop().unwrap();
                 evaluate_binary_operation(&left, *op, &right)?
             }
-            Expression::UnaryOperation(op, _) => {
+            Expression::UnaryOperation(UnaryOperation { op, .. }) => {
                 let inner = self.value_stack.pop().unwrap();
                 match (op, inner.as_ref()) {
                     (UnaryOperator::Minus, Value::FieldElement(e)) => {

pil-analyzer/src/expression_processor.rs

@@ -9,7 +9,7 @@ use powdr_ast::{
     parsed::{
         self, asm::SymbolPath, ArrayExpression, ArrayLiteral, BinaryOperation, IfExpression,
         LambdaExpression, LetStatementInsideBlock, MatchArm, NamespacedPolynomialReference, Number,
-        Pattern, SelectedExpressions, StatementInsideBlock, SymbolCategory,
+        Pattern, SelectedExpressions, StatementInsideBlock, SymbolCategory, UnaryOperation,
     },
 };
 use powdr_number::DegreeType;
@@ -98,6 +98,11 @@ impl<'a, D: AnalysisDriver> ExpressionProcessor<'a, D> {
             PExpression::LambdaExpression(lambda_expression) => {
                 Expression::LambdaExpression(self.process_lambda_expression(lambda_expression))
             }
+            PExpression::UnaryOperation(UnaryOperation { op, expr: value }) => UnaryOperation {
+                op,
+                expr: Box::new(self.process_expression(*value)),
+            }
+            .into(),
             PExpression::BinaryOperation(BinaryOperation { left, op, right }) => {
                 (BinaryOperation {
                     left: Box::new(self.process_expression(*left)),
@@ -106,9 +111,6 @@ impl<'a, D: AnalysisDriver> ExpressionProcessor<'a, D> {
                 })
                 .into()
             }
-            PExpression::UnaryOperation(op, value) => {
-                Expression::UnaryOperation(op, Box::new(self.process_expression(*value)))
-            }
             PExpression::IndexAccess(index_access) => {
                 Expression::IndexAccess(parsed::IndexAccess {
                     array: Box::new(self.process_expression(*index_access.array)),

pil-analyzer/src/type_inference.rs

@@ -8,7 +8,7 @@ use powdr_ast::{
         types::{ArrayType, FunctionType, TupleType, Type, TypeBounds, TypeScheme},
         visitor::ExpressionVisitable,
         ArrayLiteral, BinaryOperation, FunctionCall, IndexAccess, LambdaExpression,
-        LetStatementInsideBlock, MatchArm, Number, Pattern, StatementInsideBlock,
+        LetStatementInsideBlock, MatchArm, Number, Pattern, StatementInsideBlock, UnaryOperation,
     },
 };
 
@@ -553,7 +553,7 @@ impl<'a> TypeChecker<'a> {
                     || format!("applying operator {op}"),
                 )?
             }
-            Expression::UnaryOperation(op, inner) => {
+            Expression::UnaryOperation(UnaryOperation { op, expr: inner }) => {
                 // TODO at some point, also store the generic args for operators
                 let fun_type = self.instantiate_scheme(unary_operator_scheme(*op)).0;
                 self.infer_type_of_function_call(

riscv-executor/src/lib.rs

@@ -22,7 +22,9 @@ use powdr_ast::{
     asm_analysis::{
         AnalysisASMFile, CallableSymbol, FunctionStatement, Item, LabelStatement, Machine,
     },
-    parsed::{asm::DebugDirective, BinaryOperation, Expression, FunctionCall, Number},
+    parsed::{
+        asm::DebugDirective, BinaryOperation, Expression, FunctionCall, Number, UnaryOperation,
+    },
 };
 use powdr_number::{FieldElement, LargeInt};
 use powdr_riscv_syscalls::SYSCALL_REGISTERS;
@@ -901,7 +903,7 @@ impl<'a, 'b, F: FieldElement> Executor<'a, 'b, F> {
 
                 vec![result]
             }
-            Expression::UnaryOperation(op, arg) => {
+            Expression::UnaryOperation(UnaryOperation { op, expr: arg }) => {
                 let arg = self.eval_expression(arg)[0].bin();
                 let result = match op {
                     powdr_ast::parsed::UnaryOperator::Minus => -arg,