From d0350b326e541b0d08ab00545e3ebd870495d1af Mon Sep 17 00:00:00 2001 From: Anders Eknert Date: Tue, 27 Jan 2026 22:46:11 +0100 Subject: [PATCH] Add `array.flatten` built-in function (#8232) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Originally meant to be `array.concat_n`, but this name is better as the behavior of this function differs from `array.concat` — namely that `array.flatten` accepts any type of valued in the input array. Only arrays are however flattened, and the rest are appended directly to the flattened output. Note that this function only flattens at the topmost level of the input array — not recursively! A cursory look at a few other languages suggest a single level is the common case. But if others feel we should flstten more, I'm happy to make an update. The C code for a Wasm implementstion here is cowboy coded, and I did not manage to run the tests on my machine due to some `docker` <-> `container` differences. I mostly just imitated the existing code in the array category. I doubt it'll work on the first try, but only CI can judge me. Also: - Remove `opa fmt` step from the Rego CI step, as this is done by Regal anyway a little later in the list of tasks. - Replace some hard-coded `docker` names in the `Makefile` with `$(DOCKER)` - Added name of built-in function missing to the unsupportedBuiltinErr error, as it has happened a few times now that I've used `:=` in a query, and had no clue what built-in it referred to. Fixes #8226 Signed-off-by: Anders Eknert --- .github/workflows/pull-request.yaml | 3 - Makefile | 8 +-- builtin_metadata.json | 21 +++++++ capabilities.json | 20 +++++++ internal/compiler/wasm/wasm.go | 1 + v1/ast/builtins.go | 13 +++++ v1/test/cases/testdata/v0/array/flatten.yaml | 58 ++++++++++++++++++++ v1/test/cases/testdata/v1/array/flatten.yaml | 58 ++++++++++++++++++++ v1/topdown/array.go | 35 ++++++++++++ v1/topdown/errors.go | 4 +- v1/topdown/eval.go | 2 +- v1/topdown/topdown_bench_test.go | 41 ++++++++++++++ v1/topdown/topdown_test.go | 32 +++++++---- wasm/src/array.c | 39 +++++++++++++ wasm/src/array.h | 1 + wasm/tests/test.c | 22 ++++++++ 16 files changed, 336 insertions(+), 22 deletions(-) create mode 100644 v1/test/cases/testdata/v0/array/flatten.yaml create mode 100644 v1/test/cases/testdata/v1/array/flatten.yaml diff --git a/.github/workflows/pull-request.yaml b/.github/workflows/pull-request.yaml index 3b5e5ab42a..39bfd3a2cb 100644 --- a/.github/workflows/pull-request.yaml +++ b/.github/workflows/pull-request.yaml @@ -443,9 +443,6 @@ jobs: - name: Test policies run: opa test --schema build/policy/schema --bundle build/policy - - name: Ensure proper formatting - run: opa fmt --list --fail build/policy - - name: Run file policy checks on changed files run: | curl --silent --fail --header 'Authorization: Bearer ${{ secrets.GITHUB_TOKEN }}' -o files.json \ diff --git a/Makefile b/Makefile index 6213eec8b9..1eb02576ea 100644 --- a/Makefile +++ b/Makefile @@ -160,7 +160,7 @@ wasm-sdk-e2e-test: generate .PHONY: check check: ifeq ($(DOCKER_RUNNING), 1) - docker run --rm -v $(shell pwd):/app:ro,Z -w /app golangci/golangci-lint:${GOLANGCI_LINT_VERSION} golangci-lint run -v + $(DOCKER) run --rm -v $(shell pwd):/app:ro,Z -w /app golangci/golangci-lint:${GOLANGCI_LINT_VERSION} golangci-lint run -v else @echo "Docker not installed or running. Skipping golangci run." endif @@ -168,7 +168,7 @@ endif .PHONY: fmt fmt: ifeq ($(DOCKER_RUNNING), 1) - docker run --rm -v $(shell pwd):/app:Z -w /app golangci/golangci-lint:${GOLANGCI_LINT_VERSION} golangci-lint run -v --fix + $(DOCKER) run --rm -v $(shell pwd):/app:Z -w /app golangci/golangci-lint:${GOLANGCI_LINT_VERSION} golangci-lint run -v --fix else @echo "Docker not installed or running. Skipping golangci run." endif @@ -470,7 +470,7 @@ check-fuzz: fuzz # not be able to use any module cache. .PHONY: check-go-module check-go-module: - docker run \ + $(DOCKER) run \ $(DOCKER_FLAGS) \ -w /src \ -v $(PWD):/src:Z \ @@ -482,7 +482,7 @@ check-go-module: .PHONY: check-yaml-tests check-yaml-tests: ifeq ($(DOCKER_RUNNING), 1) - docker run --rm -v $(shell pwd):/data:ro,Z -w /data pipelinecomponents/yamllint:${YAML_LINT_VERSION} yamllint -f $(YAML_LINT_FORMAT) v1/test/cases/testdata + $(DOCKER) run --rm -v $(shell pwd):/data:ro,Z -w /data pipelinecomponents/yamllint:${YAML_LINT_VERSION} yamllint -f $(YAML_LINT_FORMAT) v1/test/cases/testdata else @echo "Docker not installed or running. Skipping yamllint run." endif diff --git a/builtin_metadata.json b/builtin_metadata.json index c0aadffc60..63d29908fb 100644 --- a/builtin_metadata.json +++ b/builtin_metadata.json @@ -10,6 +10,7 @@ ], "array": [ "array.concat", + "array.flatten", "array.reverse", "array.slice" ], @@ -996,6 +997,26 @@ }, "wasm": true }, + "array.flatten": { + "args": [ + { + "description": "the array to be flattened", + "name": "arr", + "type": "array[any]" + } + ], + "available": [ + "edge" + ], + "description": "Non-recursively unpacks array items in arr into the flattened array. Other types are appended as-is.", + "introduced": "edge", + "result": { + "description": "array flattened one level", + "name": "flattened", + "type": "array[any]" + }, + "wasm": true + }, "array.reverse": { "args": [ { diff --git a/capabilities.json b/capabilities.json index e4bf21abf3..9eb82c2968 100644 --- a/capabilities.json +++ b/capabilities.json @@ -125,6 +125,26 @@ "type": "function" } }, + { + "name": "array.flatten", + "decl": { + "args": [ + { + "dynamic": { + "type": "any" + }, + "type": "array" + } + ], + "result": { + "dynamic": { + "type": "any" + }, + "type": "array" + }, + "type": "function" + } + }, { "name": "array.reverse", "decl": { diff --git a/internal/compiler/wasm/wasm.go b/internal/compiler/wasm/wasm.go index 6b38b12eee..92a117112a 100644 --- a/internal/compiler/wasm/wasm.go +++ b/internal/compiler/wasm/wasm.go @@ -90,6 +90,7 @@ var builtinsFunctions = map[string]string{ ast.Floor.Name: "opa_arith_floor", ast.Rem.Name: "opa_arith_rem", ast.ArrayConcat.Name: "opa_array_concat", + ast.ArrayFlatten.Name: "opa_array_flatten", ast.ArrayReverse.Name: "opa_array_reverse", ast.ArraySlice.Name: "opa_array_slice", ast.SetDiff.Name: "opa_set_diff", diff --git a/v1/ast/builtins.go b/v1/ast/builtins.go index dffc2ace23..39921d7055 100644 --- a/v1/ast/builtins.go +++ b/v1/ast/builtins.go @@ -95,6 +95,7 @@ var DefaultBuiltins = [...]*Builtin{ // Arrays ArrayConcat, + ArrayFlatten, ArraySlice, ArrayReverse, @@ -893,6 +894,18 @@ var ArrayConcat = &Builtin{ CanSkipBctx: true, } +var ArrayFlatten = &Builtin{ + Name: "array.flatten", + Description: "Non-recursively unpacks array items in arr into the flattened array. Other types are appended as-is.", + Decl: types.NewFunction( + types.Args( + types.Named("arr", types.NewArray(nil, types.A)).Description("the array to be flattened"), + ), + types.Named("flattened", types.NewArray(nil, types.A)).Description("array flattened one level"), + ), + CanSkipBctx: true, +} + var ArraySlice = &Builtin{ Name: "array.slice", Description: "Returns a slice of a given array. If `start` is greater or equal than `stop`, `slice` is `[]`.", diff --git a/v1/test/cases/testdata/v0/array/flatten.yaml b/v1/test/cases/testdata/v0/array/flatten.yaml new file mode 100644 index 0000000000..8d267907b8 --- /dev/null +++ b/v1/test/cases/testdata/v0/array/flatten.yaml @@ -0,0 +1,58 @@ +--- +cases: + - note: array/flatten empty array + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: [] + want_result: + - x: [] + - note: array/flatten only arrays + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: [[1, 2], [3, 4], [5]] + want_result: + - x: [1, 2, 3, 4, 5] + - note: array/flatten has nesting depth 1 + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: [[[1], [2, 3, []]], [[4, 5]]] + want_result: + - x: [[1], [2, 3, []], [4, 5]] + - note: array/flatten mixed types + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: [[1, 2], "string", [3, { "a": "b" }], 4.5, [6, [7, 8]]] + want_result: + - x: [1, 2, "string", 3, { "a": "b" }, 4.5, 6, [7, 8]] + - note: array/flatten type error + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: 42 + strict_error: true + want_error: "array.flatten: operand 1 must be array but got number" + want_error_code: eval_type_error diff --git a/v1/test/cases/testdata/v1/array/flatten.yaml b/v1/test/cases/testdata/v1/array/flatten.yaml new file mode 100644 index 0000000000..8d267907b8 --- /dev/null +++ b/v1/test/cases/testdata/v1/array/flatten.yaml @@ -0,0 +1,58 @@ +--- +cases: + - note: array/flatten empty array + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: [] + want_result: + - x: [] + - note: array/flatten only arrays + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: [[1, 2], [3, 4], [5]] + want_result: + - x: [1, 2, 3, 4, 5] + - note: array/flatten has nesting depth 1 + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: [[[1], [2, 3, []]], [[4, 5]]] + want_result: + - x: [[1], [2, 3, []], [4, 5]] + - note: array/flatten mixed types + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: [[1, 2], "string", [3, { "a": "b" }], 4.5, [6, [7, 8]]] + want_result: + - x: [1, 2, "string", 3, { "a": "b" }, 4.5, 6, [7, 8]] + - note: array/flatten type error + query: data.test.p = x + modules: + - | + package test + + p := array.flatten(data.foo) + data: + foo: 42 + strict_error: true + want_error: "array.flatten: operand 1 must be array but got number" + want_error_code: eval_type_error diff --git a/v1/topdown/array.go b/v1/topdown/array.go index 3541669016..64e5006028 100644 --- a/v1/topdown/array.go +++ b/v1/topdown/array.go @@ -43,6 +43,40 @@ func builtinArrayConcat(_ BuiltinContext, operands []*ast.Term, iter func(*ast.T return iter(ast.ArrayTerm(arrC...)) } +func builtinArrayFlatten(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error { + arr, err := builtins.ArrayOperand(operands[0].Value, 1) + if err != nil { + return err + } + + size := arr.Len() + preAlloc := size + + for i := range size { + if nested, ok := arr.Elem(i).Value.(*ast.Array); ok { + preAlloc += nested.Len() - 1 + } + } + + if size == preAlloc { + return iter(operands[0]) // Empty array, or no nested arrays -> nothing to flatten. + } + + flattened := make([]*ast.Term, 0, preAlloc) + for i := range size { + elem := arr.Elem(i) + if nested, ok := elem.Value.(*ast.Array); ok { + for j := range nested.Len() { + flattened = append(flattened, nested.Elem(j)) + } + } else { + flattened = append(flattened, elem) + } + } + + return iter(ast.ArrayTerm(flattened...)) +} + func builtinArraySlice(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error { arr, err := builtins.ArrayOperand(operands[0].Value, 1) if err != nil { @@ -108,6 +142,7 @@ func builtinArrayReverse(_ BuiltinContext, operands []*ast.Term, iter func(*ast. func init() { RegisterBuiltinFunc(ast.ArrayConcat.Name, builtinArrayConcat) + RegisterBuiltinFunc(ast.ArrayFlatten.Name, builtinArrayFlatten) RegisterBuiltinFunc(ast.ArraySlice.Name, builtinArraySlice) RegisterBuiltinFunc(ast.ArrayReverse.Name, builtinArrayReverse) } diff --git a/v1/topdown/errors.go b/v1/topdown/errors.go index 51a294c190..e80339e312 100644 --- a/v1/topdown/errors.go +++ b/v1/topdown/errors.go @@ -138,11 +138,11 @@ func objectDocKeyConflictErr(loc *ast.Location) error { } } -func unsupportedBuiltinErr(loc *ast.Location) error { +func unsupportedBuiltinErr(loc *ast.Location, name string) error { return &Error{ Code: InternalErr, Location: loc, - Message: "unsupported built-in", + Message: "unsupported built-in: " + name, } } diff --git a/v1/topdown/eval.go b/v1/topdown/eval.go index 3d74f9eb8b..6d6b79bea9 100644 --- a/v1/topdown/eval.go +++ b/v1/topdown/eval.go @@ -969,7 +969,7 @@ func (e *eval) evalCall(terms []*ast.Term, iter unifyIterator) error { builtinName := ref.String() bi, f, ok := e.builtinFunc(builtinName) if !ok { - return unsupportedBuiltinErr(e.query[e.index].Location) + return unsupportedBuiltinErr(e.query[e.index].Location, builtinName) } if mocked { // value replacement of built-in call diff --git a/v1/topdown/topdown_bench_test.go b/v1/topdown/topdown_bench_test.go index 0818d3d6be..b530660fe1 100644 --- a/v1/topdown/topdown_bench_test.go +++ b/v1/topdown/topdown_bench_test.go @@ -1032,3 +1032,44 @@ func BenchmarkTemplateStringVsConcatVsSprintf(b *testing.B) { }) } } + +// array.flatten-16 969697 1269 ns/op 2186 B/op 34 allocs/op +// array.concat-16 462535 2293 ns/op 3082 B/op 58 allocs/op +func BenchmarkArrayFlattenWithAndWithoutBuiltin(b *testing.B) { + b.Run("array.flatten", func(b *testing.B) { + runFlattenBenchmark(b, `x = array.flatten( + [[1,2,3], [4,5,6], [7,8,9], 10, [11,12]] + ) == [1,2,3,4,5,6,7,8,9,10,11,12]`) + }) + + b.Run("array.concat", func(b *testing.B) { + runFlattenBenchmark(b, `x = array.concat( + array.concat([1,2,3], [4,5,6]), + array.concat([7,8,9], array.concat([10], [11,12])) + ) == [1,2,3,4,5,6,7,8,9,10,11,12]`) + }) +} + +func runFlattenBenchmark(b *testing.B, query string) { + b.Helper() + + cqr, err := ast.NewCompiler().QueryCompiler().Compile(ast.MustParseBody(query)) + if err != nil { + b.Fatal(err) + } + q := NewQuery(cqr) + + for b.Loop() { + rs, err := q.Run(b.Context()) + if err != nil { + b.Fatalf("Unexpected topdown query error: %v", err) + } + if len(rs) != 1 { + b.Fatalf("Expected one result, got: %v", rs) + } + if x := rs[0][ast.Var("x")]; !ast.Boolean(true).Equal(x.Value) { + b.Fatalf("Expected true result, got: %v", x) + } + } + +} diff --git a/v1/topdown/topdown_test.go b/v1/topdown/topdown_test.go index 6deb9f9e54..f59efc8bc1 100644 --- a/v1/topdown/topdown_test.go +++ b/v1/topdown/topdown_test.go @@ -167,23 +167,31 @@ func TestTopDownWithKeyword(t *testing.T) { // Warning(philipc): This test modifies package variables in the ast package, // which means it cannot be run in parallel with other tests. func TestTopDownUnsupportedBuiltin(t *testing.T) { - ast.RegisterBuiltin(&ast.Builtin{ - Name: "unsupported_builtin", + ast.RegisterBuiltin(&ast.Builtin{Name: "unsupported_builtin"}) + t.Cleanup(func() { + i := slices.IndexFunc(ast.Builtins, func(b *ast.Builtin) bool { + return b.Name == "unsupported_builtin" + }) + ast.Builtins = append(ast.Builtins[:i], ast.Builtins[i+1:]...) + + delete(ast.BuiltinMap, "unsupported_builtin") }) - body := ast.MustParseBody(`unsupported_builtin()`) - ctx := t.Context() - compiler := ast.NewCompiler() + query := ast.MustParseBody(`unsupported_builtin()`) store := inmem.New() - txn := storage.NewTransactionOrDie(ctx, store) - defer store.Abort(ctx, txn) - q := NewQuery(body).WithCompiler(compiler).WithStore(store).WithTransaction(txn) - _, err := q.Run(ctx) - expected := unsupportedBuiltinErr(body[0].Location) + err := storage.Txn(t.Context(), store, storage.TransactionParams{}, func(txn storage.Transaction) error { + _, err := NewQuery(query). + WithCompiler(ast.NewCompiler()). + WithStore(store). + WithTransaction(txn). + Run(t.Context()) - if err.Error() != expected.Error() { - t.Fatalf("Expected %v but got: %v", expected, err) + return err + }) + + if exp := unsupportedBuiltinErr(query[0].Location, "unsupported_builtin"); err.Error() != exp.Error() { + t.Fatalf("Expected %v but got: %v", exp, err) } } diff --git a/wasm/src/array.c b/wasm/src/array.c index 595fb988db..1dba73cf61 100644 --- a/wasm/src/array.c +++ b/wasm/src/array.c @@ -26,6 +26,45 @@ opa_value *opa_array_concat(opa_value *a, opa_value *b) return &r->hdr; } +OPA_BUILTIN +opa_value *opa_array_flatten(opa_value *a) +{ + if (opa_value_type(a) != OPA_ARRAY) + { + return NULL; + } + + opa_array_t *arr = opa_cast_array(a); + + if (arr->len == 0) { + return a; + } + + int n = arr->len; + + for (int i = 0; i < arr->len; i++) { + if (opa_value_type(arr->elems[i].v) == OPA_ARRAY) { + opa_array_t *subarr = opa_cast_array(arr->elems[i].v); + n += subarr->len - 1; + } + } + + opa_array_t *flat = opa_cast_array(opa_array_with_cap(n)); + + for (int i = 0; i < arr->len; i++) { + if (opa_value_type(arr->elems[i].v) == OPA_ARRAY) { + opa_array_t *subarr = opa_cast_array(arr->elems[i].v); + for (int j = 0; j < subarr->len; j++) { + opa_array_append(flat, subarr->elems[j].v); + } + } else { + opa_array_append(flat, arr->elems[i].v); + } + } + + return &flat->hdr; +} + OPA_BUILTIN opa_value *opa_array_slice(opa_value *a, opa_value *i, opa_value *j) { diff --git a/wasm/src/array.h b/wasm/src/array.h index e1bcdce904..c1ec73a18d 100644 --- a/wasm/src/array.h +++ b/wasm/src/array.h @@ -2,6 +2,7 @@ #define OPA_ARRAY_H opa_value *opa_array_concat(opa_value *a, opa_value *b); +opa_value *opa_array_flatten(opa_value *a); opa_value *opa_array_slice(opa_value *a, opa_value *i, opa_value *j); opa_value *opa_array_reverse(opa_value *a); diff --git a/wasm/tests/test.c b/wasm/tests/test.c index 52e807a9b4..e313a9326f 100644 --- a/wasm/tests/test.c +++ b/wasm/tests/test.c @@ -1738,6 +1738,28 @@ void test_array(void) opa_value_compare(r->elems[0].v, opa_number_int(2)) == 0 && opa_value_compare(r->elems[1].v, opa_number_int(1)) == 0 && opa_value_compare(r->elems[2].v, opa_number_int(0)) == 0); + + // array.flatten [[0,1], [2,3], 4] -> [0,1,2,3,4] + opa_array_t *sub1 = opa_cast_array(opa_array()); + opa_array_append(sub1, opa_number_int(0)); + opa_array_append(sub1, opa_number_int(1)); + + opa_array_t *sub2 = opa_cast_array(opa_array()); + opa_array_append(sub2, opa_number_int(2)); + opa_array_append(sub2, opa_number_int(3)); + + opa_array_t *nested = opa_cast_array(opa_array()); + opa_array_append(nested, &sub1->hdr); + opa_array_append(nested, &sub2->hdr); + opa_array_append(nested, opa_number_int(4)); + + r = opa_cast_array(opa_array_flatten(&nested->hdr)); + test("array_flatten", r->len == 5 && + opa_value_compare(r->elems[0].v, opa_number_int(0)) == 0 && + opa_value_compare(r->elems[1].v, opa_number_int(1)) == 0 && + opa_value_compare(r->elems[2].v, opa_number_int(2)) == 0 && + opa_value_compare(r->elems[3].v, opa_number_int(3)) == 0 && + opa_value_compare(r->elems[4].v, opa_number_int(4)) == 0); } WASM_EXPORT(test_types)