feat(parser): add support for variables to parser

parser/src/ast/boundary_constraints.rs

@@ -7,9 +7,22 @@ use std::fmt::Display;
 /// Stores the boundary constraints to be enforced on the trace column values.
 #[derive(Debug, PartialEq)]
 pub struct BoundaryConstraints {
+    pub variables: Vec<BoundaryVariable>,
     pub boundary_constraints: Vec<BoundaryConstraint>,
 }
 
+impl BoundaryConstraints {
+    pub fn new(
+        variables: Vec<BoundaryVariable>,
+        boundary_constraints: Vec<BoundaryConstraint>,
+    ) -> Self {
+        Self {
+            variables,
+            boundary_constraints,
+        }
+    }
+}
+
 /// Stores the expression corresponding to the boundary constraint.
 #[derive(Debug, PartialEq)]
 pub struct BoundaryConstraint {
@@ -77,3 +90,27 @@ pub enum BoundaryExpr {
     Mul(Box<BoundaryExpr>, Box<BoundaryExpr>),
     Exp(Box<BoundaryExpr>, u64),
 }
+
+#[derive(Debug, PartialEq)]
+pub struct BoundaryVariable {
+    name: Identifier,
+    value: BoundaryVariableType,
+}
+
+impl BoundaryVariable {
+    pub fn new(name: Identifier, value: BoundaryVariableType) -> Self {
+        Self { name, value }
+    }
+}
+
+#[derive(Debug, PartialEq)]
+pub enum BoundaryVariableType {
+    Scalar(BoundaryExpr),
+    Vector(Vec<BoundaryExpr>),
+    Matrix(Vec<Vec<BoundaryExpr>>),
+}
+
+pub enum BoundaryElem {
+    Constraint(BoundaryConstraint),
+    Variable(BoundaryVariable),
+}

parser/src/ast/transition_constraints.rs

@@ -6,9 +6,22 @@ use super::Identifier;
 /// Stores the transition constraints to be enforced on the trace column values.
 #[derive(Debug, PartialEq)]
 pub struct TransitionConstraints {
+    pub variables: Vec<TransitionVariable>,
     pub transition_constraints: Vec<TransitionConstraint>,
 }
 
+impl TransitionConstraints {
+    pub fn new(
+        variables: Vec<TransitionVariable>,
+        transition_constraints: Vec<TransitionConstraint>,
+    ) -> Self {
+        Self {
+            variables,
+            transition_constraints,
+        }
+    }
+}
+
 /// Stores the expression corresponding to the transition constraint.
 #[derive(Debug, PartialEq, Clone)]
 pub struct TransitionConstraint {
@@ -44,3 +57,27 @@ pub enum TransitionExpr {
     Mul(Box<TransitionExpr>, Box<TransitionExpr>),
     Exp(Box<TransitionExpr>, u64),
 }
+
+#[derive(Debug, PartialEq)]
+pub struct TransitionVariable {
+    name: Identifier,
+    value: TransitionVariableType,
+}
+
+impl TransitionVariable {
+    pub fn new(name: Identifier, value: TransitionVariableType) -> Self {
+        Self { name, value }
+    }
+}
+
+#[derive(Debug, PartialEq)]
+pub enum TransitionVariableType {
+    Scalar(TransitionExpr),
+    Vector(Vec<TransitionExpr>),
+    Matrix(Vec<Vec<TransitionExpr>>),
+}
+
+pub enum TransitionElem {
+    Constraint(TransitionConstraint),
+    Variable(TransitionVariable),
+}

parser/src/error.rs

@@ -11,4 +11,6 @@ pub enum ParseError {
     InvalidInt(String),
     InvalidTraceCols(String),
     MissingMainTraceCols(String),
+    MissingBoundaryConstraint(String),
+    MissingTransitionConstraint(String),
 }

parser/src/lexer/mod.rs

@@ -27,6 +27,10 @@ pub enum Token {
     #[token("def")]
     Def,
 
+    /// Used to declare intermediate variables in the AIR constraints module.
+    #[token("let")]
+    Let,
+
     /// Used to declare constants in the AIR constraints module.
     #[token("constants")]
     Constants,

parser/src/parser/grammar.lalrpop

@@ -1,10 +1,13 @@
 use crate::{
     ast::{
-        boundary_constraints::{Boundary, BoundaryConstraints, BoundaryConstraint, BoundaryExpr}, 
-        transition_constraints::{TransitionConstraint, TransitionConstraints, TransitionExpr},
+        boundary_constraints::{Boundary, BoundaryConstraints, BoundaryConstraint, BoundaryElem,
+        BoundaryExpr, BoundaryVariable, BoundaryVariableType}, 
+        transition_constraints::{TransitionConstraint, TransitionConstraints, TransitionExpr,
+        TransitionElem, TransitionVariable, TransitionVariableType},
         constants::{Constant, ConstantType},
         Identifier, Source, SourceSection, TraceCols, PublicInput, PeriodicColumn
-    }, error::{Error, ParseError::{InvalidInt, InvalidTraceCols, MissingMainTraceCols}}, lexer::Token
+    }, error::{Error, ParseError::{InvalidInt, InvalidTraceCols, MissingMainTraceCols,
+    MissingBoundaryConstraint, MissingTransitionConstraint }}, lexer::Token
 };
 use std::str::FromStr;
 use lalrpop_util::ParseError;
@@ -63,7 +66,7 @@ AuxCols: Vec<Identifier> = {
 // CONSTANTS
 // ================================================================================================
 
-// Constants section is optional but cannot be empty. Atleast one constant is required.
+// Constants section is optional but cannot be empty. At least one constant is required.
 Constants: Vec<Constant> = {
     "constants" ":" <constants: Constant+> => constants
 }
@@ -105,13 +108,38 @@ PeriodicColumn: PeriodicColumn = {
 // ================================================================================================
 
 BoundaryConstraints: BoundaryConstraints = {
-    "boundary_constraints" ":" <boundary_constraints: BoundaryConstraint+> =>
-        BoundaryConstraints { boundary_constraints }
+    "boundary_constraints" ":" <boundary_elems: BoundaryElem+> =>? {
+        let mut variables = Vec::new();
+        let mut boundary_constraints = Vec::new();
+        for elem in boundary_elems {
+            match elem {
+                BoundaryElem::Constraint(constraint) => boundary_constraints.push(constraint),
+                BoundaryElem::Variable(variable) => variables.push(variable),
+            }
+        }
+        if boundary_constraints.is_empty() {
+            return Err(ParseError::User {
+                error: Error::ParseError(MissingBoundaryConstraint("Declaration of at least one boundary constraint is required".to_string()))
+            });
+        }
+        Ok(BoundaryConstraints{ variables, boundary_constraints })
+    }
 }
 
-BoundaryConstraint: BoundaryConstraint = {
+BoundaryElem: BoundaryElem = {
     "enf" <column: Identifier> "." <boundary: Boundary> "=" <value: BoundaryExpr> =>
-        BoundaryConstraint::new(column, boundary, value)
+        BoundaryElem::Constraint(BoundaryConstraint::new(column, boundary, value)),
+    "let" <name: Identifier> "=" <boundary_variable_type: BoundaryVariableType> =>
+        BoundaryElem::Variable(BoundaryVariable::new(name, boundary_variable_type)),
+}
+
+BoundaryVariableType: BoundaryVariableType = {
+    <scalar_value: BoundaryExpr> =>
+        BoundaryVariableType::Scalar(scalar_value),
+    "[" <vector_value: CommaElems<BoundaryExpr>> "]" =>
+        BoundaryVariableType::Vector(vector_value),
+    "[" <matrix_value: CommaElems<Vector<BoundaryExpr>>> "]" =>
+        BoundaryVariableType::Matrix(matrix_value),
 }
 
 Boundary: Boundary = {
@@ -147,12 +175,37 @@ BoundaryAtom: BoundaryExpr = {
 // ================================================================================================
 
 TransitionConstraints: TransitionConstraints = {
-    "transition_constraints" ":" <transition_constraints: TransitionConstraint+> =>
-        TransitionConstraints { transition_constraints }
+    "transition_constraints" ":" <transition_elems: TransitionElem+> =>? {
+        let mut variables = Vec::new();
+        let mut transition_constraints = Vec::new();
+        for elem in transition_elems {
+            match elem {
+                TransitionElem::Constraint(constraint) => transition_constraints.push(constraint),
+                TransitionElem::Variable(variable) => variables.push(variable),
+            }
+        }
+        if transition_constraints.is_empty() {
+            return Err(ParseError::User {
+                error: Error::ParseError(MissingTransitionConstraint("Declaration of at least one transition constraint is required".to_string()))
+            });
+        }
+        Ok(TransitionConstraints{ variables, transition_constraints })
+    }
 }
 
-TransitionConstraint: TransitionConstraint = {
-    "enf" <lhs: TransitionExpr> "=" <rhs: TransitionExpr> => TransitionConstraint::new(lhs, rhs)
+TransitionElem: TransitionElem = {
+    "enf" <lhs: TransitionExpr> "=" <rhs: TransitionExpr> => TransitionElem::Constraint(TransitionConstraint::new(lhs, rhs)),
+    "let" <name: Identifier> "=" <transition_variable_type: TransitionVariableType> =>
+        TransitionElem::Variable(TransitionVariable::new(name, transition_variable_type)),
+}
+
+TransitionVariableType: TransitionVariableType = {
+    <scalar_value: TransitionExpr> =>
+        TransitionVariableType::Scalar(scalar_value),
+    "[" <vector_value: CommaElems<TransitionExpr>> "]" =>
+        TransitionVariableType::Vector(vector_value),
+    "[" <matrix_value: CommaElems<Vector<TransitionExpr>>> "]" =>
+        TransitionVariableType::Matrix(matrix_value),
 }
 
 // --- TRANSITION CONSTRAINT EXPRESSIONS WITH PRECEDENCE (LOWEST TO HIGHEST) ----------------------
@@ -221,8 +274,9 @@ extern {
         identifier => Token::Ident(<String>),
         r"[0-9]+" => Token::Num(<String>),
         "def" => Token::Def,
-        "trace_columns" => Token::TraceColumnns,
+        "let" => Token::Let,
         "constants" => Token::Constants,
+        "trace_columns" => Token::TraceColumnns,
         "main" => Token::Main,
         "aux" => Token::Aux,
         "public_inputs" => Token::PublicInputs,

parser/src/parser/tests/boundary_constraints.rs

@@ -2,9 +2,15 @@ use super::{
     build_parse_test, Boundary, BoundaryConstraint, BoundaryConstraints, BoundaryExpr, Identifier,
     Source, SourceSection,
 };
-use crate::ast::constants::{
-    Constant,
-    ConstantType::{Matrix, Scalar, Vector},
+use crate::{
+    ast::{
+        constants::{
+            Constant,
+            ConstantType::{Matrix, Scalar, Vector},
+        },
+        BoundaryVariable, BoundaryVariableType,
+    },
+    error::{Error, ParseError},
 };
 
 // BOUNDARY CONSTRAINTS
@@ -17,6 +23,7 @@ fn boundary_constraint_at_first() {
         enf clk.first = 0";
     let expected = Source(vec![SourceSection::BoundaryConstraints(
         BoundaryConstraints {
+            variables: vec![],
             boundary_constraints: vec![BoundaryConstraint::new(
                 Identifier("clk".to_string()),
                 Boundary::First,
@@ -34,6 +41,7 @@ fn boundary_constraint_at_last() {
         enf clk.last = 15";
     let expected = Source(vec![SourceSection::BoundaryConstraints(
         BoundaryConstraints {
+            variables: vec![],
             boundary_constraints: vec![BoundaryConstraint::new(
                 Identifier("clk".to_string()),
                 Boundary::Last,
@@ -60,6 +68,7 @@ fn multiple_boundary_constraints() {
         enf clk.last = 1";
     let expected = Source(vec![SourceSection::BoundaryConstraints(
         BoundaryConstraints {
+            variables: vec![],
             boundary_constraints: vec![
                 BoundaryConstraint::new(
                     Identifier("clk".to_string()),
@@ -84,6 +93,7 @@ fn boundary_constraint_with_pub_input() {
         enf clk.first = a[0]";
     let expected = Source(vec![SourceSection::BoundaryConstraints(
         BoundaryConstraints {
+            variables: vec![],
             boundary_constraints: vec![BoundaryConstraint::new(
                 Identifier("clk".to_string()),
                 Boundary::First,
@@ -101,6 +111,7 @@ fn boundary_constraint_with_expr() {
         enf clk.first = 5 + a[3] + 6";
     let expected = Source(vec![SourceSection::BoundaryConstraints(
         BoundaryConstraints {
+            variables: vec![],
             boundary_constraints: vec![BoundaryConstraint::new(
                 Identifier("clk".to_string()),
                 Boundary::First,
@@ -142,6 +153,7 @@ fn boundary_constraint_with_const() {
             },
         ]),
         SourceSection::BoundaryConstraints(BoundaryConstraints {
+            variables: vec![],
             boundary_constraints: vec![BoundaryConstraint::new(
                 Identifier("clk".to_string()),
                 Boundary::First,
@@ -158,6 +170,83 @@ fn boundary_constraint_with_const() {
     build_parse_test!(source).expect_ast(expected);
 }
 
+#[test]
+fn boundary_constraint_with_variables() {
+    let source = "
+    boundary_constraints:
+        let a = 2^2
+        let b = [a, 2 * a]
+        let c = [[a - 1, a^2], [b[0], b[1]]]
+        enf clk.first = 5 + a[3] + 6";
+    let expected = Source(vec![SourceSection::BoundaryConstraints(
+        BoundaryConstraints {
+            variables: vec![
+                BoundaryVariable::new(
+                    Identifier("a".to_string()),
+                    BoundaryVariableType::Scalar(BoundaryExpr::Exp(
+                        Box::new(BoundaryExpr::Const(2)),
+                        2,
+                    )),
+                ),
+                BoundaryVariable::new(
+                    Identifier("b".to_string()),
+                    BoundaryVariableType::Vector(vec![
+                        BoundaryExpr::Var(Identifier("a".to_string())),
+                        BoundaryExpr::Mul(
+                            Box::new(BoundaryExpr::Const(2)),
+                            Box::new(BoundaryExpr::Var(Identifier("a".to_string()))),
+                        ),
+                    ]),
+                ),
+                BoundaryVariable::new(
+                    Identifier("c".to_string()),
+                    BoundaryVariableType::Matrix(vec![
+                        vec![
+                            BoundaryExpr::Sub(
+                                Box::new(BoundaryExpr::Var(Identifier("a".to_string()))),
+                                Box::new(BoundaryExpr::Const(1)),
+                            ),
+                            BoundaryExpr::Exp(
+                                Box::new(BoundaryExpr::Var(Identifier("a".to_string()))),
+                                2,
+                            ),
+                        ],
+                        vec![
+                            BoundaryExpr::VecElem(Identifier("b".to_string()), 0),
+                            BoundaryExpr::VecElem(Identifier("b".to_string()), 1),
+                        ],
+                    ]),
+                ),
+            ],
+            boundary_constraints: vec![BoundaryConstraint::new(
+                Identifier("clk".to_string()),
+                Boundary::First,
+                BoundaryExpr::Add(
+                    Box::new(BoundaryExpr::Add(
+                        Box::new(BoundaryExpr::Const(5)),
+                        Box::new(BoundaryExpr::VecElem(Identifier("a".to_string()), 3)),
+                    )),
+                    Box::new(BoundaryExpr::Const(6)),
+                ),
+            )],
+        },
+    )]);
+    build_parse_test!(source).expect_ast(expected);
+}
+
+#[test]
+fn err_missing_boundary_constraint() {
+    let source = "
+    boundary_constraints:
+        let a = 2^2
+        let b = [a, 2 * a]
+        let c = [[a - 1, a^2], [b[0], b[1]]]";
+    let error = Error::ParseError(ParseError::MissingBoundaryConstraint(
+        "Declaration of at least one boundary constraint is required".to_string(),
+    ));
+    build_parse_test!(source).expect_error(error);
+}
+
 #[test]
 fn err_empty_boundary_constraints() {
     let source = "