From 616d9473f1b596d5b938c560b3a5efec18c7f047 Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Wed, 8 Jul 2020 16:00:32 -0400
Subject: [PATCH] topdown: Add new numbers.range built-in function

This commit adds a new built-in function to generate a range of
integers between two values (inclusive). This is useful in certain
cases where users need to enumerate a set of values (e.g., port
numbers).

Fixes #2479

Signed-off-by: Torin Sandall <torinsandall@gmail.com>
---
 ast/builtins.go                  | 23 ++++++++++++++
 docs/content/policy-reference.md |  1 +
 topdown/numbers.go               | 46 ++++++++++++++++++++++++++++
 topdown/numbers_test.go          | 51 ++++++++++++++++++++++++++++++++
 4 files changed, 121 insertions(+)
 create mode 100644 topdown/numbers.go
 create mode 100644 topdown/numbers_test.go

diff --git a/ast/builtins.go b/ast/builtins.go
index 292391bb5d..e4ce8e556f 100644
--- a/ast/builtins.go
+++ b/ast/builtins.go
@@ -120,6 +120,9 @@ var DefaultBuiltins = [...]*Builtin{
 	TrimSpace,
 	Sprintf,
 
+	// Numbers
+	NumbersRange,
+
 	// Encoding
 	JSONMarshal,
 	JSONUnmarshal,
@@ -984,6 +987,26 @@ var Sprintf = &Builtin{
 	),
 }
 
+/**
+ * Numbers
+ */
+
+// NumbersRange returns an array of numbers in the given inclusive range.
+var NumbersRange = &Builtin{
+	Name: "numbers.range",
+	Decl: types.NewFunction(
+		types.Args(
+			types.N,
+			types.N,
+		),
+		types.NewArray(nil, types.N),
+	),
+}
+
+/**
+ * Units
+ */
+
 // UnitsParseBytes converts strings like 10GB, 5K, 4mb, and the like into an
 // integer number of bytes.
 var UnitsParseBytes = &Builtin{
diff --git a/docs/content/policy-reference.md b/docs/content/policy-reference.md
index 627f5cc69e..39ea036615 100644
--- a/docs/content/policy-reference.md
+++ b/docs/content/policy-reference.md
@@ -273,6 +273,7 @@ complex types.
 | <span class="opa-keep-it-together">``z := x % y``</span>   | ``z`` is the remainder from the division of ``x`` and ``y``  |
 | <span class="opa-keep-it-together">``output := round(x)``</span>    | ``output`` is ``x`` rounded to the nearest integer |
 | <span class="opa-keep-it-together">``output := abs(x)``</span>    | ``output`` is the absolute value of ``x`` |
+| <span class="opa-keep-it-together">``output := numbers.range(a, b)``</span> | ``output`` is the range of integer numbers between ``a`` and ``b`` (inclusive). If ``a`` == ``b`` then ``output`` == ``[a]``. If ``a`` < ``b`` the range is in ascending order. If ``a`` > ``b`` the range is in descending order. |
 
 ### Aggregates
 
diff --git a/topdown/numbers.go b/topdown/numbers.go
new file mode 100644
index 0000000000..8a9ecfca45
--- /dev/null
+++ b/topdown/numbers.go
@@ -0,0 +1,46 @@
+// Copyright 2020 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package topdown
+
+import (
+	"math/big"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/topdown/builtins"
+)
+
+var one = big.NewInt(1)
+
+func builtinNumbersRange(_ BuiltinContext, operands []*ast.Term, iter func(*ast.Term) error) error {
+
+	x, err := builtins.BigIntOperand(operands[0].Value, 1)
+	if err != nil {
+		return err
+	}
+
+	y, err := builtins.BigIntOperand(operands[1].Value, 2)
+	if err != nil {
+		return err
+	}
+
+	var result ast.Array
+	cmp := x.Cmp(y)
+
+	if cmp <= 0 {
+		for i := new(big.Int).Set(x); i.Cmp(y) <= 0; i = i.Add(i, one) {
+			result = append(result, ast.NewTerm(builtins.IntToNumber(i)))
+		}
+	} else {
+		for i := new(big.Int).Set(x); i.Cmp(y) >= 0; i = i.Sub(i, one) {
+			result = append(result, ast.NewTerm(builtins.IntToNumber(i)))
+		}
+	}
+
+	return iter(ast.NewTerm(result))
+}
+
+func init() {
+	RegisterBuiltinFunc(ast.NumbersRange.Name, builtinNumbersRange)
+}
diff --git a/topdown/numbers_test.go b/topdown/numbers_test.go
new file mode 100644
index 0000000000..f03eac85af
--- /dev/null
+++ b/topdown/numbers_test.go
@@ -0,0 +1,51 @@
+// Copyright 2020 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+package topdown
+
+import (
+	"testing"
+)
+
+func TestBuiltinNumbersRange(t *testing.T) {
+	cases := []struct {
+		note string
+		stmt string
+		exp  interface{}
+	}{
+		{
+			note: "one",
+			stmt: "p = x { x := numbers.range(0, 0) }",
+			exp:  "[0]",
+		},
+		{
+			note: "ascending",
+			stmt: "p = x { x := numbers.range(-2, 3) }",
+			exp:  "[-2, -1, 0, 1, 2, 3]",
+		},
+		{
+			note: "descending",
+			stmt: "p = x { x := numbers.range(2, -3) }",
+			exp:  "[2, 1, 0, -1, -2, -3]",
+		},
+		{
+			note: "precision",
+			stmt: "p { numbers.range(49649733057, 49649733060, [49649733057, 49649733058, 49649733059, 49649733060]) }",
+			exp:  "true",
+		},
+		{
+			note: "error: floating-point number pos 1",
+			stmt: "p { numbers.range(3.14, 4) }",
+			exp:  &Error{Code: TypeErr, Message: "numbers.range: operand 1 must be integer number but got floating-point number"},
+		},
+		{
+			note: "error: floating-point number pos 2",
+			stmt: "p { numbers.range(3, 3.14) }",
+			exp:  &Error{Code: TypeErr, Message: "numbers.range: operand 2 must be integer number but got floating-point number"},
+		},
+	}
+
+	for _, tc := range cases {
+		runTopDownTestCase(t, map[string]interface{}{}, tc.note, []string{tc.stmt}, tc.exp)
+	}
+}