Add array.flatten built-in function (#8232)

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 <anders.eknert@apple.com>

.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 \

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

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": [
       {

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": {

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",

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 `[]`.",

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

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

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)
 }

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,
 	}
 }
 

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

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)
+		}
+	}
+
+}

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)
 	}
 }
 

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)
 {

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);
 

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)