From d223284b453318c5a02f064eac1c0396aae453fc Mon Sep 17 00:00:00 2001
From: Ravi Soni <soniravi829@gmail.com>
Date: Tue, 26 Mar 2019 17:14:03 +0530
Subject: [PATCH] Add array.slice builtin function

Fixes #1243

Signed-off-by: Ravi Soni <soniravi829@gmail.com>
---
 ast/builtins.go       | 14 ++++++++++++++
 topdown/array.go      | 39 +++++++++++++++++++++++++++++++++++++++
 topdown/array_test.go |  5 +++++
 3 files changed, 58 insertions(+)

diff --git a/ast/builtins.go b/ast/builtins.go
index 7c373014a8..9a201f1f54 100644
--- a/ast/builtins.go
+++ b/ast/builtins.go
@@ -65,6 +65,7 @@ var DefaultBuiltins = [...]*Builtin{
 
 	// Arrays
 	ArrayConcat,
+	ArraySlice,
 
 	// Casting
 	ToNumber,
@@ -505,6 +506,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
  */
diff --git a/topdown/array.go b/topdown/array.go
index 8c9c02b516..ab0908256d 100644
--- a/topdown/array.go
+++ b/topdown/array.go
@@ -5,6 +5,8 @@
 package topdown
 
 import (
+	"fmt"
+
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/topdown/builtins"
 )
@@ -33,6 +35,43 @@ func builtinArrayConcat(a, b ast.Value) (ast.Value, error) {
 	return arrC, nil
 }
 
+func builtinArraySlice(a, i, j ast.Value) (ast.Value, error) {
+	arrA, err := builtins.ArrayOperand(a, 1)
+	if err != nil {
+		return nil, err
+	}
+
+	startIndex, err := builtins.IntOperand(i, 1)
+	if err != nil {
+		return nil, err
+	}
+
+	stopIndex, err := builtins.IntOperand(j, 2)
+	if err != nil {
+		return nil, err
+	}
+
+	// Don't allow negative indices for slicing
+	if startIndex < 0 || stopIndex < 0 {
+		return nil, fmt.Errorf("Invalid slicing operation: negative indices not allowed")
+	}
+
+	// stopIndex can't be less than startIndex
+	if stopIndex < startIndex {
+		return nil, fmt.Errorf("Invalid slicing operation: stopIndex can't be less than startIndex")
+	}
+
+	arrB := make(ast.Array, 0, stopIndex-startIndex)
+
+	for i := startIndex; i < stopIndex; i++ {
+		arrB = append(arrB, arrA[i])
+	}
+
+	return arrB, nil
+
+}
+
 func init() {
 	RegisterFunctionalBuiltin2(ast.ArrayConcat.Name, builtinArrayConcat)
+	RegisterFunctionalBuiltin3(ast.ArraySlice.Name, builtinArraySlice)
 }
diff --git a/topdown/array_test.go b/topdown/array_test.go
index 3d9e5c12d4..d783dfff95 100644
--- a/topdown/array_test.go
+++ b/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", []string{`p = x { x = array.slice([1,2,3], 0, 0) }`}, "[]"},
+		{"slice: err negative indices", []string{`p = x { x = array.slice([3,4], -1, 0) }`}, fmt.Errorf("Invalid slicing operation: negative indices not allowed")},
+		{"slice: err stopIndex < startIndex", []string{`p = x { x = array.slice([1,2,3], 2, 1) }`}, fmt.Errorf("Invalid slicing operation: stopIndex can't be less than startIndex")},
 	}
 
 	data := loadSmallTestData()