From 1f94f392313608ea175eca0c81852afc413abf56 Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Tue, 14 Jun 2016 09:14:04 -0700
Subject: [PATCH] Add other simple arithmetic built-ins

---
 ast/builtins.go         | 30 ++++++++++++++-
 topdown/arithmetic.go   | 84 +++++++++++++++++++++++++++++++++--------
 topdown/builtins.go     |  6 ++-
 topdown/topdown_test.go |  4 ++
 4 files changed, 107 insertions(+), 17 deletions(-)

diff --git a/ast/builtins.go b/ast/builtins.go
index 8df45e985a..13a1174818 100644
--- a/ast/builtins.go
+++ b/ast/builtins.go
@@ -21,7 +21,7 @@ func RegisterBuiltin(b *Builtin) {
 var DefaultBuiltins = [...]*Builtin{
 	Equality,
 	GreaterThan, GreaterThanEq, LessThan, LessThanEq, NotEqual,
-	Plus,
+	Plus, Minus, Multiply, Divide, Round,
 	Count, Sum,
 	ToNumber,
 }
@@ -92,6 +92,34 @@ var Plus = &Builtin{
 	TargetPos: []int{2},
 }
 
+// Minus subtracts the second number from the first number.
+var Minus = &Builtin{
+	Name:      Var("minus"),
+	NumArgs:   3,
+	TargetPos: []int{2},
+}
+
+// Multiply multiplies two numbers together.
+var Multiply = &Builtin{
+	Name:      Var("mul"),
+	NumArgs:   3,
+	TargetPos: []int{2},
+}
+
+// Divide divides the first number by the second number.
+var Divide = &Builtin{
+	Name:      Var("div"),
+	NumArgs:   3,
+	TargetPos: []int{2},
+}
+
+// Round rounds the number up to the nearest integer.
+var Round = &Builtin{
+	Name:      Var("round"),
+	NumArgs:   2,
+	TargetPos: []int{1},
+}
+
 /**
  * Aggregates
  */
diff --git a/topdown/arithmetic.go b/topdown/arithmetic.go
index 53d5b394e9..40baa67be4 100644
--- a/topdown/arithmetic.go
+++ b/topdown/arithmetic.go
@@ -5,34 +5,88 @@
 package topdown
 
 import (
+	"fmt"
+	"math"
+
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/pkg/errors"
 )
 
-func evalPlus(ctx *Context, expr *ast.Expr, iter Iterator) error {
-	ops := expr.Terms.([]*ast.Term)
+type arithmeticFunc func(a, b float64) (ast.Number, error)
 
+func arithPlus(a, b float64) (ast.Number, error) {
+	return ast.Number(a + b), nil
+}
+
+func arithMinus(a, b float64) (ast.Number, error) {
+	return ast.Number(a - b), nil
+}
+
+func arithMultiply(a, b float64) (ast.Number, error) {
+	return ast.Number(a * b), nil
+}
+
+func arithDivide(a, b float64) (ast.Number, error) {
+	if b == 0 {
+		return 0, fmt.Errorf("divide: by zero")
+	}
+	return ast.Number(a / b), nil
+}
+
+func arithRound(a float64) (ast.Number, error) {
+	return ast.Number(math.Floor(a + 0.5)), nil
+}
+
+func evalRound(ctx *Context, expr *ast.Expr, iter Iterator) error {
+	ops := expr.Terms.([]*ast.Term)
 	a, err := ValueToFloat64(ops[1].Value, ctx)
 	if err != nil {
-		return errors.Wrapf(err, "plus")
+		return errors.Wrapf(err, "round")
 	}
-
-	b, err := ValueToFloat64(ops[2].Value, ctx)
-	if err != nil {
-		return errors.Wrapf(err, "plus")
-	}
-
-	c := ops[3].Value
-	r := ast.Number(a + b)
-
-	switch c := c.(type) {
+	r := ast.Number(math.Floor(a + 0.5))
+	b := ops[2].Value
+	switch b := b.(type) {
 	case ast.Var:
-		ctx = ctx.BindVar(c, r)
+		ctx = ctx.BindVar(b, r)
 		return iter(ctx)
 	default:
-		if r.Equal(c) {
+		if b.Equal(r) {
 			return iter(ctx)
 		}
 		return nil
 	}
 }
+
+func evalArithmetic(f arithmeticFunc) BuiltinFunc {
+	return func(ctx *Context, expr *ast.Expr, iter Iterator) error {
+		ops := expr.Terms.([]*ast.Term)
+
+		a, err := ValueToFloat64(ops[1].Value, ctx)
+		if err != nil {
+			return errors.Wrapf(err, "arithemtic")
+		}
+
+		b, err := ValueToFloat64(ops[2].Value, ctx)
+		if err != nil {
+			return errors.Wrapf(err, "arithemtic")
+		}
+
+		c, err := f(a, b)
+		if err != nil {
+			return err
+		}
+
+		cv := ops[3].Value
+
+		switch cv := cv.(type) {
+		case ast.Var:
+			ctx = ctx.BindVar(cv, c)
+			return iter(ctx)
+		default:
+			if cv.Equal(c) {
+				return iter(ctx)
+			}
+			return nil
+		}
+	}
+}
diff --git a/topdown/builtins.go b/topdown/builtins.go
index 187e96959a..561acd117c 100644
--- a/topdown/builtins.go
+++ b/topdown/builtins.go
@@ -31,7 +31,11 @@ var defaultBuiltinFuncs = map[ast.Var]BuiltinFunc{
 	ast.LessThan.Name:      evalIneq(compareLessThan),
 	ast.LessThanEq.Name:    evalIneq(compareLessThanEq),
 	ast.NotEqual.Name:      evalIneq(compareNotEq),
-	ast.Plus.Name:          evalPlus,
+	ast.Plus.Name:          evalArithmetic(arithPlus),
+	ast.Minus.Name:         evalArithmetic(arithMinus),
+	ast.Multiply.Name:      evalArithmetic(arithMultiply),
+	ast.Divide.Name:        evalArithmetic(arithDivide),
+	ast.Round.Name:         evalRound,
 	ast.Count.Name:         evalCount,
 	ast.Sum.Name:           evalSum,
 	ast.ToNumber.Name:      evalToNumber,
diff --git a/topdown/topdown_test.go b/topdown/topdown_test.go
index e3dd862d89..0fec7a6a7e 100644
--- a/topdown/topdown_test.go
+++ b/topdown/topdown_test.go
@@ -604,6 +604,10 @@ func TestTopDownArithmetic(t *testing.T) {
 		expected interface{}
 	}{
 		{"plus", []string{"p[y] :- a[i] = x, plus(i, x, y)"}, "[1,3,5,7]"},
+		{"minus", []string{"p[y] :- a[i] = x, minus(i, x, y)"}, "[-1,-1,-1,-1]"},
+		{"multiply", []string{"p[y] :- a[i] = x, mul(i, x, y)"}, "[0,2,6,12]"},
+		{"divide+round", []string{"p[z] :- a[i] = x, div(i, x, y), round(y, z)"}, "[0,1,1,1]"},
+		{"divide+error", []string{"p[y] :- a[i] = x, div(x, i, y)"}, fmt.Errorf("divide: by zero")},
 	}
 
 	data := loadSmallTestData()