internal/astutil/cursor: Cursor.Ancestors iterator

This CL adds an iterator for the strict ancestors
(transitive parents) of the current node.
It accepts a type filter.

Also, use it in two places.

This may slow down Cursor.Stack slightly, but see
comments in benchmark for justification. The
basic Cursor traversal is still competitive so
long as Stack or Ancestors are called sparingly.

BenchmarkInspectCalls/Preorder-8         	   12069	    106995 ns/op
BenchmarkInspectCalls/WithStack-8        	    7447	    153992 ns/op
BenchmarkInspectCalls/CursorStack-8      	    2545	    460433 ns/op (was 178907: 2.6x increase)
BenchmarkInspectCalls/Cursor-8           	   12225	     99472 ns/op
BenchmarkInspectCalls/CursorAncestors-8  	    3249	    350503 ns/op (=3x WithStack)

Change-Id: I79941423888028a622b20b7ab63b37f8435dce33
Reviewed-on: https://go-review.googlesource.com/c/tools/+/641435
Auto-Submit: Alan Donovan <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Robert Findley <[email protected]>

Author: adonovan

gopls/internal/analysis/hostport/hostport.go

@@ -164,8 +164,7 @@ func run(pass *analysis.Pass) (any, error) {
 				// Search for decl of addrVar within common ancestor of addrVar and Dial call.
 				if addrVar, ok := info.Uses[address].(*types.Var); ok {
 					pos := addrVar.Pos()
-					// TODO(adonovan): use Cursor.Ancestors iterator when available.
-					for _, curAncestor := range curCall.Stack(nil) {
+					for curAncestor := range curCall.Ancestors() {
 						if curIdent, ok := curAncestor.FindPos(pos, pos); ok {
 							// curIdent is the declaring ast.Ident of addr.
 							switch parent := curIdent.Parent().Node().(type) {

gopls/internal/analysis/modernize/bloop.go

@@ -36,12 +36,12 @@ func bloop(pass *analysis.Pass) {
 	// edits computes the text edits for a matched for/range loop
 	// at the specified cursor. b is the *testing.B value, and
 	// (start, end) is the portion using b.N to delete.
-	edits := func(cur cursor.Cursor, b ast.Expr, start, end token.Pos) (edits []analysis.TextEdit) {
+	edits := func(curLoop cursor.Cursor, b ast.Expr, start, end token.Pos) (edits []analysis.TextEdit) {
+		curFn, _ := enclosingFunc(curLoop)
 		// Within the same function, delete all calls to
 		// b.{Start,Stop,Timer} that precede the loop.
 		filter := []ast.Node{(*ast.ExprStmt)(nil), (*ast.FuncLit)(nil)}
-		fn, _ := enclosingFunc(cur)
-		fn.Inspect(filter, func(cur cursor.Cursor, push bool) (descend bool) {
+		curFn.Inspect(filter, func(cur cursor.Cursor, push bool) (descend bool) {
 			if push {
 				node := cur.Node()
 				if is[*ast.FuncLit](node) {
@@ -162,22 +162,10 @@ func uses(info *types.Info, cur cursor.Cursor, obj types.Object) bool {
 }
 
 // enclosingFunc returns the cursor for the innermost Func{Decl,Lit}
-// that encloses (or is) c, if any.
-//
-// TODO(adonovan): consider adding:
-//
-//	func (Cursor) AnyEnclosing(filter ...ast.Node) (Cursor bool)
-//	func (Cursor) Enclosing[N ast.Node]() (Cursor, bool)
-//
-// See comments at [cursor.Cursor.Stack].
+// that encloses c, if any.
 func enclosingFunc(c cursor.Cursor) (cursor.Cursor, bool) {
-	for {
-		switch c.Node().(type) {
-		case *ast.FuncLit, *ast.FuncDecl:
-			return c, true
-		case nil:
-			return cursor.Cursor{}, false
-		}
-		c = c.Parent()
+	for curAncestor := range c.Ancestors((*ast.FuncLit)(nil), (*ast.FuncDecl)(nil)) {
+		return curAncestor, true
 	}
+	return cursor.Cursor{}, false
 }

internal/astutil/cursor/cursor.go

@@ -157,13 +157,6 @@ func (c Cursor) Inspect(types []ast.Node, f func(c Cursor, push bool) (descend b
 // must be empty.
 //
 // Stack must not be called on the Root node.
-//
-// TODO(adonovan): perhaps this should be replaced by:
-//
-//	func (Cursor) Ancestors(filter []ast.Node) iter.Seq[Cursor]
-//
-// returning a filtering iterator up the parent chain.
-// This finesses the question of allocation entirely.
 func (c Cursor) Stack(stack []Cursor) []Cursor {
 	if len(stack) > 0 {
 		panic("stack is non-empty")
@@ -172,14 +165,37 @@ func (c Cursor) Stack(stack []Cursor) []Cursor {
 		panic("Cursor.Stack called on Root node")
 	}
 
-	events := c.events()
-	for i := c.index; i >= 0; i = events[i].parent {
-		stack = append(stack, Cursor{c.in, i})
-	}
+	stack = append(stack, c)
+	stack = slices.AppendSeq(stack, c.Ancestors())
 	slices.Reverse(stack)
 	return stack
 }
 
+// Ancestors returns an iterator over the ancestors of the current
+// node, starting with [Cursor.Parent].
+//
+// Ancestors must not be called on the Root node (whose [Cursor.Node] returns nil).
+//
+// The types argument, if non-empty, enables type-based filtering of
+// events: the sequence includes only ancestors whose type matches an
+// element of the types slice.
+func (c Cursor) Ancestors(types ...ast.Node) iter.Seq[Cursor] {
+	if c.index < 0 {
+		panic("Cursor.Ancestors called on Root node")
+	}
+
+	mask := maskOf(types)
+
+	return func(yield func(Cursor) bool) {
+		events := c.events()
+		for i := events[c.index].parent; i >= 0; i = events[i].parent {
+			if events[i].typ&mask != 0 && !yield(Cursor{c.in, i}) {
+				break
+			}
+		}
+	}
+}
+
 // Parent returns the parent of the current node.
 //
 // Parent must not be called on the Root node (whose [Cursor.Node] returns nil).

internal/astutil/cursor/cursor_test.go

@@ -15,6 +15,7 @@ import (
 	"iter"
 	"log"
 	"path/filepath"
+	"reflect"
 	"slices"
 	"strings"
 	"testing"
@@ -152,6 +153,13 @@ func g() {
 			if got, want := fmt.Sprint(stack), "[*ast.File *ast.FuncDecl *ast.BlockStmt *ast.ExprStmt *ast.CallExpr]"; got != want {
 				t.Errorf("curCall.Stack() = %q, want %q", got, want)
 			}
+
+			// Ancestors = Reverse(Stack[:last]).
+			stack = stack[:len(stack)-1]
+			slices.Reverse(stack)
+			if got, want := slices.Collect(curCall.Ancestors()), stack; !reflect.DeepEqual(got, want) {
+				t.Errorf("Ancestors = %v, Reverse(Stack - last element) = %v", got, want)
+			}
 		}
 
 		// nested Inspect traversal
@@ -381,6 +389,15 @@ func BenchmarkInspectCalls(b *testing.B) {
 	// And if the calls to Stack are very selective,
 	// or are replaced by 2 calls to Parent, it runs
 	// 27% faster than WithStack.
+	//
+	// But the purpose of inspect.WithStack is not to obtain the
+	// stack on every node, but to perform a traversal in which it
+	// one as the _option_ to access the stack if it should be
+	// needed, but the need is rare and usually only for a small
+	// portion. Arguably, because Cursor traversals always
+	// provide, at no extra cost, the option to access the
+	// complete stack, the right comparison is the plain Cursor
+	// benchmark below.
 	b.Run("CursorStack", func(b *testing.B) {
 		var ncalls int
 		for range b.N {
@@ -392,6 +409,28 @@ func BenchmarkInspectCalls(b *testing.B) {
 			}
 		}
 	})
+
+	b.Run("Cursor", func(b *testing.B) {
+		var ncalls int
+		for range b.N {
+			for cur := range cursor.Root(inspect).Preorder(callExprs...) {
+				_ = cur.Node().(*ast.CallExpr)
+				ncalls++
+			}
+		}
+	})
+
+	b.Run("CursorAncestors", func(b *testing.B) {
+		var ncalls int
+		for range b.N {
+			for cur := range cursor.Root(inspect).Preorder(callExprs...) {
+				_ = cur.Node().(*ast.CallExpr)
+				for range cur.Ancestors() {
+				}
+				ncalls++
+			}
+		}
+	})
 }
 
 // This benchmark compares methods for finding a known node in a tree.