wasm: Set operations and builtins.

This is for #1114.

Signed-off-by: Teemu Koponen <koponen@styra.com>

internal/planner/planner.go

@@ -24,13 +24,18 @@ type wasmBuiltin struct {
 
 // internalBuiltins are the built-in functions implemented in wasm.
 var internalBuiltins = map[string]wasmBuiltin{
-	ast.Plus.Name:     wasmBuiltin{ast.Plus, "opa_arith_plus"},
-	ast.Minus.Name:    wasmBuiltin{ast.Minus, "opa_arith_minus"},
-	ast.Multiply.Name: wasmBuiltin{ast.Multiply, "opa_arith_multiply"},
-	ast.Divide.Name:   wasmBuiltin{ast.Divide, "opa_arith_divide"},
-	ast.Abs.Name:      wasmBuiltin{ast.Abs, "opa_arith_abs"},
-	ast.Round.Name:    wasmBuiltin{ast.Abs, "opa_arith_round"},
-	ast.Rem.Name:      wasmBuiltin{ast.Rem, "opa_arith_rem"},
+	ast.Plus.Name:         wasmBuiltin{ast.Plus, "opa_arith_plus"},
+	ast.Minus.Name:        wasmBuiltin{ast.Minus, "opa_arith_minus"},
+	ast.Multiply.Name:     wasmBuiltin{ast.Multiply, "opa_arith_multiply"},
+	ast.Divide.Name:       wasmBuiltin{ast.Divide, "opa_arith_divide"},
+	ast.Abs.Name:          wasmBuiltin{ast.Abs, "opa_arith_abs"},
+	ast.Round.Name:        wasmBuiltin{ast.Abs, "opa_arith_round"},
+	ast.Rem.Name:          wasmBuiltin{ast.Rem, "opa_arith_rem"},
+	ast.SetDiff.Name:      wasmBuiltin{ast.SetDiff, "opa_set_diff"},
+	ast.And.Name:          wasmBuiltin{ast.And, "opa_set_intersection"},
+	ast.Or.Name:           wasmBuiltin{ast.Or, "opa_set_union"},
+	ast.Intersection.Name: wasmBuiltin{ast.Intersection, "opa_sets_intersection"},
+	ast.Union.Name:        wasmBuiltin{ast.Union, "opa_sets_union"},
 }
 
 // Planner implements a query planner for Rego queries.

test/wasm/assets/018_builtins.yaml

@@ -20,6 +20,21 @@ cases:
   - note: remainder built-in
     query: rem(4,3,x)
     want_result: [{'x': 1}]
+  - note: set diff built-in
+    query: set_diff({0,1},{0},x)
+    want_result: [{'x': [1]}]
+  - note: and built-in
+    query: and({0,1},{0},x)
+    want_result: [{'x': [0]}]
+  - note: or built-in
+    query: or({0,1},{0},x)
+    want_result: [{'x': [1,0]}]
+  - note: intersection built-in
+    query: intersection({{0,1},{0}},x)
+    want_result: [{'x': [0]}]
+  - note: union built-in
+    query: union({{0,1},{0}},x)
+    want_result: [{'x': [1,0]}]
   - note: custom built-in
     query: x = custom_builtin_test(100)
     want_result: [{'x': 101}]

wasm/src/set.c

@@ -1,4 +1,4 @@
-#include "value.h"
+#include "set.h"
 
 opa_value *opa_set_diff(opa_value *a, opa_value *b)
 {
@@ -27,3 +27,157 @@ opa_value *opa_set_diff(opa_value *a, opa_value *b)
 
 	return &r->hdr;
 }
+
+opa_value *opa_set_intersection(opa_value *a, opa_value *b)
+{
+    if (opa_value_type(a) != OPA_SET || opa_value_type(b) != OPA_SET)
+    {
+        return NULL;
+    }
+
+    opa_set_t *x = opa_cast_set(a);
+    opa_set_t *y = opa_cast_set(b);
+    opa_set_t *r = opa_cast_set(opa_set_with_cap(x->len < y->len ? x->len : y->len));
+
+    if (y->len < x->len)
+    {
+        x = opa_cast_set(b);
+        y = opa_cast_set(a);
+    }
+
+    for (int i = 0; i < x->n; i++)
+    {
+        opa_set_elem_t *elem = x->buckets[i];
+
+        while (elem != NULL)
+        {
+            if (opa_set_get(y, elem->v) != NULL)
+            {
+                opa_set_add(r, elem->v);
+            }
+            elem = elem->next;
+        }
+    }
+
+    return &r->hdr;
+}
+
+opa_value *opa_sets_intersection(opa_value *v)
+{
+    if (opa_value_type(v) != OPA_SET)
+    {
+        return NULL;
+    }
+
+    opa_set_t *s = opa_cast_set(v);
+
+    if (s->len == 0)
+    {
+        return opa_set();
+    }
+
+    opa_value *r = NULL;
+
+    for (int i = 0; i < s->n; i++)
+    {
+        opa_set_elem_t *elem = s->buckets[i];
+
+        while (elem != NULL)
+        {
+            if (opa_value_type(elem->v) != OPA_SET)
+            {
+                return NULL;
+            }
+
+            if (r == NULL)
+            {
+                r = opa_set_union(opa_set(), elem->v);
+            } else {
+                opa_value *x = opa_set_intersection(r, elem->v);
+                opa_value_free(r);
+                if (x == NULL)
+                {
+                    return NULL;
+                }
+
+                r = x;
+            }
+
+            elem = elem->next;
+        }
+    }
+
+    return r;
+}
+
+opa_value *opa_set_union(opa_value *a, opa_value *b)
+{
+    if (opa_value_type(a) != OPA_SET || opa_value_type(b) != OPA_SET)
+    {
+        return NULL;
+    }
+
+    opa_set_t *x = opa_cast_set(a);
+    opa_set_t *y = opa_cast_set(b);
+    opa_set_t *r = opa_cast_set(opa_set());
+
+    for (int i = 0; i < x->n; i++)
+    {
+        opa_set_elem_t *elem = x->buckets[i];
+
+        while (elem != NULL)
+        {
+            opa_set_add(r, elem->v);
+            elem = elem->next;
+        }
+    }
+
+    for (int i = 0; i < y->n; i++)
+    {
+        opa_set_elem_t *elem = y->buckets[i];
+
+        while (elem != NULL)
+        {
+            opa_set_add(r, elem->v);
+            elem = elem->next;
+        }
+    }
+
+	return &r->hdr;
+}
+
+opa_value *opa_sets_union(opa_value *v)
+{
+    if (opa_value_type(v) != OPA_SET)
+    {
+        return NULL;
+    }
+
+    opa_set_t *s = opa_cast_set(v);
+    opa_value *r = opa_set();
+
+    for (int i = 0; i < s->n; i++)
+    {
+        opa_set_elem_t *elem = s->buckets[i];
+
+        while (elem != NULL)
+        {
+            if (opa_value_type(elem->v) != OPA_SET)
+            {
+                return NULL;
+            }
+
+            opa_value *x = opa_set_union(r, elem->v);
+            opa_value_free(r);
+            if (x == NULL)
+            {
+                return NULL;
+            }
+
+            r = x;
+            elem = elem->next;
+        }
+    }
+
+    return r;
+}

wasm/src/set.h

@@ -1,6 +1,13 @@
 #ifndef OPA_SET_H
 #define OPA_SET_H
 
+#include "value.h"
+
 opa_value *opa_set_diff(opa_value *a, opa_value *b);
+opa_value *opa_set_intersection(opa_value *a, opa_value *b);
+opa_value *opa_set_union(opa_value *a, opa_value *b);
+
+opa_value *opa_sets_intersection(opa_value *v);
+opa_value *opa_sets_union(opa_value *v);
 
 #endif

wasm/src/value.c

@@ -1035,6 +1035,18 @@ opa_value *opa_set()
     return __opa_set_with_buckets(OPA_SET_MIN_BUCKETS);
 }
 
+opa_value *opa_set_with_cap(size_t n)
+{
+    size_t buckets = OPA_SET_MIN_BUCKETS;
+
+    while (n > (buckets * OPA_SET_LOAD_FACTOR))
+    {
+        buckets *= 2;
+    }
+
+    return __opa_set_with_buckets(buckets);
+}
+
 void opa_value_boolean_set(opa_value *v, int b)
 {
     opa_boolean_t *ret = opa_cast_boolean(v);

wasm/src/value.h

@@ -134,8 +134,8 @@ opa_value *opa_array();
 opa_value *opa_array_with_cap(size_t cap);
 opa_value *opa_array_with_elems(opa_array_elem_t *elems, size_t len, size_t cap);
 opa_value *opa_object();
-opa_value *opa_object_with_cap(size_t cap);
 opa_value *opa_set();
+opa_value *opa_set_with_cap(size_t cap);
 
 void opa_value_boolean_set(opa_value *v, int b);
 void opa_value_number_set_int(opa_value *v, long long i);

wasm/tests/test.c

@@ -2,6 +2,7 @@
 #include "json.h"
 #include "malloc.h"
 #include "arithmetic.h"
+#include "set.h"
 
 void opa_test_fail(const char *note, const char *func, const char *file, int line);
 void opa_test_pass(const char *note, const char *func);
@@ -970,3 +971,59 @@ void test_arithmetic(void)
     test("divide 3/2", opa_number_as_float(opa_cast_number(opa_arith_divide(opa_number_float(3), opa_number_float(2)))) == 1.5);
     test("remainder 5 % 2", opa_number_as_float(opa_cast_number(opa_arith_rem(opa_number_float(5), opa_number_float(2)))) == 1);
 }
+
+void test_set_diff(void)
+{
+    // test_arithmetic covers the diff.
+}
+
+void test_set_intersection_union(void)
+{
+    opa_set_t *s1 = opa_cast_set(opa_set());
+    opa_set_add(s1, opa_number_int(0));
+    opa_set_add(s1, opa_number_int(1));
+    opa_set_add(s1, opa_number_int(2));
+
+    opa_set_t *s2 = opa_cast_set(opa_set());
+    opa_set_add(s2, opa_number_int(0));
+    opa_set_add(s2, opa_number_int(1));
+
+    opa_set_t *r = opa_cast_set(opa_set_intersection(&s1->hdr, &s2->hdr));
+    test("set/intersection", r->len == 2 && opa_set_get(r, opa_number_int(0)) != NULL && opa_set_get(r, opa_number_int(1)) != NULL);
+
+    r = opa_cast_set(opa_set_union(&s1->hdr, &s2->hdr));
+    test("set/union", r->len == 3 &&
+         opa_set_get(r, opa_number_int(0)) != NULL &&
+         opa_set_get(r, opa_number_int(1)) != NULL &&
+         opa_set_get(r, opa_number_int(2)) != NULL);
+}
+
+
+void test_sets_intersection_union(void)
+{
+    opa_set_t *s1 = opa_cast_set(opa_set());
+    opa_set_add(s1, opa_number_int(0));
+    opa_set_add(s1, opa_number_int(1));
+    opa_set_add(s1, opa_number_int(2));
+
+    opa_set_t *s2 = opa_cast_set(opa_set());
+    opa_set_add(s2, opa_number_int(0));
+    opa_set_add(s2, opa_number_int(1));
+
+    opa_set_t *s3 = opa_cast_set(opa_set());
+    opa_set_add(s3, opa_number_int(0));
+
+    opa_set_t *sets = opa_cast_set(opa_set());
+    opa_set_add(sets, &s1->hdr);
+    opa_set_add(sets, &s2->hdr);
+    opa_set_add(sets, &s3->hdr);
+
+    opa_set_t *r = opa_cast_set(opa_sets_intersection(&sets->hdr));
+    test("sets/intersection", r->len == 1 && opa_set_get(r, opa_number_int(0)) != NULL);
+
+    r = opa_cast_set(opa_sets_union(&sets->hdr));
+    test("sets/union", r->len == 3 &&
+         opa_set_get(r, opa_number_int(0)) != NULL &&
+         opa_set_get(r, opa_number_int(1)) != NULL &&
+         opa_set_get(r, opa_number_int(2)) != NULL);
+}