test: add list comprehension parser tests

0xPolygonMiden · Jan 30, 2023 · 8db9bf4 · 8db9bf4
1 parent 76079c7
commit 8db9bf4
Show file tree

Hide file tree

Showing 3 changed files with 523 additions and 9 deletions.
diff --git a/parser/src/parser/grammar.lalrpop b/parser/src/parser/grammar.lalrpop
@@ -157,7 +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),
+    "[" <list_comprehension: ListComprehension<BoundaryExpr>> "]" =>
+        VariableType::ListComprehension(list_comprehension),
 }
 
 Boundary: Boundary = {
@@ -233,7 +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),
+    "[" <list_comprehension: ListComprehension<IntegrityExpr>> "]" =>
+        VariableType::ListComprehension(list_comprehension),
 }
 
 // --- INTEGRITY CONSTRAINT EXPRESSIONS WITH PRECEDENCE (LOWEST TO HIGHEST) ----------------------
@@ -323,15 +325,10 @@ IndexedTraceAccess: IndexedTraceAccess = {
 }
 
 ListComprehension<T>: ListComprehension = {
-    "[" <expr: T> "for" <members: Members> "in" <iterables: Iterables> "]" =>?
+    <expr: T> "for" <members: Members> "in" <iterables: Iterables> =>?
     match (members, iterables)
     {
-        (Members::Single(..), Iterables::Double(..)) => Err(ParseError::User {
-            error: Error::ParseError(
-                InvalidListComprehension("For a single member, a single iterable is required".to_string())
-            )
-        }),
-        (Members::Single(..), Iterables::Range(..)) => Err(ParseError::User {
+        (Members::Single(..), Iterables::Double(..) | Iterables::Range(..)) => Err(ParseError::User {
             error: Error::ParseError(
                 InvalidListComprehension("For a single member, a single iterable is required".to_string())
             )