Code from github.com/anakrish/rego-rs

Authored by anakrish and mingweishih

Signed-off-by: Anand Krishnamoorthi <anakrish@microsoft.com>

tests/interpreter/cases/compr/array-vs-compr-tricky.yaml

@@ -0,0 +1,132 @@
+# Copyright (c) Rego-Rs Authors.
+# Licensed under the Apache 2.0 license.
+
+# Captures tricky cases where similar looking expressions will be
+# treated as arrays or array-comprs.
+cases:
+  - note: case1-compr
+    data: {}
+    modules:
+      - |
+        package test
+        x = {1}
+        y = {2}
+        # If there is one term and the term has an | operator, it will be parsed
+        # as a compr.
+        z = [ x | y ] #, 5] # in {5}, 6]
+    query: data.test
+    sort_bindings: true
+    want_result:
+      x:
+        set!: [1]
+      y:
+        set!: [2]
+      z:
+        - set!: [1]
+
+  - note: case2-array
+    data: {}
+    modules:
+      - |
+        package test
+        x = {1}
+        y = {2}
+        # If there is a comma following the expression to the right of |
+        # ie "y ," then enclosing expression will be parsed as array.
+        z = [ x | y, 5] # in {5}, 6]
+    query: data.test
+    sort_bindings: true
+    want_result:
+      x:
+        set!: [1]
+      y:
+        set!: [2]
+      z:
+        - set!: [2, 1]
+        - 5
+
+  - note: case3-compr
+    data: {}
+    modules:
+      - |
+        package test
+        import future.keywords.in
+        x = {1}
+        y = {2}
+        # In the following case, the "y, 5 in {5}" which is to the right of |
+        # will be parsed as a membership expression. Since this expression is not
+        # followed by a ",", the enclosing expression will be parsed as a compr.
+        z = [ x | y, 5 in {5}] #, 6]
+    query: data.test
+    sort_bindings: true
+    want_result:
+      x:
+        set!: [1]
+      y:
+        set!: [2]
+      z: []
+
+  - note: case4-array
+    data: {}
+    modules:
+      - |
+        package test
+        import future.keywords.in
+        x = {1}
+        y = {2}
+        # Unlike the previous case, the "y, 5 in {5}" is followed by a ",".
+        # Therefore, the enclosing expression will be parsed as an array.
+        z = [ x | y, 5 in {5}, 6]
+    query: data.test
+    sort_bindings: true
+    want_result:
+      x:
+        set!: [1]
+      y:
+        set!: [2]
+      z:
+        - set!: [2, 1]
+        - true
+        - 6
+
+  - note: case5-array
+    data: {}
+    modules:
+      - |
+        package test
+        import future.keywords.in
+        x = {1}
+        y = {2}
+        # If the first expression "x - {10}"is something higher than an in-expr
+        # (e.g arithexpr) then the enclosing expression will be parsed as an array.
+        z = [ x - {10} | y, 5 in {5}] #, 6]
+    query: data.test
+    sort_bindings: true
+    want_result:
+      x:
+        set!: [1]
+      y:
+        set!: [2]
+      z:
+        - set!: [2, 1]
+        - true
+
+  - note: case6-compr
+    data: {}
+    modules:
+      - |
+        package test
+        import future.keywords.in
+        x = {1}
+        y = {2}
+        # The way to get the previous case parsed as a compr is to enclose the first
+        # expression in parentheses.
+        z = [ (x - {10}) | y, 5 in {5}] #, 6]
+    query: data.test
+    sort_bindings: true
+    want_result:
+      x:
+        set!: [1]
+      y:
+        set!: [2]
+      z: []

tests/interpreter/cases/compr/mod.rs

@@ -0,0 +1,123 @@
+// Copyright (c) Rego-Rs Authors.
+// Licensed under the Apache 2.0 license.
+#![cfg(test)]
+
+use crate::interpreter::*;
+
+#[test]
+fn basic_array() -> Result<()> {
+    let rego = r#"
+    package test
+
+    array = [1, 2, 3]
+
+    array_compr_0 = [ x | x = 1 ]
+
+    array_compr_1 = [ array | true ]
+
+    array_compr_2 = [ x | x = array[_] ]
+
+    array_compr_3 = [ x | x = array[_]; x != 2 ]
+
+    # This produces 3 values
+    array_compr_4 = [ 1 | [1, 2, 3][_] ]
+
+    # This produces 6 values.
+    array_compr_5 = [ 1 | [1, 2, 3][_]; {1, 2}[_] ]
+
+    # This also produces 6 values.
+    array_compr_6 = [ 1 | [1, 2, 3][_]; {"a":1, "b":2}[_] ]
+
+    # This produces 3 values.
+    array_compr_7 = [ 1 | [1, 2, 3][_]; [1, 2][_] >= 2 ]
+"#;
+
+    let expected = Value::from_json_str(
+        r#" {
+            "array": [1, 2, 3],
+            "array_compr_0": [1],
+            "array_compr_1": [[1, 2, 3]],
+            "array_compr_2": [1, 2, 3],
+            "array_compr_3": [1, 3],
+            "array_compr_4": [1, 1, 1],
+            "array_compr_5": [1, 1, 1, 1, 1, 1],
+            "array_compr_6": [1, 1, 1, 1, 1, 1],
+            "array_compr_7": [1, 1, 1]
+}"#,
+    )?;
+
+    assert_match(
+        eval_file(&[rego.to_owned()], None, None, "data.test")?,
+        expected,
+    );
+    Ok(())
+}
+
+#[test]
+fn basic_set() -> Result<()> {
+    let rego = r#"
+    package test
+
+    set = { 1,  "string", 1, [2, 3, 4], 567, false, 1 }
+
+    set_compr_0 = { x | x = 1 }
+
+    set_compr_1 = { set | true }
+
+    set_compr_2 = { x | x = set[_] }
+
+    set_compr_3 = { x | x = set[_]; x != [2, 3, 4] }
+
+    # This produces 1 value
+    set_compr_4 = { 1 | [1, 2, 3][_] }
+
+    # This produces 4 values.
+    set_compr_5 = { (a+b) | a=[1, 2, 3][_]; b={1, 2}[_] }
+
+    # This also produces 2 values.
+    set_compr_6 = { a | [1, 2, 3][_]; a={"a":1, "b":2}[_] }
+
+    # This produces 3 values.
+    set_compr_7 = { a | a = [1, 2, 3][_]; [1, 2][_] >= 2 }
+"#;
+
+    let expected = Value::from_json_str(
+        r#" {
+            "set": {
+                "set!": [1, "string", [2, 3, 4], 567, false]
+            },
+            "set_compr_0": {
+                "set!": [1]
+            },
+           "set_compr_1":  {
+                "set!" : [{
+                    "set!": [1, "string", [2, 3, 4], 567, false]
+                }]
+            },
+            "set_compr_2":  {
+                "set!": [1, "string", [2, 3, 4], 567, false]
+            },
+            "set_compr_3":  {
+                "set!": [1, "string", 567, false]
+            },
+            "set_compr_4":  {
+                "set!": [1]
+            },
+            "set_compr_5":  {
+                "set!": [2, 3, 4, 5]
+            },
+            "set_compr_6":  {
+                "set!": [1, 2]
+            },
+            "set_compr_7":  {
+                "set!": [1, 2, 3]
+            }
+ }"#,
+    )?;
+
+    assert_match(
+        eval_file(&[rego.to_owned()], None, None, "data.test")?,
+        expected,
+    );
+    Ok(())
+}

tests/interpreter/cases/compr/object.yaml

@@ -0,0 +1,67 @@
+# Copyright (c) Rego-Rs Authors.
+# Licensed under the Apache 2.0 license.
+
+cases:
+  - note: simple
+    data: {}
+    modules:
+      - |
+        package test
+
+        x = { k:1 | k = ["Hello", "world", 1][_] }
+    query: data.test
+    want_result:
+      x:
+        object!:
+          - key: "Hello"
+            value: 1
+          - key: "world"
+            value: 1
+          - key: 1
+            value: 1
+
+
+  - note: key-loop
+    data: {}
+    modules:
+      - |
+        package test
+
+        x = { ["Hello", "world", 1][_]:1 | true }
+    query: data.test
+    want_result:
+      x:
+        object!:
+          - key: "Hello"
+            value: 1
+          - key: "world"
+            value: 1
+          - key: 1
+            value: 1
+
+  - note: multiple-occurance-of-same-key-value-pair
+    data: {}
+    modules:
+      - |
+        package test
+        x = { k:v | k = ["Hello", "world", 1][_]; v = [1, 1][_] }
+    query: data.test
+    want_result:
+      x:
+        object!:
+          - key: "Hello"
+            value: 1
+          - key: "world"
+            value: 1
+          - key: 1
+            value: 1
+
+  - note: different-values-for-same-key
+    data: {}
+    modules:
+      - |
+        package test
+        x = { k:v | k = ["Hello", "world", 1][_]; v = [1, 2][_] }
+    query: data.test
+    error: "value for key `\"Hello\"` generated multiple times: `1` and `2`"
+    want_result: