SourceRef in expressions (powdr-labs#1357)

This PR solves issue powdr-labs#1293. Includes PRs that modify some Expression in order to match fields. It is recommended to review/merge those first: - powdr-labs#1351 - powdr-labs#1352 - powdr-labs#1353 - powdr-labs#1354 - powdr-labs#1355 All comments are welcome :) --------- Co-authored-by: chriseth <[email protected]>
Maddiaa0 · May 26, 2024 · 0d9669e · 0d9669e
1 parent f04fb8c
commit 0d9669e
Show file tree

Hide file tree

Showing 25 changed files with 775 additions and 498 deletions.
diff --git a/analysis/src/vm/inference.rs b/analysis/src/vm/inference.rs
@@ -39,9 +39,9 @@ fn infer_machine(mut machine: Machine) -> Result<Machine, Vec<String>> {
             if let FunctionStatement::Assignment(a) = s {
                 // Map function calls to the list of assignment registers and all other expressions to a list of None.
                 let expr_regs = match &*a.rhs {
-                    Expression::FunctionCall(c) => {
+                    Expression::FunctionCall(_, c) => {
                         let instr_name =
-                            if let Expression::Reference(reference) = c.function.as_ref() {
+                            if let Expression::Reference(_, reference) = c.function.as_ref() {
                                 reference.try_to_identifier().unwrap()
                             } else {
                                 panic!("Only instructions allowed.");

diff --git a/asm-to-pil/src/romgen.rs b/asm-to-pil/src/romgen.rs
@@ -32,7 +32,7 @@ fn substitute_name_in_statement_expressions(
     substitution: &HashMap<String, String>,
 ) {
     fn substitute(e: &mut Expression, substitution: &HashMap<String, String>) {
-        if let Expression::Reference(r) = e {
+        if let Expression::Reference(_, r) = e {
             if let Some(n) = r.try_to_identifier() {
                 if let Some(v) = substitution.get(n).cloned() {
                     *r = NamespacedPolynomialReference::from_identifier(v);

diff --git a/asm-to-pil/src/vm_to_constrained.rs b/asm-to-pil/src/vm_to_constrained.rs
@@ -243,7 +243,7 @@ impl<T: FieldElement> VMConverter<T> {
                     .collect();
 
                 match *rhs {
-                    Expression::FunctionCall(c) => {
+                    Expression::FunctionCall(_, c) => {
                         self.handle_functional_instruction(lhs_with_reg, *c.function, c.arguments)
                     }
                     _ => self.handle_non_functional_assignment(source, lhs_with_reg, *rhs),
@@ -432,7 +432,7 @@ impl<T: FieldElement> VMConverter<T> {
             .collect::<HashMap<_, _>>();
         body.iter_mut().for_each(|s| {
             s.post_visit_expressions_mut(&mut |e| {
-                if let Expression::Reference(r) = e {
+                if let Expression::Reference(_, r) = e {
                     if let Some(name) = r.try_to_identifier() {
                         if let Some(sub) = substitutions.get(name) {
                             *r.path.try_last_part_mut().unwrap() = sub.to_string();
@@ -529,17 +529,20 @@ impl<T: FieldElement> VMConverter<T> {
             // collect assignment registers and next references to write registers used on rhs
             for expr in rhs.inputs_and_outputs() {
                 expr.pre_visit_expressions(&mut |e| match e {
-                    Expression::Reference(poly) => {
+                    Expression::Reference(_, poly) => {
                         poly.try_to_identifier()
                             .and_then(|name| self.registers.get(name).map(|reg| (name, reg)))
                             .filter(|(_, reg)| reg.ty == RegisterTy::Assignment)
                             .map(|(name, _)| rhs_assignment_registers.insert(name.clone()));
                     }
-                    Expression::UnaryOperation(UnaryOperation {
-                        op: UnaryOperator::Next,
-                        expr,
-                    }) => {
-                        if let Expression::Reference(poly) = expr.as_ref() {
+                    Expression::UnaryOperation(
+                        _,
+                        UnaryOperation {
+                            op: UnaryOperator::Next,
+                            expr: e,
+                        },
+                    ) => {
+                        if let Expression::Reference(_, poly) = e.as_ref() {
                             poly.try_to_identifier()
                                 .and_then(|name| self.registers.get(name).map(|reg| (name, reg)))
                                 .filter(|(_, reg)| {
@@ -615,7 +618,7 @@ impl<T: FieldElement> VMConverter<T> {
         function: Expression,
         mut args: Vec<Expression>,
     ) -> CodeLine<T> {
-        let Expression::Reference(reference) = function else {
+        let Expression::Reference(_, reference) = function else {
             panic!("Expected instruction name");
         };
         let instr_name = reference.try_to_identifier().unwrap();
@@ -660,7 +663,7 @@ impl<T: FieldElement> VMConverter<T> {
                             value.insert(reg.clone(), self.process_assignment_value(a));
                         }
                         Input::Literal(_, LiteralKind::Label) => {
-                            if let Expression::Reference(r) = a {
+                            if let Expression::Reference(_, r) = a {
                                 instruction_literal_arg.push(InstructionLiteralArg::LabelRef(
                                     r.try_to_identifier().unwrap().clone(),
                                 ));
@@ -670,25 +673,25 @@ impl<T: FieldElement> VMConverter<T> {
                         }
                         Input::Literal(_, LiteralKind::UnsignedConstant) => {
                             // TODO evaluate expression
-                            if let Expression::Number(Number {value: n, type_: _}) = a {
+                            if let Expression::Number(_, Number {value, ..}) = a {
                                 let half_modulus = T::modulus().to_arbitrary_integer() / BigUint::from(2u64);
-                                assert!(n < half_modulus, "Number passed to unsigned parameter is negative or too large: {n}");
+                                assert!(value < half_modulus, "Number passed to unsigned parameter is negative or too large: {value}");
                                 instruction_literal_arg.push(InstructionLiteralArg::Number(
-                                    T::from(n),
+                                    T::from(value),
                                 ));
                             } else {
                                 panic!("expected unsigned number, received {a}");
                             }
                         }
                         Input::Literal(_, LiteralKind::SignedConstant) => {
                             // TODO evaluate expression
-                            if let Expression::Number(Number {value, ..}) = a {
+                            if let Expression::Number(_, Number {value, ..}) = a {
                                 instruction_literal_arg.push(InstructionLiteralArg::Number(
                                     T::checked_from(value).unwrap(),
                                 ));
-                            } else if let Expression::UnaryOperation(UnaryOperation { op: UnaryOperator::Minus, expr }) = a
+                            } else if let Expression::UnaryOperation(_, UnaryOperation { op: UnaryOperator::Minus, expr }) = a
                             {
-                                if let Expression::Number(Number {value, ..}) = *expr {
+                                if let Expression::Number(_, Number {value, ..}) = *expr {
                                     instruction_literal_arg.push(InstructionLiteralArg::Number(
                                         -T::checked_from(value).unwrap(),
                                     ))
@@ -710,7 +713,7 @@ impl<T: FieldElement> VMConverter<T> {
             .zip(&mut args)
             .map(|(reg, a)| {
                 // Output a value trough assignment register "reg"
-                if let Expression::Reference(r) = a {
+                if let Expression::Reference(_, r) = a {
                     (reg.clone(), vec![r.try_to_identifier().unwrap().clone()])
                 } else {
                     panic!("Expected direct register to assign to in instruction call.");
@@ -730,30 +733,30 @@ impl<T: FieldElement> VMConverter<T> {
 
     fn process_assignment_value(&self, value: Expression) -> Vec<(T, AffineExpressionComponent)> {
         match value {
-            Expression::PublicReference(_) => panic!(),
-            Expression::IndexAccess(_) => panic!(),
-            Expression::FunctionCall(_) => panic!(),
-            Expression::Reference(reference) => {
+            Expression::PublicReference(_, _) => panic!(),
+            Expression::IndexAccess(_, _) => panic!(),
+            Expression::FunctionCall(_, _) => panic!(),
+            Expression::Reference(_, reference) => {
                 // TODO check it actually is a register
                 let name = reference.try_to_identifier().unwrap();
                 vec![(1.into(), AffineExpressionComponent::Register(name.clone()))]
             }
-            Expression::Number(Number { value, .. }) => {
+            Expression::Number(_, Number { value, .. }) => {
                 vec![(T::from(value), AffineExpressionComponent::Constant)]
             }
-            Expression::String(_) => panic!(),
-            Expression::Tuple(_) => panic!(),
-            Expression::ArrayLiteral(_) => panic!(),
-            Expression::MatchExpression(_) => panic!(),
-            Expression::IfExpression(_) => panic!(),
-            Expression::BlockExpression(_) => panic!(),
-            Expression::FreeInput(expr) => {
+            Expression::String(_, _) => panic!(),
+            Expression::Tuple(_, _) => panic!(),
+            Expression::ArrayLiteral(_, _) => panic!(),
+            Expression::MatchExpression(_, _) => panic!(),
+            Expression::IfExpression(_, _) => panic!(),
+            Expression::BlockExpression(_, _) => panic!(),
+            Expression::FreeInput(_, expr) => {
                 vec![(1.into(), AffineExpressionComponent::FreeInput(*expr))]
             }
-            Expression::LambdaExpression(_) => {
+            Expression::LambdaExpression(_, _) => {
                 unreachable!("lambda expressions should have been removed")
             }
-            Expression::BinaryOperation(BinaryOperation { left, op, right }) => match op {
+            Expression::BinaryOperation(_, BinaryOperation { left, op, right }) => match op {
                 BinaryOperator::Add => self.add_assignment_value(
                     self.process_assignment_value(*left),
                     self.process_assignment_value(*right),
@@ -816,7 +819,7 @@ impl<T: FieldElement> VMConverter<T> {
                     panic!("Invalid operation in expression {left} {op} {right}")
                 }
             },
-            Expression::UnaryOperation(UnaryOperation { op, expr }) => {
+            Expression::UnaryOperation(_, UnaryOperation { op, expr }) => {
                 assert!(op == UnaryOperator::Minus);
                 self.negate_assignment_value(self.process_assignment_value(*expr))
             }
@@ -1092,11 +1095,14 @@ impl<T: FieldElement> VMConverter<T> {
         expr: Expression,
     ) -> (usize, Expression) {
         match expr {
-            Expression::BinaryOperation(BinaryOperation {
-                left,
-                op: operator,
-                right,
-            }) => match operator {
+            Expression::BinaryOperation(
+                _,
+                BinaryOperation {
+                    left,
+                    op: operator,
+                    right,
+                },
+            ) => match operator {
                 BinaryOperator::Add => {
                     let (counter, left) = self.linearize_rec(prefix, counter, *left);
                     let (counter, right) = self.linearize_rec(prefix, counter, *right);
@@ -1226,29 +1232,38 @@ fn witness_column<S: Into<String>>(
 }
 
 fn extract_update(expr: Expression) -> (Option<String>, Expression) {
-    let Expression::BinaryOperation(BinaryOperation {
-        left,
-        op: BinaryOperator::Identity,
-        right,
-    }) = expr
+    let Expression::BinaryOperation(
+        _,
+        BinaryOperation {
+            left,
+            op: BinaryOperator::Identity,
+            right,
+        },
+    ) = expr
     else {
         panic!("Invalid statement for instruction body, expected constraint: {expr}");
     };
     // TODO check that there are no other "next" references in the expression
     match *left {
-        Expression::UnaryOperation(UnaryOperation {
-            op: UnaryOperator::Next,
-            expr,
-        }) => match *expr {
-            Expression::Reference(column) => {
+        Expression::UnaryOperation(
+            source_ref,
+            UnaryOperation {
+                op: UnaryOperator::Next,
+                expr: column,
+            },
+        ) => match *column {
+            Expression::Reference(_, column) => {
                 (Some(column.try_to_identifier().unwrap().clone()), *right)
             }
             _ => (
                 None,
-                Expression::UnaryOperation(UnaryOperation {
-                    op: UnaryOperator::Next,
-                    expr,
-                }) - *right,
+                Expression::UnaryOperation(
+                    source_ref,
+                    UnaryOperation {
+                        op: UnaryOperator::Next,
+                        expr: column,
+                    },
+                ) - *right,
             ),
         },
         _ => (None, *left - *right),

diff --git a/ast/src/analyzed/display.rs b/ast/src/analyzed/display.rs
@@ -232,9 +232,9 @@ impl Display for FunctionValueDefinition {
 
 fn format_outer_function(e: &Expression, f: &mut Formatter<'_>) -> Result {
     match e {
-        parsed::Expression::LambdaExpression(lambda) if lambda.params.len() == 1 => {
+        parsed::Expression::LambdaExpression(_, lambda) if lambda.params.len() == 1 => {
             let body = if lambda.kind == FunctionKind::Pure
-                && !matches!(lambda.body.as_ref(), Expression::BlockExpression(_))
+                && !matches!(lambda.body.as_ref(), Expression::BlockExpression(_, _))
             {
                 format!("{{ {} }}", lambda.body)
             } else {

diff --git a/ast/src/analyzed/mod.rs b/ast/src/analyzed/mod.rs
@@ -251,7 +251,7 @@ impl<T> Analyzed<T> {
                 poly.id = replacements[&poly_id].id;
             });
         let visitor = &mut |expr: &mut Expression| {
-            if let Expression::Reference(Reference::Poly(poly)) = expr {
+            if let Expression::Reference(_, Reference::Poly(poly)) = expr {
                 poly.poly_id = poly.poly_id.map(|poly_id| replacements[&poly_id]);
             }
         };
@@ -694,11 +694,14 @@ impl<R> Identity<parsed::Expression<R>> {
     ) -> (&parsed::Expression<R>, Option<&parsed::Expression<R>>) {
         assert_eq!(self.kind, IdentityKind::Polynomial);
         match self.expression_for_poly_id() {
-            parsed::Expression::BinaryOperation(parsed::BinaryOperation {
-                left,
-                op: BinaryOperator::Sub,
-                right,
-            }) => (left.as_ref(), Some(right.as_ref())),
+            parsed::Expression::BinaryOperation(
+                _,
+                parsed::BinaryOperation {
+                    left,
+                    op: BinaryOperator::Sub,
+                    right,
+                },
+            ) => (left.as_ref(), Some(right.as_ref())),
             a => (a, None),
         }
     }

diff --git a/ast/src/parsed/build.rs b/ast/src/parsed/build.rs
@@ -1,6 +1,6 @@
 use powdr_number::BigUint;
 
-use crate::parsed::Expression;
+use crate::parsed::{Expression, SourceReference};
 
 use super::{
     asm::{parse_absolute_path, Part, SymbolPath},
@@ -36,10 +36,13 @@ pub fn next_reference<S: Into<String>>(name: S) -> Expression {
 /// Returns an index access operation to expr if the index is Some, otherwise returns expr itself.
 pub fn index_access(expr: Expression, index: Option<BigUint>) -> Expression {
     match index {
-        Some(i) => Expression::IndexAccess(IndexAccess {
-            array: Box::new(expr),
-            index: Box::new(i.into()),
-        }),
+        Some(i) => Expression::IndexAccess(
+            expr.source_reference().clone(),
+            IndexAccess {
+                array: Box::new(expr),
+                index: Box::new(i.into()),
+            },
+        ),
         None => expr,
     }
 }

diff --git a/ast/src/parsed/display.rs b/ast/src/parsed/display.rs
@@ -531,7 +531,7 @@ impl Display for FunctionDefinition {
             FunctionDefinition::Array(array_expression) => {
                 write!(f, " = {array_expression}")
             }
-            FunctionDefinition::Expression(Expression::LambdaExpression(lambda))
+            FunctionDefinition::Expression(Expression::LambdaExpression(_, lambda))
                 if lambda.params.len() == 1 =>
             {
                 write!(
@@ -598,27 +598,27 @@ fn format_list<L: IntoIterator<Item = I>, I: Display>(list: L) -> String {
 impl<Ref: Display> Display for Expression<Ref> {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         match self {
-            Expression::Reference(reference) => write!(f, "{reference}"),
-            Expression::PublicReference(name) => write!(f, ":{name}"),
-            Expression::Number(Number { value, .. }) => write!(f, "{value}"),
-            Expression::String(value) => write!(f, "{}", quote(value)),
-            Expression::Tuple(items) => write!(f, "({})", format_list(items)),
-            Expression::LambdaExpression(lambda) => write!(f, "{lambda}"),
-            Expression::ArrayLiteral(array) => write!(f, "{array}"),
-            Expression::BinaryOperation(binaryop) => {
+            Expression::Reference(_, reference) => write!(f, "{reference}"),
+            Expression::PublicReference(_, name) => write!(f, ":{name}"),
+            Expression::Number(_, Number { value, .. }) => write!(f, "{value}"),
+            Expression::String(_, value) => write!(f, "{}", quote(value)),
+            Expression::Tuple(_, items) => write!(f, "({})", format_list(items)),
+            Expression::LambdaExpression(_, lambda) => write!(f, "{lambda}"),
+            Expression::ArrayLiteral(_, array) => write!(f, "{array}"),
+            Expression::BinaryOperation(_, binaryop) => {
                 write!(f, "{binaryop}")
             }
-            Expression::UnaryOperation(unaryop) => {
+            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} }}"),
-            Expression::MatchExpression(match_expr) => {
+            Expression::IndexAccess(_, index_access) => write!(f, "{index_access}"),
+            Expression::FunctionCall(_, fun_call) => write!(f, "{fun_call}"),
+            Expression::FreeInput(_, input) => write!(f, "${{ {input} }}"),
+            Expression::MatchExpression(_, match_expr) => {
                 write!(f, "{match_expr}")
             }
-            Expression::IfExpression(e) => write!(f, "{e}"),
-            Expression::BlockExpression(block_expr) => {
+            Expression::IfExpression(_, e) => write!(f, "{e}"),
+            Expression::BlockExpression(_, block_expr) => {
                 write!(f, "{block_expr}")
             }
         }