Merge pull request #733 from powdr-labs/parse_next_as_operator

Parse next as operator

analysis/src/macro_expansion.rs

@@ -173,9 +173,6 @@ where
     fn process_expression(&mut self, e: &mut Expression<T>) {
         if let Expression::Reference(poly) = e {
             if poly.namespace().is_none() && self.parameter_names.contains_key(poly.name()) {
-                // TODO to make this work inside macros, "next" and "index" need to be
-                // their own ast nodes / operators.
-                assert!(!poly.shift());
                 assert!(poly.index().is_none());
                 *e = self.arguments[self.parameter_names[poly.name()]].clone()
             }

asm_to_pil/src/vm_to_constrained.rs

@@ -538,7 +538,6 @@ impl<T: FieldElement> ASMPILConverter<T> {
             .map(|(reg, a)| {
                 // Output a value trough assignment register "reg"
                 if let Expression::Reference(r) = a {
-                    assert!(!r.shift());
                     assert!(r.index().is_none());
                     (reg.clone(), vec![r.name().into()])
                 } else {
@@ -567,7 +566,6 @@ impl<T: FieldElement> ASMPILConverter<T> {
             Expression::Reference(reference) => {
                 assert!(reference.namespace().is_none());
                 assert!(reference.index().is_none());
-                assert!(!reference.shift());
                 // TODO check it actually is a register
                 vec![(
                     1.into(),
@@ -1029,11 +1027,17 @@ fn extract_update<T: FieldElement>(expr: Expression<T>) -> (Option<String>, Expr
     // TODO check that there are no other "next" references in the expression
     if let Expression::BinaryOperation(left, BinaryOperator::Sub, right) = expr {
         match *left {
-            Expression::Reference(r) if r.shift() => {
-                assert!(r.namespace().is_none());
-                assert!(r.index().is_none());
-                (Some(r.name().into()), *right)
-            }
+            Expression::UnaryOperation(UnaryOperator::Next, column) => match *column {
+                Expression::Reference(column) => {
+                    assert!(column.namespace().is_none());
+                    assert!(column.index().is_none());
+                    return (Some(column.name().into()), *right);
+                }
+                _ => (
+                    None,
+                    Expression::UnaryOperation(UnaryOperator::Next, column) - *right,
+                ),
+            },
             _ => (None, *left - *right),
         }
     } else {

ast/src/analyzed/display.rs

@@ -223,13 +223,12 @@ impl Display for PolynomialReference {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         write!(
             f,
-            "{}{}{}",
+            "{}{}",
             self.name,
             self.index
                 .as_ref()
                 .map(|s| format!("[{s}]"))
                 .unwrap_or_default(),
-            if self.next { "'" } else { "" },
         )
     }
 }

ast/src/analyzed/mod.rs

@@ -630,7 +630,6 @@ pub struct PolynomialReference {
     /// TODO make this non-optional
     pub poly_id: Option<PolyID>,
     pub index: Option<u64>,
-    pub next: bool,
 }
 
 #[derive(Debug, Copy, Clone, PartialOrd, Ord, PartialEq, Eq, Hash)]

ast/src/asm_analysis/mod.rs

@@ -18,7 +18,7 @@ use crate::parsed::{
         AbsoluteSymbolPath, AssignmentRegister, CallableRef, InstructionBody, OperationId, Params,
     },
     visitor::{ExpressionVisitable, VisitOrder},
-    PilStatement, ShiftedPolynomialReference,
+    NamespacedPolynomialReference, PilStatement,
 };
 
 pub use crate::parsed::Expression;
@@ -540,10 +540,12 @@ pub enum FunctionStatement<T> {
     Return(Return<T>),
 }
 
-impl<T> ExpressionVisitable<Expression<T, ShiftedPolynomialReference<T>>> for FunctionStatement<T> {
+impl<T> ExpressionVisitable<Expression<T, NamespacedPolynomialReference<T>>>
+    for FunctionStatement<T>
+{
     fn visit_expressions_mut<F, B>(&mut self, f: &mut F, o: VisitOrder) -> std::ops::ControlFlow<B>
     where
-        F: FnMut(&mut Expression<T, ShiftedPolynomialReference<T>>) -> std::ops::ControlFlow<B>,
+        F: FnMut(&mut Expression<T, NamespacedPolynomialReference<T>>) -> std::ops::ControlFlow<B>,
     {
         match self {
             FunctionStatement::Assignment(assignment) => {
@@ -565,7 +567,7 @@ impl<T> ExpressionVisitable<Expression<T, ShiftedPolynomialReference<T>>> for Fu
 
     fn visit_expressions<F, B>(&self, f: &mut F, o: VisitOrder) -> std::ops::ControlFlow<B>
     where
-        F: FnMut(&Expression<T, ShiftedPolynomialReference<T>>) -> std::ops::ControlFlow<B>,
+        F: FnMut(&Expression<T, NamespacedPolynomialReference<T>>) -> std::ops::ControlFlow<B>,
     {
         match self {
             FunctionStatement::Assignment(assignment) => {

ast/src/lib.rs

@@ -72,6 +72,7 @@ pub fn evaluate_unary_operation<T: FieldElement>(op: UnaryOperator, v: T) -> T {
         UnaryOperator::Plus => v,
         UnaryOperator::Minus => -v,
         UnaryOperator::LogicalNot => v.is_zero().into(),
+        UnaryOperator::Next => panic!("Cannot evaluate \"'\"."),
     }
 }
 

ast/src/parsed/build.rs

@@ -1,27 +1,18 @@
 use crate::parsed::Expression;
 
-use super::PolynomialReference;
+use super::{PolynomialReference, UnaryOperator};
 
 pub fn direct_reference<S: Into<String>, T>(name: S) -> Expression<T> {
-    PolynomialReference::new(name)
-        .single()
-        .local()
-        .current()
-        .into()
+    PolynomialReference::new(name).single().local().into()
 }
 
 pub fn namespaced_reference<S: Into<String>, T>(namespace: String, name: S) -> Expression<T> {
     PolynomialReference::new(name)
         .single()
         .namespaced(namespace)
-        .current()
         .into()
 }
 
 pub fn next_reference<T>(name: &str) -> Expression<T> {
-    PolynomialReference::new(name)
-        .single()
-        .local()
-        .next()
-        .into()
+    Expression::UnaryOperation(UnaryOperator::Next, Box::new(direct_reference(name)))
 }

ast/src/parsed/display.rs

@@ -434,7 +434,13 @@ impl<T: Display, Ref: Display> Display for Expression<T, Ref> {
             Expression::LambdaExpression(lambda) => write!(f, "{}", lambda),
             Expression::ArrayLiteral(array) => write!(f, "{array}"),
             Expression::BinaryOperation(left, op, right) => write!(f, "({left} {op} {right})"),
-            Expression::UnaryOperation(op, exp) => write!(f, "{op}{exp}"),
+            Expression::UnaryOperation(op, exp) => {
+                if op.is_prefix() {
+                    write!(f, "{op}{exp}")
+                } else {
+                    write!(f, "{exp}{op}")
+                }
+            }
             Expression::FunctionCall(fun_call) => write!(f, "{fun_call}"),
             Expression::FreeInput(input) => write!(f, "${{ {input} }}"),
             Expression::MatchExpression(scrutinee, arms) => {
@@ -458,12 +464,6 @@ impl<T: Display> Display for PolynomialName<T> {
     }
 }
 
-impl<T: Display> Display for ShiftedPolynomialReference<T> {
-    fn fmt(&self, f: &mut Formatter<'_>) -> Result {
-        write!(f, "{}{}", self.pol, if self.is_next { "'" } else { "" })
-    }
-}
-
 impl<T: Display> Display for NamespacedPolynomialReference<T> {
     fn fmt(&self, f: &mut Formatter<'_>) -> Result {
         write!(
@@ -549,6 +549,7 @@ impl Display for UnaryOperator {
                 UnaryOperator::Minus => "-",
                 UnaryOperator::Plus => "+",
                 UnaryOperator::LogicalNot => "!",
+                UnaryOperator::Next => "'",
             }
         )
     }

ast/src/parsed/mod.rs

@@ -68,7 +68,7 @@ impl<Expr> Default for SelectedExpressions<Expr> {
 }
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
-pub enum Expression<T, Ref = ShiftedPolynomialReference<T>> {
+pub enum Expression<T, Ref = NamespacedPolynomialReference<T>> {
     Reference(Ref),
     PublicReference(String),
     Number(T),
@@ -82,6 +82,7 @@ pub enum Expression<T, Ref = ShiftedPolynomialReference<T>> {
         Box<Expression<T, Ref>>,
     ),
     UnaryOperation(UnaryOperator, Box<Expression<T, Ref>>),
+
     FunctionCall(FunctionCall<T, Ref>),
     FreeInput(Box<Expression<T, Ref>>),
     MatchExpression(Box<Expression<T, Ref>>, Vec<MatchArm<T, Ref>>),
@@ -144,8 +145,8 @@ impl<T: FieldElement, Ref> From<T> for Expression<T, Ref> {
     }
 }
 
-impl<T> From<ShiftedPolynomialReference<T>> for Expression<T> {
-    fn from(value: ShiftedPolynomialReference<T>) -> Self {
+impl<T> From<NamespacedPolynomialReference<T>> for Expression<T> {
+    fn from(value: NamespacedPolynomialReference<T>) -> Self {
         Self::Reference(value)
     }
 }
@@ -156,47 +157,6 @@ pub struct PolynomialName<T> {
     pub array_size: Option<Expression<T>>,
 }
 
-/// A polynomial with an optional shift
-#[derive(Debug, PartialEq, Eq, Default, Clone, PartialOrd, Ord)]
-pub struct ShiftedPolynomialReference<T> {
-    /// Whether we shift or not
-    is_next: bool,
-    /// The underlying polynomial
-    pol: NamespacedPolynomialReference<T>,
-}
-
-impl<T> ShiftedPolynomialReference<T> {
-    /// Returns the underlying namespaced polynomial
-    pub fn into_namespaced(self) -> NamespacedPolynomialReference<T> {
-        self.pol
-    }
-
-    /// Returns the shift of this polynomial
-    pub fn shift(&self) -> bool {
-        self.is_next
-    }
-
-    /// Returns the optional namespace of the underlying polynomial
-    pub fn namespace(&self) -> &Option<String> {
-        self.pol.namespace()
-    }
-
-    /// Returns the optional index of the underlying polynomial in its declaration array
-    pub fn index(&self) -> &Option<Box<Expression<T>>> {
-        self.pol.index()
-    }
-
-    /// Returns the name of the declared polynomial or array of polynomials
-    pub fn name(&self) -> &str {
-        self.pol.name()
-    }
-
-    /// Returns a mutable reference to the declared polynomial or array of polynomials
-    pub fn name_mut(&mut self) -> &mut String {
-        self.pol.name_mut()
-    }
-}
-
 #[derive(Debug, PartialEq, Eq, Default, Clone, PartialOrd, Ord)]
 /// A polynomial with an optional namespace
 pub struct NamespacedPolynomialReference<T> {
@@ -207,14 +167,6 @@ pub struct NamespacedPolynomialReference<T> {
 }
 
 impl<T> NamespacedPolynomialReference<T> {
-    /// Return a shifted polynomial based on this namespaced polynomial and a boolean shift
-    pub fn with_shift(self, next: bool) -> ShiftedPolynomialReference<T> {
-        ShiftedPolynomialReference {
-            is_next: next,
-            pol: self,
-        }
-    }
-
     /// Returns the optional namespace of this polynomial
     pub fn namespace(&self) -> &Option<String> {
         &self.namespace
@@ -234,16 +186,6 @@ impl<T> NamespacedPolynomialReference<T> {
     pub fn name_mut(&mut self) -> &mut String {
         self.pol.name_mut()
     }
-
-    /// Return a shifted polynomial based on this namespaced polynomial with a shift of 1
-    pub fn next(self) -> ShiftedPolynomialReference<T> {
-        self.with_shift(true)
-    }
-
-    /// Return a shifted polynomial based on this namespaced polynomial with a shift of 0
-    pub fn current(self) -> ShiftedPolynomialReference<T> {
-        self.with_shift(false)
-    }
 }
 
 #[derive(Debug, PartialEq, Eq, Default, Clone, PartialOrd, Ord)]
@@ -333,13 +275,13 @@ impl PolynomialReference {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
-pub struct LambdaExpression<T, Ref = ShiftedPolynomialReference<T>> {
+pub struct LambdaExpression<T, Ref = NamespacedPolynomialReference<T>> {
     pub params: Vec<String>,
     pub body: Box<Expression<T, Ref>>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
-pub struct ArrayLiteral<T, Ref = ShiftedPolynomialReference<T>> {
+pub struct ArrayLiteral<T, Ref = NamespacedPolynomialReference<T>> {
     pub items: Vec<Expression<T, Ref>>,
 }
 
@@ -348,6 +290,17 @@ pub enum UnaryOperator {
     Plus,
     Minus,
     LogicalNot,
+    Next,
+}
+
+impl UnaryOperator {
+    /// Returns true if the operator is a prefix-operator and false if it is a postfix operator.
+    pub fn is_prefix(&self) -> bool {
+        match self {
+            UnaryOperator::Plus | UnaryOperator::Minus | UnaryOperator::LogicalNot => true,
+            UnaryOperator::Next => false,
+        }
+    }
 }
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
@@ -374,20 +327,20 @@ pub enum BinaryOperator {
 }
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
-pub struct FunctionCall<T, Ref = ShiftedPolynomialReference<T>> {
+pub struct FunctionCall<T, Ref = NamespacedPolynomialReference<T>> {
     pub id: String,
     pub arguments: Vec<Expression<T, Ref>>,
 }
 
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
-pub struct MatchArm<T, Ref = ShiftedPolynomialReference<T>> {
+pub struct MatchArm<T, Ref = NamespacedPolynomialReference<T>> {
     pub pattern: MatchPattern<T, Ref>,
     pub value: Expression<T, Ref>,
 }
 
 /// A pattern for a match arm. We could extend this in the future.
 #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone)]
-pub enum MatchPattern<T, Ref = ShiftedPolynomialReference<T>> {
+pub enum MatchPattern<T, Ref = NamespacedPolynomialReference<T>> {
     CatchAll,
     Pattern(Expression<T, Ref>),
 }

ast/src/parsed/visitor.rs

@@ -2,7 +2,7 @@ use std::ops::ControlFlow;
 
 use super::{
     ArrayExpression, ArrayLiteral, Expression, FunctionCall, FunctionDefinition, LambdaExpression,
-    MatchArm, MatchPattern, PilStatement, SelectedExpressions, ShiftedPolynomialReference,
+    MatchArm, MatchPattern, NamespacedPolynomialReference, PilStatement, SelectedExpressions,
 };
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -180,10 +180,10 @@ impl<T, Ref> ExpressionVisitable<Expression<T, Ref>> for Expression<T, Ref> {
     }
 }
 
-impl<T> ExpressionVisitable<Expression<T, ShiftedPolynomialReference<T>>> for PilStatement<T> {
+impl<T> ExpressionVisitable<Expression<T, NamespacedPolynomialReference<T>>> for PilStatement<T> {
     fn visit_expressions_mut<F, B>(&mut self, f: &mut F, o: VisitOrder) -> ControlFlow<B>
     where
-        F: FnMut(&mut Expression<T, ShiftedPolynomialReference<T>>) -> ControlFlow<B>,
+        F: FnMut(&mut Expression<T, NamespacedPolynomialReference<T>>) -> ControlFlow<B>,
     {
         match self {
             PilStatement::FunctionCall(_, _, arguments) => arguments

backend/src/pilstark/json_exporter/mod.rs

@@ -59,7 +59,6 @@ pub fn export<T: FieldElement>(analyzed: &Analyzed<T>) -> PIL {
             StatementIdentifier::PublicDeclaration(name) => {
                 let pub_def = &analyzed.public_declarations[name];
                 let pub_ref = &pub_def.polynomial;
-                assert!(!pub_ref.next);
                 let (_, expr) = exporter.polynomial_reference_to_json(&AlgebraicReference {
                     name: pub_ref.name.clone(),
                     poly_id: pub_ref.poly_id.unwrap(),