From d0cff2974991736b383e9617f41b7d0bbb081a54 Mon Sep 17 00:00:00 2001
From: Tobias Guggenmos <tguggenm@redhat.com>
Date: Mon, 25 Nov 2019 11:41:59 +0000
Subject: [PATCH] PromQL: Reduce Code duplication in AST traversion (#6362)

* promql: Add Children function for easier AST traversion

Signed-off-by: Tobias Guggenmos <tguggenm@redhat.com>
---
 promql/ast.go     | 103 ++++++++++++++++++++++------------------------
 promql/printer.go |  35 ++--------------
 2 files changed, 53 insertions(+), 85 deletions(-)

diff --git a/promql/ast.go b/promql/ast.go
index 3cc699aa3..973971d29 100644
--- a/promql/ast.go
+++ b/promql/ast.go
@@ -248,61 +248,10 @@ func Walk(v Visitor, node Node, path []Node) error {
 	}
 	path = append(path, node)
 
-	switch n := node.(type) {
-	case *EvalStmt:
-		if err := Walk(v, n.Expr, path); err != nil {
+	for _, e := range Children(node) {
+		if err := Walk(v, e, path); err != nil {
 			return err
 		}
-
-	case Expressions:
-		for _, e := range n {
-			if err := Walk(v, e, path); err != nil {
-				return err
-			}
-		}
-	case *AggregateExpr:
-		if n.Param != nil {
-			if err := Walk(v, n.Param, path); err != nil {
-				return err
-			}
-		}
-		if err := Walk(v, n.Expr, path); err != nil {
-			return err
-		}
-
-	case *BinaryExpr:
-		if err := Walk(v, n.LHS, path); err != nil {
-			return err
-		}
-		if err := Walk(v, n.RHS, path); err != nil {
-			return err
-		}
-
-	case *Call:
-		if err := Walk(v, n.Args, path); err != nil {
-			return err
-		}
-
-	case *SubqueryExpr:
-		if err := Walk(v, n.Expr, path); err != nil {
-			return err
-		}
-
-	case *ParenExpr:
-		if err := Walk(v, n.Expr, path); err != nil {
-			return err
-		}
-
-	case *UnaryExpr:
-		if err := Walk(v, n.Expr, path); err != nil {
-			return err
-		}
-
-	case *MatrixSelector, *NumberLiteral, *StringLiteral, *VectorSelector:
-		// nothing to do
-
-	default:
-		panic(errors.Errorf("promql.Walk: unhandled node type %T", node))
 	}
 
 	_, err = v.Visit(nil, nil)
@@ -326,3 +275,51 @@ func Inspect(node Node, f inspector) {
 	//nolint: errcheck
 	Walk(inspector(f), node, nil)
 }
+
+// Children returns a list of all child nodes of a syntax tree node.
+func Children(node Node) []Node {
+	// For some reasons these switches have significantly better performance than interfaces
+	switch n := node.(type) {
+	case *EvalStmt:
+		return []Node{n.Expr}
+	case Expressions:
+		// golang cannot convert slices of interfaces
+		ret := make([]Node, len(n))
+		for i, e := range n {
+			ret[i] = e
+		}
+		return ret
+	case *AggregateExpr:
+		// While this does not look nice, it should avoid unnecessary allocations
+		// caused by slice resizing
+		if n.Expr == nil && n.Param == nil {
+			return nil
+		} else if n.Expr == nil {
+			return []Node{n.Param}
+		} else if n.Param == nil {
+			return []Node{n.Expr}
+		} else {
+			return []Node{n.Expr, n.Param}
+		}
+	case *BinaryExpr:
+		return []Node{n.LHS, n.RHS}
+	case *Call:
+		// golang cannot convert slices of interfaces
+		ret := make([]Node, len(n.Args))
+		for i, e := range n.Args {
+			ret[i] = e
+		}
+		return ret
+	case *SubqueryExpr:
+		return []Node{n.Expr}
+	case *ParenExpr:
+		return []Node{n.Expr}
+	case *UnaryExpr:
+		return []Node{n.Expr}
+	case *MatrixSelector, *NumberLiteral, *StringLiteral, *VectorSelector:
+		// nothing to do
+		return []Node{}
+	default:
+		panic(errors.Errorf("promql.Children: unhandled node type %T", node))
+	}
+}
diff --git a/promql/printer.go b/promql/printer.go
index adcf2e699..167b3fb8e 100644
--- a/promql/printer.go
+++ b/promql/printer.go
@@ -38,39 +38,10 @@ func tree(node Node, level string) string {
 
 	level += " · · ·"
 
-	switch n := node.(type) {
-	case *EvalStmt:
-		t += tree(n.Expr, level)
-
-	case Expressions:
-		for _, e := range n {
-			t += tree(e, level)
-		}
-	case *AggregateExpr:
-		t += tree(n.Expr, level)
-
-	case *BinaryExpr:
-		t += tree(n.LHS, level)
-		t += tree(n.RHS, level)
-
-	case *Call:
-		t += tree(n.Args, level)
-
-	case *ParenExpr:
-		t += tree(n.Expr, level)
-
-	case *UnaryExpr:
-		t += tree(n.Expr, level)
-
-	case *SubqueryExpr:
-		t += tree(n.Expr, level)
-
-	case *MatrixSelector, *NumberLiteral, *StringLiteral, *VectorSelector:
-		// nothing to do
-
-	default:
-		panic("promql.Tree: not all node types covered")
+	for _, e := range Children(node) {
+		t += tree(e, level)
 	}
+
 	return t
 }