From 2fdc5717a3db2afc3b165afaa9237a7fd9d7cafd Mon Sep 17 00:00:00 2001
From: royels <docmentz@gmail.com>
Date: Sun, 29 May 2016 05:06:14 -0400
Subject: [PATCH] promql: add power binary operation

---
 promql/engine.go               |  4 ++++
 promql/lex.go                  | 16 ++++++++++++++++
 promql/lex_test.go             |  3 +++
 promql/parse.go                | 25 ++++++++++++++-----------
 promql/testdata/operators.test | 25 ++++++++++++++++++++++++-
 5 files changed, 61 insertions(+), 12 deletions(-)

diff --git a/promql/engine.go b/promql/engine.go
index 0bfcbe2d8..23dfeef25 100644
--- a/promql/engine.go
+++ b/promql/engine.go
@@ -991,6 +991,8 @@ func scalarBinop(op itemType, lhs, rhs model.SampleValue) model.SampleValue {
 		return lhs * rhs
 	case itemDIV:
 		return lhs / rhs
+	case itemPOW:
+		return model.SampleValue(math.Pow(float64(lhs), float64(rhs)))
 	case itemMOD:
 		if rhs != 0 {
 			return model.SampleValue(int(lhs) % int(rhs))
@@ -1023,6 +1025,8 @@ func vectorElemBinop(op itemType, lhs, rhs model.SampleValue) (model.SampleValue
 		return lhs * rhs, true
 	case itemDIV:
 		return lhs / rhs, true
+	case itemPOW:
+		return model.SampleValue(math.Pow(float64(lhs), float64(rhs))), true
 	case itemMOD:
 		if rhs != 0 {
 			return model.SampleValue(int(lhs) % int(rhs)), true
diff --git a/promql/lex.go b/promql/lex.go
index e37460e3b..d62501198 100644
--- a/promql/lex.go
+++ b/promql/lex.go
@@ -100,11 +100,23 @@ func (i itemType) precedence() int {
 		return 4
 	case itemMUL, itemDIV, itemMOD:
 		return 5
+	case itemPOW:
+		return 6
 	default:
 		return LowestPrec
 	}
 }
 
+func (i itemType) isRightAssociative() bool {
+	switch i {
+	case itemPOW:
+		return true
+	default:
+		return false
+	}
+
+}
+
 type itemType int
 
 const (
@@ -146,6 +158,7 @@ const (
 	itemGTR
 	itemEQLRegex
 	itemNEQRegex
+	itemPOW
 	operatorsEnd
 
 	aggregatorsStart
@@ -248,6 +261,7 @@ var itemTypeStr = map[itemType]string{
 	itemGTR:      ">",
 	itemEQLRegex: "=~",
 	itemNEQRegex: "!~",
+	itemPOW:      "^",
 }
 
 func init() {
@@ -471,6 +485,8 @@ func lexStatements(l *lexer) stateFn {
 		l.emit(itemADD)
 	case r == '-':
 		l.emit(itemSUB)
+	case r == '^':
+		l.emit(itemPOW)
 	case r == '=':
 		if t := l.peek(); t == '=' {
 			l.next()
diff --git a/promql/lex_test.go b/promql/lex_test.go
index 492a3aefc..824b4e00c 100644
--- a/promql/lex_test.go
+++ b/promql/lex_test.go
@@ -208,6 +208,9 @@ var tests = []struct {
 	}, {
 		input:    `/`,
 		expected: []item{{itemDIV, 0, `/`}},
+	}, {
+		input:    `^`,
+		expected: []item{{itemPOW, 0, `^`}},
 	}, {
 		input:    `%`,
 		expected: []item{{itemMOD, 0, `%`}},
diff --git a/promql/parse.go b/promql/parse.go
index cac58617e..e2e03b18f 100644
--- a/promql/parse.go
+++ b/promql/parse.go
@@ -500,17 +500,20 @@ func (p *parser) expr() Expr {
 }
 
 func (p *parser) balance(lhs Expr, op itemType, rhs Expr, vecMatching *VectorMatching, returnBool bool) *BinaryExpr {
-	if lhsBE, ok := lhs.(*BinaryExpr); ok && lhsBE.Op.precedence() < op.precedence() {
-		balanced := p.balance(lhsBE.RHS, op, rhs, vecMatching, returnBool)
-		if lhsBE.Op.isComparisonOperator() && !lhsBE.ReturnBool && balanced.Type() == model.ValScalar && lhsBE.LHS.Type() == model.ValScalar {
-			p.errorf("comparisons between scalars must use BOOL modifier")
-		}
-		return &BinaryExpr{
-			Op:             lhsBE.Op,
-			LHS:            lhsBE.LHS,
-			RHS:            balanced,
-			VectorMatching: lhsBE.VectorMatching,
-			ReturnBool:     lhsBE.ReturnBool,
+	if lhsBE, ok := lhs.(*BinaryExpr); ok {
+		precd := lhsBE.Op.precedence() - op.precedence()
+		if (precd < 0) || (precd == 0 && op.isRightAssociative()) {
+			balanced := p.balance(lhsBE.RHS, op, rhs, vecMatching, returnBool)
+			if lhsBE.Op.isComparisonOperator() && !lhsBE.ReturnBool && balanced.Type() == model.ValScalar && lhsBE.LHS.Type() == model.ValScalar {
+				p.errorf("comparisons between scalars must use BOOL modifier")
+			}
+			return &BinaryExpr{
+				Op:             lhsBE.Op,
+				LHS:            lhsBE.LHS,
+				RHS:            balanced,
+				VectorMatching: lhsBE.VectorMatching,
+				ReturnBool:     lhsBE.ReturnBool,
+			}
 		}
 	}
 	if op.isComparisonOperator() && !returnBool && rhs.Type() == model.ValScalar && lhs.Type() == model.ValScalar {
diff --git a/promql/testdata/operators.test b/promql/testdata/operators.test
index c24e7db70..8169ff6f4 100644
--- a/promql/testdata/operators.test
+++ b/promql/testdata/operators.test
@@ -8,7 +8,6 @@ load 5m
 	http_requests{job="app-server", instance="0", group="canary"}		0+70x10
 	http_requests{job="app-server", instance="1", group="canary"}		0+80x10
 
-
 load 5m
 	vector_matching_a{l="x"} 0+1x100
 	vector_matching_a{l="y"} 0+2x50
@@ -35,6 +34,30 @@ eval instant at 50m SUM(http_requests) BY (job) % 3
 	{job="api-server"} 1
 	{job="app-server"} 2
 
+eval instant at 50m SUM(http_requests) BY (job) ^ 2
+	{job="api-server"} 1000000
+	{job="app-server"} 6760000
+
+eval instant at 50m SUM(http_requests) BY (job) % 3 ^ 2
+	{job="api-server"} 1
+	{job="app-server"} 8
+
+eval instant at 50m SUM(http_requests) BY (job) % 2 ^ (3 ^ 2)
+	{job="api-server"} 488
+	{job="app-server"} 40
+
+eval instant at 50m SUM(http_requests) BY (job) % 2 ^ 3 ^ 2
+	{job="api-server"} 488
+	{job="app-server"} 40
+
+eval instant at 50m SUM(http_requests) BY (job) % 2 ^ 3 ^ 2 ^ 2
+	{job="api-server"} 1000
+	{job="app-server"} 2600
+
+eval instant at 50m COUNT(http_requests) BY (job) ^ COUNT(http_requests) BY (job)
+	{job="api-server"} 256
+	{job="app-server"} 256
+
 eval instant at 50m SUM(http_requests) BY (job) / 0
 	{job="api-server"} +Inf
 	{job="app-server"} +Inf