Repository: microsoft/regorus (language: python)

Issue:
Parse error when rule  body is parsed as a set instead of a comprehension
```rego
import future.keywords

p if { 1 | [1,2][_] }
```

the rule body is parsed as a set resulting in 
```
  |
5 | p if { 1 | [1,2][_] }
  |        ^
error: `union` expects set argument. Got `1` instead
```

It must be parsed as a comprehension instead.


Discussion / hints:

Fix patch:
diff --git a/src/parser.rs b/src/parser.rs
index 1b820958..0c782602 100644
--- a/src/parser.rs
+++ b/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),


Test patch:
diff --git a/tests/interpreter/cases/compr/tricky.yaml b/tests/interpreter/cases/compr/tricky.yaml
new file mode 100644
index 00000000..5b8f7aba
--- /dev/null
+++ b/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
diff --git a/tests/opa.passing b/tests/opa.passing
index 26177ffa..ff90f36a 100644
--- a/tests/opa.passing
+++ b/tests/opa.passing
@@ -227,3 +227,4 @@ v1/uuid
 v1/varreferences
 v1/virtualdocs
 v1/walkbuiltin
+v1/withkeyword