internal/astutil: PreorderStack: a safer ast.Inspect for stacks

This CL defines PreorderStack, a safer function than ast.Inspect
for when you need to maintain a stack.

Beware, the stack that it produces does not include n itself--a
half-open interval--so that nested traversals compose correctly.

The CL also uses the new function in various places in x/tools
where appropriate; in some cases it was clearer to rewrite
using cursor.Cursor.

+ test

Updates golang/go#73319

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

Author: adonovan

go/analysis/passes/lostcancel/lostcancel.go

@@ -17,6 +17,7 @@ import (
 	"golang.org/x/tools/go/ast/inspector"
 	"golang.org/x/tools/go/cfg"
 	"golang.org/x/tools/internal/analysisinternal"
+	"golang.org/x/tools/internal/astutil"
 )
 
 //go:embed doc.go
@@ -83,30 +84,22 @@ func runFunc(pass *analysis.Pass, node ast.Node) {
 	// {FuncDecl,FuncLit,CallExpr,SelectorExpr}.
 
 	// Find the set of cancel vars to analyze.
-	stack := make([]ast.Node, 0, 32)
-	ast.Inspect(node, func(n ast.Node) bool {
-		switch n.(type) {
-		case *ast.FuncLit:
-			if len(stack) > 0 {
-				return false // don't stray into nested functions
-			}
-		case nil:
-			stack = stack[:len(stack)-1] // pop
-			return true
+	astutil.PreorderStack(node, nil, func(n ast.Node, stack []ast.Node) bool {
+		if _, ok := n.(*ast.FuncLit); ok && len(stack) > 0 {
+			return false // don't stray into nested functions
 		}
-		stack = append(stack, n) // push
 
-		// Look for [{AssignStmt,ValueSpec} CallExpr SelectorExpr]:
+		// Look for n=SelectorExpr beneath stack=[{AssignStmt,ValueSpec} CallExpr]:
 		//
 		//   ctx, cancel    := context.WithCancel(...)
 		//   ctx, cancel     = context.WithCancel(...)
 		//   var ctx, cancel = context.WithCancel(...)
 		//
-		if !isContextWithCancel(pass.TypesInfo, n) || !isCall(stack[len(stack)-2]) {
+		if !isContextWithCancel(pass.TypesInfo, n) || !isCall(stack[len(stack)-1]) {
 			return true
 		}
 		var id *ast.Ident // id of cancel var
-		stmt := stack[len(stack)-3]
+		stmt := stack[len(stack)-2]
 		switch stmt := stmt.(type) {
 		case *ast.ValueSpec:
 			if len(stmt.Names) > 1 {

gopls/internal/golang/codeaction.go

@@ -713,33 +713,24 @@ func refactorRewriteEliminateDotImport(ctx context.Context, req *codeActionsRequ
 
 	// Go through each use of the dot imported package, checking its scope for
 	// shadowing and calculating an edit to qualify the identifier.
-	var stack []ast.Node
-	ast.Inspect(req.pgf.File, func(n ast.Node) bool {
-		if n == nil {
-			stack = stack[:len(stack)-1] // pop
-			return false
-		}
-		stack = append(stack, n) // push
+	for curId := range req.pgf.Cursor.Preorder((*ast.Ident)(nil)) {
+		ident := curId.Node().(*ast.Ident)
 
-		ident, ok := n.(*ast.Ident)
-		if !ok {
-			return true
-		}
 		// Only keep identifiers that use a symbol from the
 		// dot imported package.
 		use := req.pkg.TypesInfo().Uses[ident]
 		if use == nil || use.Pkg() == nil {
-			return true
+			continue
 		}
 		if use.Pkg() != imported {
-			return true
+			continue
 		}
 
 		// Only qualify unqualified identifiers (due to dot imports).
 		// All other references to a symbol imported from another package
 		// are nested within a select expression (pkg.Foo, v.Method, v.Field).
-		if is[*ast.SelectorExpr](stack[len(stack)-2]) {
-			return true
+		if is[*ast.SelectorExpr](curId.Parent().Node()) {
+			continue
 		}
 
 		// Make sure that the package name will not be shadowed by something else in scope.
@@ -750,24 +741,22 @@ func refactorRewriteEliminateDotImport(ctx context.Context, req *codeActionsRequ
 		// allowed to go through.
 		sc := fileScope.Innermost(ident.Pos())
 		if sc == nil {
-			return true
+			continue
 		}
 		_, obj := sc.LookupParent(newName, ident.Pos())
 		if obj != nil {
-			return true
+			continue
 		}
 
 		rng, err := req.pgf.PosRange(ident.Pos(), ident.Pos()) // sic, zero-width range before ident
 		if err != nil {
-			return true
+			continue
 		}
 		edits = append(edits, protocol.TextEdit{
 			Range:   rng,
 			NewText: newName + ".",
 		})
-
-		return true
-	})
+	}
 
 	req.addEditAction("Eliminate dot import", nil, protocol.DocumentChangeEdit(
 		req.fh,

gopls/internal/golang/hover.go

@@ -38,6 +38,7 @@ import (
 	gastutil "golang.org/x/tools/gopls/internal/util/astutil"
 	"golang.org/x/tools/gopls/internal/util/bug"
 	"golang.org/x/tools/gopls/internal/util/safetoken"
+	internalastutil "golang.org/x/tools/internal/astutil"
 	"golang.org/x/tools/internal/event"
 	"golang.org/x/tools/internal/stdlib"
 	"golang.org/x/tools/internal/tokeninternal"
@@ -1502,16 +1503,10 @@ func findDeclInfo(files []*ast.File, pos token.Pos) (decl ast.Decl, spec ast.Spe
 	stack := make([]ast.Node, 0, 20)
 
 	// Allocate the closure once, outside the loop.
-	f := func(n ast.Node) bool {
+	f := func(n ast.Node, stack []ast.Node) bool {
 		if found {
 			return false
 		}
-		if n != nil {
-			stack = append(stack, n) // push
-		} else {
-			stack = stack[:len(stack)-1] // pop
-			return false
-		}
 
 		// Skip subtrees (incl. files) that don't contain the search point.
 		if !(n.Pos() <= pos && pos < n.End()) {
@@ -1596,7 +1591,7 @@ func findDeclInfo(files []*ast.File, pos token.Pos) (decl ast.Decl, spec ast.Spe
 		return true
 	}
 	for _, file := range files {
-		ast.Inspect(file, f)
+		internalastutil.PreorderStack(file, stack, f)
 		if found {
 			return decl, spec, field
 		}

gopls/internal/golang/rename_check.go

@@ -45,6 +45,8 @@ import (
 	"golang.org/x/tools/go/ast/astutil"
 	"golang.org/x/tools/gopls/internal/cache"
 	"golang.org/x/tools/gopls/internal/util/safetoken"
+	"golang.org/x/tools/internal/astutil/cursor"
+	"golang.org/x/tools/internal/astutil/edge"
 	"golang.org/x/tools/internal/typeparams"
 	"golang.org/x/tools/internal/typesinternal"
 	"golang.org/x/tools/refactor/satisfy"
@@ -338,64 +340,61 @@ func deeper(x, y *types.Scope) bool {
 // lexical block enclosing the reference.  If fn returns false the
 // iteration is terminated and findLexicalRefs returns false.
 func forEachLexicalRef(pkg *cache.Package, obj types.Object, fn func(id *ast.Ident, block *types.Scope) bool) bool {
+	filter := []ast.Node{
+		(*ast.Ident)(nil),
+		(*ast.SelectorExpr)(nil),
+		(*ast.CompositeLit)(nil),
+	}
 	ok := true
-	var stack []ast.Node
-
-	var visit func(n ast.Node) bool
-	visit = func(n ast.Node) bool {
-		if n == nil {
-			stack = stack[:len(stack)-1] // pop
-			return false
-		}
+	var visit func(cur cursor.Cursor, push bool) (descend bool)
+	visit = func(cur cursor.Cursor, push bool) (descend bool) {
 		if !ok {
 			return false // bail out
 		}
-
-		stack = append(stack, n) // push
-		switch n := n.(type) {
+		if !push {
+			return false
+		}
+		switch n := cur.Node().(type) {
 		case *ast.Ident:
 			if pkg.TypesInfo().Uses[n] == obj {
-				block := enclosingBlock(pkg.TypesInfo(), stack)
+				block := enclosingBlock(pkg.TypesInfo(), cur)
 				if !fn(n, block) {
 					ok = false
 				}
 			}
-			return visit(nil) // pop stack
 
 		case *ast.SelectorExpr:
 			// don't visit n.Sel
-			ast.Inspect(n.X, visit)
-			return visit(nil) // pop stack, don't descend
+			cur.ChildAt(edge.SelectorExpr_X, -1).Inspect(filter, visit)
+			return false // don't descend
 
 		case *ast.CompositeLit:
 			// Handle recursion ourselves for struct literals
 			// so we don't visit field identifiers.
 			tv, ok := pkg.TypesInfo().Types[n]
 			if !ok {
-				return visit(nil) // pop stack, don't descend
+				return false // don't descend
 			}
 			if is[*types.Struct](typeparams.CoreType(typeparams.Deref(tv.Type))) {
 				if n.Type != nil {
-					ast.Inspect(n.Type, visit)
+					cur.ChildAt(edge.CompositeLit_Type, -1).Inspect(filter, visit)
 				}
-				for _, elt := range n.Elts {
-					if kv, ok := elt.(*ast.KeyValueExpr); ok {
-						ast.Inspect(kv.Value, visit)
-					} else {
-						ast.Inspect(elt, visit)
+				for i, elt := range n.Elts {
+					curElt := cur.ChildAt(edge.CompositeLit_Elts, i)
+					if _, ok := elt.(*ast.KeyValueExpr); ok {
+						// skip kv.Key
+						curElt = curElt.ChildAt(edge.KeyValueExpr_Value, -1)
 					}
+					curElt.Inspect(filter, visit)
 				}
-				return visit(nil) // pop stack, don't descend
+				return false // don't descend
 			}
 		}
 		return true
 	}
 
-	for _, f := range pkg.Syntax() {
-		ast.Inspect(f, visit)
-		if len(stack) != 0 {
-			panic(stack)
-		}
+	for _, pgf := range pkg.CompiledGoFiles() {
+		pgf.Cursor.Inspect(filter, visit)
 		if !ok {
 			break
 		}
@@ -404,11 +403,10 @@ func forEachLexicalRef(pkg *cache.Package, obj types.Object, fn func(id *ast.Ide
 }
 
 // enclosingBlock returns the innermost block logically enclosing the
-// specified AST node (an ast.Ident), specified in the form of a path
-// from the root of the file, [file...n].
-func enclosingBlock(info *types.Info, stack []ast.Node) *types.Scope {
-	for i := range stack {
-		n := stack[len(stack)-1-i]
+// AST node (an ast.Ident), specified as a Cursor.
+func enclosingBlock(info *types.Info, curId cursor.Cursor) *types.Scope {
+	for cur := range curId.Enclosing() {
+		n := cur.Node()
 		// For some reason, go/types always associates a
 		// function's scope with its FuncType.
 		// See comments about scope above.

internal/astutil/util.go

@@ -57,3 +57,35 @@ func PosInStringLiteral(lit *ast.BasicLit, offset int) (token.Pos, error) {
 	}
 	return pos, nil
 }
+
+// PreorderStack traverses the tree rooted at root,
+// calling f before visiting each node.
+//
+// Each call to f provides the current node and traversal stack,
+// consisting of the original value of stack appended with all nodes
+// from root to n, excluding n itself. (This design allows calls
+// to PreorderStack to be nested without double counting.)
+//
+// If f returns false, the traversal skips over that subtree. Unlike
+// [ast.Inspect], no second call to f is made after visiting node n.
+// In practice, the second call is nearly always used only to pop the
+// stack, and it is surprisingly tricky to do this correctly; see
+// https://go.dev/issue/73319.
+func PreorderStack(root ast.Node, stack []ast.Node, f func(n ast.Node, stack []ast.Node) bool) {
+	before := len(stack)
+	ast.Inspect(root, func(n ast.Node) bool {
+		if n != nil {
+			if !f(n, stack) {
+				// Do not push, as there will be no corresponding pop.
+				return false
+			}
+			stack = append(stack, n) // push
+		} else {
+			stack = stack[:len(stack)-1] // pop
+		}
+		return true
+	})
+	if len(stack) != before {
+		panic("push/pop mismatch")
+	}
+}

internal/astutil/util_test.go

@@ -0,0 +1,67 @@
+// Copyright 2025 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.
+
+package astutil_test
+
+import (
+	"fmt"
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"reflect"
+	"strings"
+	"testing"
+
+	"golang.org/x/tools/internal/astutil"
+)
+
+func TestPreorderStack(t *testing.T) {
+	const src = `package a
+func f() {
+	print("hello")
+}
+func g() {
+	print("goodbye")
+	panic("oops")
+}
+`
+	fset := token.NewFileSet()
+	f, _ := parser.ParseFile(fset, "a.go", src, 0)
+
+	str := func(n ast.Node) string {
+		return strings.TrimPrefix(reflect.TypeOf(n).String(), "*ast.")
+	}
+
+	var events []string
+	var gotStack []string
+	astutil.PreorderStack(f, nil, func(n ast.Node, stack []ast.Node) bool {
+		events = append(events, str(n))
+		if decl, ok := n.(*ast.FuncDecl); ok && decl.Name.Name == "f" {
+			return false // skip subtree of f()
+		}
+		if lit, ok := n.(*ast.BasicLit); ok && lit.Value == `"oops"` {
+			for _, n := range stack {
+				gotStack = append(gotStack, str(n))
+			}
+		}
+		return true
+	})
+
+	// Check sequence of events.
+	const wantEvents = `[File Ident ` + // package a
+		`FuncDecl ` + // func f()  [pruned]
+		`FuncDecl Ident FuncType FieldList BlockStmt ` + // func g()
+		`ExprStmt CallExpr Ident BasicLit ` + // print...
+		`ExprStmt CallExpr Ident BasicLit]` // panic...
+	if got := fmt.Sprint(events); got != wantEvents {
+		t.Errorf("PreorderStack events:\ngot:  %s\nwant: %s", got, wantEvents)
+	}
+
+	// Check captured stack.
+	const wantStack = `[File FuncDecl BlockStmt ExprStmt CallExpr]`
+	if got := fmt.Sprint(gotStack); got != wantStack {
+		t.Errorf("PreorderStack stack:\ngot:  %s\nwant: %s", got, wantStack)
+	}
+
+}