refactor: replaces calls to astutils.PickNodes with calls to astutils.SeekNode when possible (#1482)

internal/astutils/ast_utils.go

@@ -152,6 +152,49 @@ func (p picker) Visit(node ast.Node) ast.Visitor {
 	return p
 }
 
+// SeekNode yields the first node selected by the given selector function in the AST subtree with root n.
+// The function returns nil if no matching node is found in the subtree.
+func SeekNode[T ast.Node](n ast.Node, selector func(n ast.Node) bool) T {
+	var result T
+
+	if n == nil {
+		return result
+	}
+
+	if selector == nil {
+		return result
+	}
+
+	onSelect := func(n ast.Node) {
+		result, _ = n.(T)
+	}
+
+	p := &seeker{selector: selector, onSelect: onSelect, found: false}
+	ast.Walk(p, n)
+
+	return result
+}
+
+type seeker struct {
+	selector func(n ast.Node) bool
+	onSelect func(n ast.Node)
+	found    bool
+}
+
+func (s *seeker) Visit(node ast.Node) ast.Visitor {
+	if s.found {
+		return nil // stop visiting subtree
+	}
+
+	if s.selector(node) {
+		s.onSelect(node)
+		s.found = true
+		return nil // skip visiting node children
+	}
+
+	return s
+}
+
 var gofmtConfig = &printer.Config{Tabwidth: 8}
 
 // GoFmt returns a string representation of an AST subtree.