From 3ec7f1515e2856667b2ec8a1025de4e00ca6d02c Mon Sep 17 00:00:00 2001 From: Teemu Koponen Date: Wed, 6 May 2020 11:30:49 -0700 Subject: [PATCH] wasm: Set operations and builtins. This is for #1114. Signed-off-by: Teemu Koponen --- internal/planner/planner.go | 19 ++-- test/wasm/assets/018_builtins.yaml | 15 +++ wasm/src/set.c | 156 ++++++++++++++++++++++++++++- wasm/src/set.h | 7 ++ wasm/src/value.c | 12 +++ wasm/src/value.h | 2 +- wasm/tests/test.c | 57 +++++++++++ 7 files changed, 259 insertions(+), 9 deletions(-) diff --git a/internal/planner/planner.go b/internal/planner/planner.go index 6e393b67db..2910eeedcd 100644 --- a/internal/planner/planner.go +++ b/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. diff --git a/test/wasm/assets/018_builtins.yaml b/test/wasm/assets/018_builtins.yaml index 7f99b3e412..e1e4754712 100644 --- a/test/wasm/assets/018_builtins.yaml +++ b/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}] diff --git a/wasm/src/set.c b/wasm/src/set.c index b7db58030b..3f15a98ca5 100644 --- a/wasm/src/set.c +++ b/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; +} diff --git a/wasm/src/set.h b/wasm/src/set.h index 61b039ca5e..2490c41d82 100644 --- a/wasm/src/set.h +++ b/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 diff --git a/wasm/src/value.c b/wasm/src/value.c index 598f5dfffb..db65cc62a0 100644 --- a/wasm/src/value.c +++ b/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); diff --git a/wasm/src/value.h b/wasm/src/value.h index b85451833f..4ea544a6ba 100644 --- a/wasm/src/value.h +++ b/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); diff --git a/wasm/tests/test.c b/wasm/tests/test.c index 3a90849540..5fd991e9e1 100644 --- a/wasm/tests/test.c +++ b/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); +}