wasm: Add native support for object.remove and object.union builtin f…

…unction

Signed-off-by: Ashutosh Narkar <[email protected]>

docs/content/policy-reference.md

@@ -309,8 +309,8 @@ complex types.
 | Built-in | Description | Wasm Support |
 | ------- |-------------|---------------|
 | <span class="opa-keep-it-together">`value := object.get(object, key, default)`</span> | `value` is the value stored by the `object` at `key`. If no value is found, `default` is returned. | ✅ |
-| <span class="opa-keep-it-together">`output := object.remove(object, keys)`</span> | `output` is a new object which is the result of removing the specified `keys` from `object`. `keys` must be either an array, object, or set of keys. | ``SDK-dependent`` |
-| <span class="opa-keep-it-together">`output := object.union(objectA, objectB)`</span> | `output` is a new object which is the result of an asymmetric recursive union of two objects where conflicts are resolved by choosing the key from the right-hand object (`objectB`). For example: `object.union({"a": 1, "b": 2, "c": {"d": 3}}, {"a": 7, "c": {"d": 4, "e": 5}})` will result in `{"a": 7, "b": 2, "c": {"d": 4, "e": 5}}`  | ``SDK-dependent`` |
+| <span class="opa-keep-it-together">`output := object.remove(object, keys)`</span> | `output` is a new object which is the result of removing the specified `keys` from `object`. `keys` must be either an array, object, or set of keys. | ✅ |
+| <span class="opa-keep-it-together">`output := object.union(objectA, objectB)`</span> | `output` is a new object which is the result of an asymmetric recursive union of two objects where conflicts are resolved by choosing the key from the right-hand object (`objectB`). For example: `object.union({"a": 1, "b": 2, "c": {"d": 3}}, {"a": 7, "c": {"d": 4, "e": 5}})` will result in `{"a": 7, "b": 2, "c": {"d": 4, "e": 5}}`  | ✅ |
 | <span class="opa-keep-it-together">`filtered := object.filter(object, keys)`</span> | `filtered` is a new object with the remaining data from `object` with only keys specified in `keys` which is an array, object, or set of keys. For example: `object.filter({"a": {"b": "x", "c": "y"}, "d": "z"}, ["a"])` will result in `{"a": {"b": "x", "c": "y"}}`). | ✅ |
 | <span class="opa-keep-it-together">`filtered := json.filter(object, paths)`</span> | `filtered` is the remaining data from `object` with only keys specified in `paths` which is an array or set of JSON string paths. For example: `json.filter({"a": {"b": "x", "c": "y"}}, ["a/b"])` will result in `{"a": {"b": "x"}}`). Paths are not filtered in-order and are deduplicated before being evaluated. | ``SDK-dependent`` |
 | <span class="opa-keep-it-together">`output := json.remove(object, paths)`</span> | `output` is a new object which is the result of removing all keys specified in `paths` which is an array or set of JSON string paths. For example: `json.remove({"a": {"b": "x", "c": "y"}}, ["a/b"])` will result in `{"a": {"c": "y"}}`. Paths are not removed in-order and are deduplicated before being evaluated. | ``SDK-dependent`` |

internal/compiler/wasm/wasm.go

@@ -113,6 +113,8 @@ var builtinsFunctions = map[string]string{
 	ast.JSONUnmarshal.Name:              "opa_json_unmarshal",
 	ast.ObjectFilter.Name:               "builtin_object_filter",
 	ast.ObjectGet.Name:                  "builtin_object_get",
+	ast.ObjectRemove.Name:               "builtin_object_remove",
+	ast.ObjectUnion.Name:                "builtin_object_union",
 	ast.Concat.Name:                     "opa_strings_concat",
 	ast.FormatInt.Name:                  "opa_strings_format_int",
 	ast.IndexOf.Name:                    "opa_strings_indexof",

internal/wasm/sdk/test/e2e/exceptions.yaml

@@ -25,6 +25,17 @@
 "objectfilter/error invalid object param type nil input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
 "objectfilter/error invalid object param type number input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
 "objectfilter/error invalid object param type string input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectunion/error wrong lhs type input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectunion/error wrong rhs type input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid object param type array input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid object param type bool input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid object param type number input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid object param type string input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid object param type nil input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid key param type string input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid key param type boolean input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid key param type number input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
+"objectremove/error invalid key param type nil input": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
 "regexmatch/re_match: bad pattern err": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
 "replacen/bad pattern object operand/non-string key": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"
 "replacen/bad pattern object operand/non-string value": "expected error missing - https://github.com/open-policy-agent/opa/issues/2954"

test/wasm/assets/018_builtins.yaml

@@ -158,9 +158,21 @@ cases:
   - note: object.get built-in
     query: 'object.get({"a": "b"}, "c", "c", x)'
     want_result: [{'x': "c"}]
+  - note: object.remove built-in
+    query: 'object.remove({"x": 0, "y": 1}, ["y"], x)'
+    want_result: [{'x': {'x': 0}}]
+  - note: object.remove built-in
+    query: 'object.remove({"x": 0, "y": 1}, {"z"}, x)'
+    want_result: [{'x': {'x': 0, 'y': 1}}]
   - note: object.filter built-in
     query: 'object.filter({"x": 0, "y": 1}, ["x"], x)'
     want_result: [{'x': {'x': 0}}]
+  - note: object.union built-in
+    query: 'object.union({"a": 1}, {"b": 2}, x)'
+    want_result: [{'x': {'a': 1, 'b': 2}}]
+  - note: object.union built-in
+    query: 'object.union({"a": 1}, {"a": {"b": {"c": 1}}, "d": 7}, x)'
+    want_result: [{'x': {"a": {"b": {"c": 1}}, "d": 7}}]
   - note: concat built-in
     query: concat(",",["a","b"],x)
     want_result: [{'x': "a,b"}]

wasm/src/object.c

@@ -1,5 +1,76 @@
 #include "object.h"
 
+
+static opa_value *__merge(opa_value *a, opa_value *b);
+static opa_value *__merge_with_overwrite(opa_value *a, opa_value *b);
+static void __copy_object_elem(opa_object_t *result, opa_value *a, opa_value *b);
+
+opa_value *__merge(opa_value *a, opa_value *b)
+{
+
+    opa_object_t *merged = opa_cast_object(opa_object());
+    opa_object_t *obj = opa_cast_object(a);
+    opa_object_t *other = opa_cast_object(b);
+
+    for (opa_value *key = opa_value_iter(a, NULL); key != NULL;
+         key = opa_value_iter(a, key))
+    {
+
+        opa_object_elem_t *original = opa_object_get(obj, key);
+        opa_object_elem_t *elem = opa_object_get(other, key);
+
+        // The key didn't exist in other, keep the original value.
+        if (elem == NULL)
+        {
+            opa_object_insert(merged, key, original->v);
+            continue;
+        }
+
+        // The key exists in both, resolve the conflict.
+        opa_value *merged_value = __merge_with_overwrite(original->v, elem->v);
+        opa_object_insert(merged, key, merged_value);
+
+    }
+
+    // Copy in any values from other for keys that don't exist in obj.
+    __copy_object_elem(merged, a, b);
+
+    return &merged->hdr;
+}
+
+opa_value *__merge_with_overwrite(opa_value *a, opa_value *b)
+{
+    if (opa_value_type(a) != OPA_OBJECT || opa_value_type(b) != OPA_OBJECT)
+    {
+        // If we can't merge, stick with the right-hand value.
+        return b;
+    }
+
+    return __merge(a, b);
+}
+
+static void __copy_object_elem(opa_object_t *result, opa_value *a, opa_value *b)
+{
+    opa_object_t *obj = opa_cast_object(b);
+
+    for (int i = 0; i < obj->n; i++)
+    {
+        opa_object_elem_t *elem = obj->buckets[i];
+
+        while (elem != NULL)
+        {
+            opa_value *other = opa_value_get(a, elem->k);
+
+            if (other == NULL)
+            {
+                opa_object_insert(result, elem->k, elem->v);
+            }
+
+            elem = elem->next;
+        }
+    }
+}
+
 opa_value *builtin_object_filter(opa_value *obj, opa_value *keys)
 {
     if (opa_value_type(obj) != OPA_OBJECT)
@@ -54,3 +125,63 @@ opa_value *builtin_object_get(opa_value *obj, opa_value *key, opa_value *value)
 
     return value;
 }
+
+opa_value *builtin_object_remove(opa_value *obj, opa_value *keys)
+{
+    if (opa_value_type(obj) != OPA_OBJECT)
+    {
+        return NULL;
+    }
+
+    opa_set_t *keys_to_remove = opa_cast_set(opa_set());
+
+    for (opa_value *key = opa_value_iter(keys, NULL); key != NULL;
+         key = opa_value_iter(keys, key))
+    {
+        opa_value* k;
+        switch (opa_value_type(keys))
+        {
+        case OPA_OBJECT:
+        case OPA_SET:
+            k = key;
+            break;
+        case OPA_ARRAY:
+            k = opa_value_get(keys, key);
+        }
+        opa_set_add(keys_to_remove, k);
+    }
+
+    opa_object_t *r = opa_cast_object(opa_object());
+
+    for (opa_value *key = opa_value_iter(obj, NULL); key != NULL;
+         key = opa_value_iter(obj, key))
+    {
+        if (opa_set_get(keys_to_remove, key) == NULL)
+        {
+            opa_object_elem_t *elem = opa_object_get(opa_cast_object(obj), key);
+            if (elem != NULL)
+            {
+                opa_object_insert(r, key, elem->v);
+            }
+        }
+    }
+
+    return &r->hdr;
+}
+
+opa_value *builtin_object_union(opa_value *a, opa_value *b)
+{
+    if (opa_value_type(a) != OPA_OBJECT)
+    {
+        return NULL;
+    }
+
+    if (opa_value_type(b) != OPA_OBJECT)
+    {
+        return NULL;
+    }
+
+    opa_value *r = __merge(a, b);
+
+    return r;
+}

wasm/src/object.h

@@ -5,5 +5,7 @@
 
 opa_value *builtin_object_filter(opa_value *obj, opa_value *keys);
 opa_value *builtin_object_get(opa_value *obj, opa_value *key, opa_value *value);
+opa_value *builtin_object_remove(opa_value *obj, opa_value *keys);
+opa_value *builtin_object_union(opa_value *a, opa_value *b);
 
 #endif