From 8003cfc5b7a1e8e912f3f2ebbd945f2b3e73fd92 Mon Sep 17 00:00:00 2001
From: Anand Krishnamoorthi <35780660+anakrish@users.noreply.github.com>
Date: Wed, 4 Sep 2024 10:02:14 -0700
Subject: [PATCH] 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 <anakrish@microsoft.com>
---
 src/parser.rs                             | 17 +++++++++
 tests/interpreter/cases/compr/tricky.yaml | 45 +++++++++++++++++++++++
 tests/opa.passing                         |  1 +
 3 files changed, 63 insertions(+)
 create mode 100644 tests/interpreter/cases/compr/tricky.yaml

diff --git a/src/parser.rs b/src/parser.rs
index 1b82095..0c78260 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),
diff --git a/tests/interpreter/cases/compr/tricky.yaml b/tests/interpreter/cases/compr/tricky.yaml
new file mode 100644
index 0000000..5b8f7ab
--- /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 26177ff..ff90f36 100644
--- a/tests/opa.passing
+++ b/tests/opa.passing
@@ -227,3 +227,4 @@ v1/uuid
 v1/varreferences
 v1/virtualdocs
 v1/walkbuiltin
+v1/withkeyword