fix: Handle parsing corner cases (#309)

A block with a single or expression needs to be treated as a comprehension instead of a
set/array with 1 item. e.g.: {1 | 1 }, [2 | foo]

Allow successfully parsing object comprehensions as rule body
x if { 1:2 | 1 }

fixes #306, fixes #307

Signed-off-by: Anand Krishnamoorthi <[email protected]>

src/parser.rs

@@ -1038,6 +1038,23 @@ impl<'source> Parser<'source> {
         let mut literals = vec![];
 
         let stmt = match self.parse_literal_stmt() {
+            Ok(_) if self.token_text() == ":" => {
+                // This is likely an object comprehension.
+                // Restore the state and return.
+                *self = state;
+                bail!("try parsing as comprehension");
+            }
+            Ok(stmt) if self.token_text() == end_delim => {
+                // Treat { 1 | 1 } as a comprehension instead of a
+                // set of 1 element.
+                if let Literal::Expr { expr: e, .. } = &stmt.literal {
+                    if matches!(e.as_ref(), Expr::BinExpr { op: BinOp::Or, .. }) {
+                        *self = state;
+                        bail!("try parse as comprehension");
+                    }
+                }
+                stmt
+            }
             Ok(stmt) => stmt,
             Err(e) if is_definite_query => return Err(e),
             Err(e) if matches!(self.token_text(), "=" | ":=") => return Err(e),

tests/interpreter/cases/compr/tricky.yaml

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+cases:
+  - note: treat { 1 | 1 } as comprehension
+    data: {}
+    modules:
+      - |
+        package test
+        import rego.v1
+
+        x if { 1 | 1 }
+        y := { 1 | 1 }
+        z := { (1) | 1}
+
+        # Parsed as a set
+        a := { ({1} | {2}) }
+
+        b := v if { { 1 | 1 } = v } 
+      
+    query: data.test
+    want_result:
+      x: true
+      y:
+        set!: [1]
+      z:
+        set!: [1]
+      a:
+        set!:
+          - set!: [1, 2]
+      b:
+        set!: [1]
+
+  - note: rule body is object comprehension
+    data: {}
+    modules:
+      - |
+        package test
+        import future.keywords
+        x if { 1:2 | some p in [1,2] }
+        y := 2 if { 1:2 | some p in [1,2] }
+    query: data.test
+    want_result:
+      x: true
+      y: 2

tests/opa.passing

@@ -227,3 +227,4 @@ v1/uuid
 v1/varreferences
 v1/virtualdocs
 v1/walkbuiltin
+v1/withkeyword