Add support for list comprehensions to parser

air-script-core/src/access.rs

@@ -58,3 +58,23 @@ impl MatrixAccess {
         self.col_idx
     }
 }
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct Range {
+    start: usize,
+    end: usize,
+}
+
+impl Range {
+    pub fn new(start: usize, end: usize) -> Self {
+        Self { start, end }
+    }
+
+    pub fn start(&self) -> usize {
+        self.start
+    }
+
+    pub fn end(&self) -> usize {
+        self.end
+    }
+}

air-script-core/src/lib.rs

@@ -1,5 +1,5 @@
 mod access;
-pub use access::{MatrixAccess, VectorAccess};
+pub use access::{MatrixAccess, Range, VectorAccess};
 
 mod constant;
 pub use constant::{Constant, ConstantType};
@@ -14,4 +14,4 @@ mod trace;
 pub use trace::{IndexedTraceAccess, NamedTraceAccess, TraceSegment};
 
 mod variable;
-pub use variable::{Variable, VariableType};
+pub use variable::{Iterable, ListComprehension, Variable, VariableType};

air-script-core/src/variable.rs

@@ -1,4 +1,4 @@
-use super::{Expression, Identifier};
+use super::{Expression, Identifier, Range};
 
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Variable {
@@ -25,4 +25,38 @@ pub enum VariableType {
     Scalar(Expression),
     Vector(Vec<Expression>),
     Matrix(Vec<Vec<Expression>>),
+    ListComprehension(ListComprehension),
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub struct ListComprehension {
+    expression: Expression,
+    context: Vec<(Identifier, Iterable)>,
+}
+
+impl ListComprehension {
+    /// Creates a new list comprehension.
+    pub fn new(expression: Expression, context: Vec<(Identifier, Iterable)>) -> Self {
+        Self {
+            expression,
+            context,
+        }
+    }
+
+    /// Returns the expression that is evaluated for each member of the list.
+    pub fn expression(&self) -> &Expression {
+        &self.expression
+    }
+
+    /// Returns the context of the list comprehension.
+    pub fn context(&self) -> &[(Identifier, Iterable)] {
+        &self.context
+    }
+}
+
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub enum Iterable {
+    Identifier(Identifier),
+    Range(Range),
+    Slice(Identifier, Range),
 }

parser/src/ast/mod.rs

@@ -1,6 +1,6 @@
 pub(crate) use air_script_core::{
-    Constant, ConstantType, Expression, Identifier, IndexedTraceAccess, MatrixAccess,
-    NamedTraceAccess, Variable, VariableType, VectorAccess,
+    Constant, ConstantType, Expression, Identifier, IndexedTraceAccess, Iterable,
+    ListComprehension, MatrixAccess, NamedTraceAccess, Range, Variable, VariableType, VectorAccess,
 };
 
 pub mod pub_inputs;

parser/src/error.rs

@@ -12,6 +12,7 @@ pub enum ParseError {
     InvalidTraceCols(String),
     MissingMainTraceCols(String),
     InvalidConst(String),
+    InvalidListComprehension(String),
     MissingBoundaryConstraint(String),
     MissingIntegrityConstraint(String),
 }

parser/src/lexer/mod.rs

@@ -91,6 +91,14 @@ pub enum Token {
     #[token("$rand")]
     Rand,
 
+    // LIST COMPREHENSION KEYWORDS
+    // --------------------------------------------------------------------------------------------
+    #[token("for")]
+    For,
+
+    #[token("in")]
+    In,
+
     // GENERAL KEYWORDS
     // --------------------------------------------------------------------------------------------
     /// Keyword to signify that a constraint needs to be enforced
@@ -139,6 +147,9 @@ pub enum Token {
     #[token(")")]
     Rparen,
 
+    #[token("..")]
+    Range,
+
     // UNDEFINED TOKENS AND TOKENS TO IGNORE
     // --------------------------------------------------------------------------------------------
     /// Error is returned on encountering unrecognized tokens.

parser/src/lexer/tests/list_comprehension.rs

@@ -0,0 +1,74 @@
+use super::{expect_valid_tokenization, Token};
+
+// LIST COMPREHENSION VALID TOKENIZATION
+// ================================================================================================
+
+#[test]
+fn one_iterable_comprehension() {
+    let source = "let y = [x for x in x]";
+    let tokens = vec![
+        Token::Let,
+        Token::Ident("y".to_string()),
+        Token::Equal,
+        Token::Lsqb,
+        Token::Ident("x".to_string()),
+        Token::For,
+        Token::Ident("x".to_string()),
+        Token::In,
+        Token::Ident("x".to_string()),
+        Token::Rsqb,
+    ];
+    expect_valid_tokenization(source, tokens);
+}
+
+#[test]
+fn multiple_iterables_comprehension() {
+    let source = "let a = [w + x - y - z for (w, x, y, z) in (0..3, x, y[0..3], z[0..3])]";
+    let tokens = vec![
+        Token::Let,
+        Token::Ident("a".to_string()),
+        Token::Equal,
+        Token::Lsqb,
+        Token::Ident("w".to_string()),
+        Token::Plus,
+        Token::Ident("x".to_string()),
+        Token::Minus,
+        Token::Ident("y".to_string()),
+        Token::Minus,
+        Token::Ident("z".to_string()),
+        Token::For,
+        Token::Lparen,
+        Token::Ident("w".to_string()),
+        Token::Comma,
+        Token::Ident("x".to_string()),
+        Token::Comma,
+        Token::Ident("y".to_string()),
+        Token::Comma,
+        Token::Ident("z".to_string()),
+        Token::Rparen,
+        Token::In,
+        Token::Lparen,
+        Token::Num("0".to_string()),
+        Token::Range,
+        Token::Num("3".to_string()),
+        Token::Comma,
+        Token::Ident("x".to_string()),
+        Token::Comma,
+        Token::Ident("y".to_string()),
+        Token::Lsqb,
+        Token::Num("0".to_string()),
+        Token::Range,
+        Token::Num("3".to_string()),
+        Token::Rsqb,
+        Token::Comma,
+        Token::Ident("z".to_string()),
+        Token::Lsqb,
+        Token::Num("0".to_string()),
+        Token::Range,
+        Token::Num("3".to_string()),
+        Token::Rsqb,
+        Token::Rparen,
+        Token::Rsqb,
+    ];
+    expect_valid_tokenization(source, tokens);
+}

parser/src/lexer/tests/mod.rs

@@ -7,6 +7,7 @@ mod arithmetic_ops;
 mod boundary_constraints;
 mod constants;
 mod identifiers;
+mod list_comprehension;
 mod periodic_columns;
 mod pub_inputs;
 mod variables;

parser/src/parser/grammar.lalrpop

@@ -2,11 +2,13 @@ use crate::{
     ast::{
         boundary_constraints::{Boundary, BoundaryConstraint, BoundaryStmt}, 
         integrity_constraints::{IntegrityConstraint, IntegrityStmt},
-        Constant, ConstantType, Expression, Identifier, Variable, VariableType,
-        Source, SourceSection, Trace, TraceCols, PublicInput, PeriodicColumn, 
-        IndexedTraceAccess, NamedTraceAccess, MatrixAccess, VectorAccess
+        Constant, ConstantType, Expression, Identifier, Variable, VariableType, ListComprehension,
+        Iterable, Source, SourceSection, Trace, TraceCols,
+        PublicInput, PeriodicColumn, IndexedTraceAccess, NamedTraceAccess, Range, MatrixAccess,
+        VectorAccess
     }, error::{Error, ParseError::{InvalidInt, InvalidTraceCols, MissingMainTraceCols,
-    InvalidConst, MissingBoundaryConstraint, MissingIntegrityConstraint}}, lexer::Token
+    InvalidConst, InvalidListComprehension, MissingBoundaryConstraint,
+    MissingIntegrityConstraint}}, lexer::Token
 };
 use std::str::FromStr;
 use lalrpop_util::ParseError;
@@ -155,6 +157,8 @@ BoundaryVariableType: VariableType = {
         VariableType::Vector(vector_value),
     "[" <matrix_value: CommaElems<Vector<BoundaryExpr>>> "]" =>
         VariableType::Matrix(matrix_value),
+    "[" <list_comprehension: ListComprehension<BoundaryExpr>> "]" =>
+        VariableType::ListComprehension(list_comprehension),
 }
 
 Boundary: Boundary = {
@@ -190,7 +194,7 @@ BoundaryAtom: Expression = {
     <n: Num_u64> => Expression::Const(n),
     <ident: Identifier> => Expression::Elem(ident),
     <vector_access: VectorAccess> => Expression::VectorAccess(vector_access),
-    <matrix_access: MatrixAccess> => Expression::MatrixAccess(matrix_access)
+    <matrix_access: MatrixAccess> => Expression::MatrixAccess(matrix_access),
 }
 
 // INTEGRITY CONSTRAINTS
@@ -230,6 +234,8 @@ IntegrityVariableType: VariableType = {
         VariableType::Vector(vector_value),
     "[" <matrix_value: CommaElems<Vector<IntegrityExpr>>> "]" =>
         VariableType::Matrix(matrix_value),
+    "[" <list_comprehension: ListComprehension<IntegrityExpr>> "]" =>
+        VariableType::ListComprehension(list_comprehension),
 }
 
 // --- INTEGRITY CONSTRAINT EXPRESSIONS WITH PRECEDENCE (LOWEST TO HIGHEST) ----------------------
@@ -262,7 +268,7 @@ IntegrityAtom: Expression = {
     <ident: Identifier> => Expression::Elem(ident),
     <vector_access: VectorAccess> => Expression::VectorAccess(vector_access),
     <matrix_access: MatrixAccess> => Expression::MatrixAccess(matrix_access),
-    <trace_access: NamedTraceAccessWithOffset> => Expression::NamedTraceAccess(trace_access)
+    <trace_access: NamedTraceAccessWithOffset> => Expression::NamedTraceAccess(trace_access),
 }
 
 // ATOMS
@@ -277,11 +283,7 @@ CommaElems<T>: Vec<T> = {
 }
 
 Vector<T>: Vec<T> = {
-    "[" <e:T> <v:("," <T>)*> "]" => {
-        let mut v = v;
-        v.insert(0, e);
-        v
-    }
+    "[" <elems: CommaElems<T>> "]" => elems
 }
 
 Size: u64 = {
@@ -318,6 +320,39 @@ IndexedTraceAccess: IndexedTraceAccess = {
     "$aux" <idx: Index> "'" => IndexedTraceAccess::new(1, idx, 1),
 }
 
+ListComprehension<T>: ListComprehension = {
+    <expr: T> "for" <members: Members> "in" <iterables: Iterables> =>?
+        if members.len() != iterables.len() {
+            Err(ParseError::User {
+                error: Error::ParseError(InvalidListComprehension(
+                    "Number of members and iterables must match".to_string()
+                ))
+            })
+        } else {
+            Ok(ListComprehension::new(expr, members.into_iter().zip(iterables).collect::<Vec<_>>()))
+        }
+}
+
+Members: Vec<Identifier> = {
+    <member: Identifier> => vec![member],
+    "(" <members: CommaElems<Identifier>> ")" => members
+}
+
+Iterables: Vec<Iterable> = {
+    <iterable: Iterable> => vec![iterable],
+    "(" <iterables: CommaElems<Iterable>> ")" => iterables
+}
+
+Iterable: Iterable = {
+    <ident: Identifier> => Iterable::Identifier(ident),
+    <range: Range> => Iterable::Range(range),
+    <ident: Identifier> "[" <range: Range> "]" => Iterable::Slice(ident, range)
+}
+
+Range: Range = {
+    <start: Num_u64> ".." <end: Num_u64> => Range::new(start as usize, end as usize)
+}
+
 Identifier: Identifier = {
     <n:identifier> => Identifier(n.to_string())
 }
@@ -341,6 +376,8 @@ extern {
         r"[0-9]+" => Token::Num(<String>),
         "def" => Token::Def,
         "let" => Token::Let,
+        "for" => Token::For,
+        "in" => Token::In,
         "const" => Token::Const,
         "trace_columns" => Token::TraceColumns,
         "main" => Token::MainDecl,
@@ -368,5 +405,6 @@ extern {
         "(" => Token::Lparen,
         ")" => Token::Rparen,
         "." => Token::Dot,
+        ".." => Token::Range,
     }
 }