From fffde1483253138e25a8700750d52763b20aec16 Mon Sep 17 00:00:00 2001
From: darshanime <deathbullet@gmail.com>
Date: Thu, 1 Feb 2024 07:25:08 +0530
Subject: [PATCH] unwrap paren expr with stepinvariant expr

Signed-off-by: darshanime <deathbullet@gmail.com>
---
 promql/engine.go    | 24 +++++++++---------------
 promql/functions.go |  3 +--
 2 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/promql/engine.go b/promql/engine.go
index ea4bc1af8..7db23674f 100644
--- a/promql/engine.go
+++ b/promql/engine.go
@@ -1442,9 +1442,7 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 		sortedGrouping := e.Grouping
 		slices.Sort(sortedGrouping)
 
-		unwrapParenExpr(&e.Param)
-		param := unwrapStepInvariantExpr(e.Param)
-		unwrapParenExpr(&param)
+		param := unwrapStepInvariantAndParenExprs(&e.Param)
 
 		if e.Op == parser.COUNT_VALUES {
 			valueLabel := param.(*parser.StringLiteral)
@@ -1487,9 +1485,7 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 			// Matrix evaluation always returns the evaluation time,
 			// so this function needs special handling when given
 			// a vector selector.
-			unwrapParenExpr(&e.Args[0])
-			arg := unwrapStepInvariantExpr(e.Args[0])
-			unwrapParenExpr(&arg)
+			arg := unwrapStepInvariantAndParenExprs(&e.Args[0])
 			vs, ok := arg.(*parser.VectorSelector)
 			if ok {
 				return ev.rangeEvalTimestampFunctionOverVectorSelector(vs, call, e)
@@ -1503,9 +1499,7 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 			warnings       annotations.Annotations
 		)
 		for i := range e.Args {
-			unwrapParenExpr(&e.Args[i])
-			a := unwrapStepInvariantExpr(e.Args[i])
-			unwrapParenExpr(&a)
+			a := unwrapStepInvariantAndParenExprs(&e.Args[i])
 			if _, ok := a.(*parser.MatrixSelector); ok {
 				matrixArgIndex = i
 				matrixArg = true
@@ -1558,9 +1552,7 @@ func (ev *evaluator) eval(expr parser.Expr) (parser.Value, annotations.Annotatio
 			}
 		}
 
-		unwrapParenExpr(&e.Args[matrixArgIndex])
-		arg := unwrapStepInvariantExpr(e.Args[matrixArgIndex])
-		unwrapParenExpr(&arg)
+		arg := unwrapStepInvariantAndParenExprs(&e.Args[matrixArgIndex])
 		sel := arg.(*parser.MatrixSelector)
 		selVS := sel.VectorSelector.(*parser.VectorSelector)
 
@@ -3169,11 +3161,13 @@ func unwrapParenExpr(e *parser.Expr) {
 	}
 }
 
-func unwrapStepInvariantExpr(e parser.Expr) parser.Expr {
-	if p, ok := e.(*parser.StepInvariantExpr); ok {
+func unwrapStepInvariantAndParenExprs(e *parser.Expr) parser.Expr {
+	unwrapParenExpr(e)
+	if p, ok := (*e).(*parser.StepInvariantExpr); ok {
+		unwrapParenExpr(&p.Expr)
 		return p.Expr
 	}
-	return e
+	return *e
 }
 
 // PreprocessExpr wraps all possible step invariant parts of the given expression with
diff --git a/promql/functions.go b/promql/functions.go
index 9b3be2287..85acf751e 100644
--- a/promql/functions.go
+++ b/promql/functions.go
@@ -1698,8 +1698,7 @@ func createLabelsForAbsentFunction(expr parser.Expr) labels.Labels {
 }
 
 func stringFromArg(e parser.Expr) string {
-	tmp := unwrapStepInvariantExpr(e) // Unwrap StepInvariant
-	unwrapParenExpr(&tmp)             // Optionally unwrap ParenExpr
+	tmp := unwrapStepInvariantAndParenExprs(&e) // Unwrap StepInvariant
 	return tmp.(*parser.StringLiteral).Val
 }