go/ast/inspector: add PreorderSeq and All iterators

This CL adds two functions that return go1.23 iterators
over selected nodes:
- inspect.PreorderSeq(types...), which is like Preorder but
  honors the break statement;
- All[N](inspect), which iterates over nodes of a single
  type, determined by the type parameter.

+ Tests, benchmark.

Fixes golang/go#67795

Change-Id: I77817f2e595846cf3ce29dc347ae895e927fc805
Reviewed-on: https://go-review.googlesource.com/c/tools/+/616218
Auto-Submit: Alan Donovan <adonovan@google.com>
Commit-Queue: Alan Donovan <adonovan@google.com>
LUCI-TryBot-Result: Go LUCI <golang-scoped@luci-project-accounts.iam.gserviceaccount.com>
Reviewed-by: Robert Findley <rfindley@google.com>

Author: adonovan

go/ast/inspector/inspector.go

@@ -73,6 +73,15 @@ func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node)) {
 	// check, Preorder is almost twice as fast as Nodes. The two
 	// features seem to contribute similar slowdowns (~1.4x each).
 
+	// This function is equivalent to the PreorderSeq call below,
+	// but to avoid the additional dynamic call (which adds 13-35%
+	// to the benchmarks), we expand it out.
+	//
+	// in.PreorderSeq(types...)(func(n ast.Node) bool {
+	// 	f(n)
+	// 	return true
+	// })
+
 	mask := maskOf(types)
 	for i := 0; i < len(in.events); {
 		ev := in.events[i]

go/ast/inspector/inspector_test.go

@@ -160,7 +160,8 @@ func TestInspectPruning(t *testing.T) {
 	compare(t, nodesA, nodesB)
 }
 
-func compare(t *testing.T, nodesA, nodesB []ast.Node) {
+// compare calls t.Error if !slices.Equal(nodesA, nodesB).
+func compare[N comparable](t *testing.T, nodesA, nodesB []N) {
 	if len(nodesA) != len(nodesB) {
 		t.Errorf("inconsistent node lists: %d vs %d", len(nodesA), len(nodesB))
 	} else {

go/ast/inspector/iter.go

@@ -0,0 +1,85 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build go1.23
+
+package inspector
+
+import (
+	"go/ast"
+	"iter"
+)
+
+// PreorderSeq returns an iterator that visits all the
+// nodes of the files supplied to New in depth-first order.
+// It visits each node n before n's children.
+// The complete traversal sequence is determined by ast.Inspect.
+//
+// The types argument, if non-empty, enables type-based
+// filtering of events: only nodes whose type matches an
+// element of the types slice are included in the sequence.
+func (in *Inspector) PreorderSeq(types ...ast.Node) iter.Seq[ast.Node] {
+
+	// This implementation is identical to Preorder,
+	// except that it supports breaking out of the loop.
+
+	return func(yield func(ast.Node) bool) {
+		mask := maskOf(types)
+		for i := 0; i < len(in.events); {
+			ev := in.events[i]
+			if ev.index > i {
+				// push
+				if ev.typ&mask != 0 {
+					if !yield(ev.node) {
+						break
+					}
+				}
+				pop := ev.index
+				if in.events[pop].typ&mask == 0 {
+					// Subtrees do not contain types: skip them and pop.
+					i = pop + 1
+					continue
+				}
+			}
+			i++
+		}
+	}
+}
+
+// All[N] returns an iterator over all the nodes of type N.
+// N must be a pointer-to-struct type that implements ast.Node.
+//
+// Example:
+//
+//	for call := range All[*ast.CallExpr](in) { ... }
+func All[N interface {
+	*S
+	ast.Node
+}, S any](in *Inspector) iter.Seq[N] {
+
+	// To avoid additional dynamic call overheads,
+	// we duplicate rather than call the logic of PreorderSeq.
+
+	mask := typeOf((N)(nil))
+	return func(yield func(N) bool) {
+		for i := 0; i < len(in.events); {
+			ev := in.events[i]
+			if ev.index > i {
+				// push
+				if ev.typ&mask != 0 {
+					if !yield(ev.node.(N)) {
+						break
+					}
+				}
+				pop := ev.index
+				if in.events[pop].typ&mask == 0 {
+					// Subtrees do not contain types: skip them and pop.
+					i = pop + 1
+					continue
+				}
+			}
+			i++
+		}
+	}
+}

go/ast/inspector/iter_test.go

@@ -0,0 +1,83 @@
+// Copyright 2024 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+//go:build go1.23
+
+package inspector_test
+
+import (
+	"go/ast"
+	"iter"
+	"slices"
+	"testing"
+
+	"golang.org/x/tools/go/ast/inspector"
+)
+
+// TestPreorderSeq checks PreorderSeq against Preorder.
+func TestPreorderSeq(t *testing.T) {
+	inspect := inspector.New(netFiles)
+
+	nodeFilter := []ast.Node{(*ast.FuncDecl)(nil), (*ast.FuncLit)(nil)}
+
+	// reference implementation
+	var want []ast.Node
+	inspect.Preorder(nodeFilter, func(n ast.Node) {
+		want = append(want, n)
+	})
+
+	// Check entire sequence.
+	got := slices.Collect(inspect.PreorderSeq(nodeFilter...))
+	compare(t, got, want)
+
+	// Check that break works.
+	got = firstN(10, inspect.PreorderSeq(nodeFilter...))
+	compare(t, got, want[:10])
+}
+
+// TestAll checks All against Preorder.
+func TestAll(t *testing.T) {
+	inspect := inspector.New(netFiles)
+
+	// reference implementation
+	var want []*ast.CallExpr
+	inspect.Preorder([]ast.Node{(*ast.CallExpr)(nil)}, func(n ast.Node) {
+		want = append(want, n.(*ast.CallExpr))
+	})
+
+	// Check entire sequence.
+	got := slices.Collect(inspector.All[*ast.CallExpr](inspect))
+	compare(t, got, want)
+
+	// Check that break works.
+	got = firstN(10, inspector.All[*ast.CallExpr](inspect))
+	compare(t, got, want[:10])
+}
+
+// firstN(n, seq), returns a slice of up to n elements of seq.
+func firstN[T any](n int, seq iter.Seq[T]) (res []T) {
+	for x := range seq {
+		res = append(res, x)
+		if len(res) == n {
+			break
+		}
+	}
+	return res
+}
+
+// BenchmarkAllCalls is like BenchmarkInspectCalls,
+// but using the single-type filtering iterator, All.
+// (The iterator adds about 5-15%.)
+func BenchmarkAllCalls(b *testing.B) {
+	inspect := inspector.New(netFiles)
+	b.ResetTimer()
+
+	// Measure marginal cost of traversal.
+	var ncalls int
+	for range b.N {
+		for range inspector.All[*ast.CallExpr](inspect) {
+			ncalls++
+		}
+	}
+}