[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,58 +375,66 @@ 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.
+// ChildrenIter returns an iterator over all child nodes of a syntax tree node.
-func Children(node Node) []Node {
+func ChildrenIter(node Node) func(func(Node) bool) {
-	// For some reasons these switches have significantly better performance than interfaces
+	return func(yield func(Node) bool) {
+		// According to lore, these switches have significantly better performance than interfaces
 		switch n := node.(type) {
 		case *EvalStmt:
-		return []Node{n.Expr}
+			yield(n.Expr)
 		case Expressions:
-		// golang cannot convert slices of interfaces
+			for _, e := range n {
-		ret := make([]Node, len(n))
+				if !yield(e) {
-		for i, e := range n {
+					return
-			ret[i] = e
+				}
 			}
-		return ret
 		case *AggregateExpr:
-		// While this does not look nice, it should avoid unnecessary allocations
+			if n.Param != nil {
-		// caused by slice resizing
+				if !yield(n.Param) {
-		switch {
+					return
-		case n.Expr == nil && n.Param == nil:
+				}
-			return nil
+			}
-		case n.Expr == nil:
+			if n.Expr != nil {
-			return []Node{n.Param}
+				yield(n.Expr)
-		case n.Param == nil:
-			return []Node{n.Expr}
-		default:
-			return []Node{n.Expr, n.Param}
 			}
 		case *BinaryExpr:
-		return []Node{n.LHS, n.RHS}
+			if !yield(n.LHS) {
-	case *Call:
+				return
-		// golang cannot convert slices of interfaces
+			}
-		ret := make([]Node, len(n.Args))
+			yield(n.RHS)
-		for i, e := range n.Args {
+		case *Call:
-			ret[i] = e
+			for _, e := range n.Args {
+				if !yield(e) {
+					return
+				}
 			}
-		return ret
 		case *SubqueryExpr:
-		return []Node{n.Expr}
+			yield(n.Expr)
 		case *ParenExpr:
-		return []Node{n.Expr}
+			yield(n.Expr)
 		case *UnaryExpr:
-		return []Node{n.Expr}
+			yield(n.Expr)
 		case *MatrixSelector:
-		return []Node{n.VectorSelector}
+			yield(n.VectorSelector)
 		case *StepInvariantExpr:
-		return []Node{n.Expr}
+			yield(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

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)
 	}