[PERF] PromQL: Walk syntax tree using iterator

Saves memory allocations. This is called a few times from the PromQL Engine.

Signed-off-by: Bryan Boreham <bjboreham@gmail.com>

promql/parser/ast.go

@@ -336,7 +336,7 @@ func Walk(v Visitor, node Node, path []Node) error {
 	}
 	path = append(path, node)
 
-	for _, e := range Children(node) {
+	for e := range ChildrenIter(node) {
 		if err := Walk(v, e, path); err != nil {
 			return err
 		}
@@ -375,59 +375,67 @@ func Inspect(node Node, f inspector) {
 	Walk(f, node, nil) //nolint:errcheck
 }
 
-// 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
-		switch {
-		case n.Expr == nil && n.Param == nil:
-			return nil
-		case n.Expr == nil:
-			return []Node{n.Param}
-		case n.Param == nil:
-			return []Node{n.Expr}
+// ChildrenIter returns an iterator over all child nodes of a syntax tree node.
+func ChildrenIter(node Node) func(func(Node) bool) {
+	return func(yield func(Node) bool) {
+		// According to lore, these switches have significantly better performance than interfaces
+		switch n := node.(type) {
+		case *EvalStmt:
+			yield(n.Expr)
+		case Expressions:
+			for _, e := range n {
+				if !yield(e) {
+					return
+				}
+			}
+		case *AggregateExpr:
+			if n.Param != nil {
+				if !yield(n.Param) {
+					return
+				}
+			}
+			if n.Expr != nil {
+				yield(n.Expr)
+			}
+		case *BinaryExpr:
+			if !yield(n.LHS) {
+				return
+			}
+			yield(n.RHS)
+		case *Call:
+			for _, e := range n.Args {
+				if !yield(e) {
+					return
+				}
+			}
+		case *SubqueryExpr:
+			yield(n.Expr)
+		case *ParenExpr:
+			yield(n.Expr)
+		case *UnaryExpr:
+			yield(n.Expr)
+		case *MatrixSelector:
+			yield(n.VectorSelector)
+		case *StepInvariantExpr:
+			yield(n.Expr)
+		case *NumberLiteral, *StringLiteral, *VectorSelector:
+			// nothing to do
 		default:
-			return []Node{n.Expr, n.Param}
+			panic(fmt.Errorf("promql.Children: unhandled node type %T", node))
 		}
-	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:
-		return []Node{n.VectorSelector}
-	case *StepInvariantExpr:
-		return []Node{n.Expr}
-	case *NumberLiteral, *StringLiteral, *VectorSelector:
-		// nothing to do
-		return []Node{}
-	default:
-		panic(fmt.Errorf("promql.Children: unhandled node type %T", node))
 	}
 }
 
+// Children returns a list of all child nodes of a syntax tree node.
+// Implemented for backwards-compatibility; prefer ChildrenIter().
+func Children(node Node) []Node {
+	ret := []Node{}
+	for e := range ChildrenIter(node) {
+		ret = append(ret, e)
+	}
+	return ret
+}
+
 // mergeRanges is a helper function to merge the PositionRanges of two Nodes.
 // Note that the arguments must be in the same order as they
 // occur in the input string.

promql/parser/printer.go

@@ -41,7 +41,7 @@ func tree(node Node, level string) string {
 
 	level += " · · ·"
 
-	for _, e := range Children(node) {
+	for e := range ChildrenIter(node) {
 		t += tree(e, level)
 	}