Add array.slice built-in function

Fixes open-policy-agent#1243

Signed-off-by: Ravi Soni <[email protected]>

ast/builtins.go

@@ -65,6 +65,7 @@ var DefaultBuiltins = [...]*Builtin{
 
 	// Arrays
 	ArrayConcat,
+	ArraySlice,
 
 	// Casting
 	ToNumber,
@@ -507,6 +508,19 @@ var ArrayConcat = &Builtin{
 	),
 }
 
+// ArraySlice returns a slice of a given array
+var ArraySlice = &Builtin{
+	Name: "array.slice",
+	Decl: types.NewFunction(
+		types.Args(
+			types.NewArray(nil, types.A),
+			types.NewNumber(),
+			types.NewNumber(),
+		),
+		types.NewArray(nil, types.A),
+	),
+}
+
 /**
  * Casting
  */

docs/content/docs/language-reference.md

@@ -51,10 +51,10 @@ complex types.
 | <span class="opa-keep-it-together">``output := any(array_or_set)``</span> | ``output`` is ``true`` if any of the values in ``array_or_set`` is ``true``. A collection of length 0 returns ``false``.|
 
 ### Arrays
-
 | Built-in | Description |
 | ------- |-------------|
 | <span class="opa-keep-it-together">``output := array.concat(array, array)``</span> | ``output`` is the result of concatenating the two input arrays together. |
+  <span class="opa-keep-it-together">``output := array.slice(array, startIndex, stopIndex)``</span> | ``output`` is the part of the ``array`` from ``startIndex`` to ``stopIndex`` including the first but excluding the last. If any of the provided indices are negative or ``stopIndex`` is less than ``startIndex`` then the original array is returned. |
 
 ### Sets
 

topdown/array.go

@@ -33,6 +33,36 @@ func builtinArrayConcat(a, b ast.Value) (ast.Value, error) {
 	return arrC, nil
 }
 
+func builtinArraySlice(a, i, j ast.Value) (ast.Value, error) {
+	arr, err := builtins.ArrayOperand(a, 1)
+	if err != nil {
+		return nil, err
+	}
+
+	startIndex, err := builtins.IntOperand(i, 2)
+	if err != nil {
+		return nil, err
+	}
+
+	stopIndex, err := builtins.IntOperand(j, 3)
+	if err != nil {
+		return nil, err
+	}
+
+	// If any of the provided indices are negative or stopIndex is less than startIndex
+	// then return a copy of the original array.
+	if startIndex < 0 || stopIndex < 0 || (stopIndex < startIndex) {
+		startIndex = 0
+		stopIndex = len(arr)
+	}
+
+	arrb := arr[startIndex:stopIndex]
+
+	return arrb, nil
+
+}
+
 func init() {
 	RegisterFunctionalBuiltin2(ast.ArrayConcat.Name, builtinArrayConcat)
+	RegisterFunctionalBuiltin3(ast.ArraySlice.Name, builtinArraySlice)
 }

topdown/array_test.go

@@ -19,6 +19,11 @@ func TestTopDownArray(t *testing.T) {
 		{"concat", []string{`p = x { x = array.concat([1,2], [3,4]) }`}, "[1,2,3,4]"},
 		{"concat: err", []string{`p = x { x = array.concat(data.b, [3,4]) }`}, fmt.Errorf("operand 1 must be array")},
 		{"concat: err rhs", []string{`p = x { x = array.concat([1,2], data.b) }`}, fmt.Errorf("operand 2 must be array")},
+
+		{"slice", []string{`p = x { x = array.slice([1,2,3,4,5], 1, 3) }`}, "[2,3]"},
+		{"slice: empty slice", []string{`p = x { x = array.slice([1,2,3], 0, 0) }`}, "[]"},
+		{"slice: negative indices", []string{`p = x { x = array.slice([1,2,3,4,5], -4, -1) }`}, "[1,2,3,4,5]"},
+		{"slice: stopIndex < startIndex", []string{`p = x { x = array.slice([1,2,3,4,5], 4, 1) }`}, "[1,2,3,4,5]"},
 	}
 
 	data := loadSmallTestData()