go/ast: add PreorderStack, a variant of Inspect that builds a stack

+ doc, test, relnote

Fixes #73319

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

api/next/73319.txt

@@ -0,0 +1 @@
+pkg go/ast, func PreorderStack(Node, []Node, func(Node, []Node) bool) #73319

doc/next/6-stdlib/99-minor/go/ast/73319.md

@@ -0,0 +1,4 @@
+The new [PreorderStack] function, like [Inspect], traverses a syntax
+tree and provides control over descent into subtrees, but as a
+convenience it also provides the stack of enclosing nodes at each
+point.

src/go/ast/walk.go

@@ -368,6 +368,11 @@ func (f inspector) Visit(node Node) Visitor {
 // f(node); node must not be nil. If f returns true, Inspect invokes f
 // recursively for each of the non-nil children of node, followed by a
 // call of f(nil).
+//
+// In many cases it may be more convenient to use [Preorder], which
+// returns an iterator over the sqeuence of nodes, or [PreorderStack],
+// which (like [Inspect]) provides control over descent into subtrees,
+// but additionally reports the stack of enclosing nodes.
 func Inspect(node Node, f func(Node) bool) {
 	Walk(inspector(f), node)
 }
@@ -376,7 +381,8 @@ func Inspect(node Node, f func(Node) bool) {
 // beneath (and including) the specified root, in depth-first
 // preorder.
 //
-// For greater control over the traversal of each subtree, use [Inspect].
+// For greater control over the traversal of each subtree, use
+// [Inspect] or [PreorderStack].
 func Preorder(root Node) iter.Seq[Node] {
 	return func(yield func(Node) bool) {
 		ok := true
@@ -389,3 +395,34 @@ func Preorder(root Node) iter.Seq[Node] {
 		})
 	}
 }
+
+// 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
+// [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.)
+func PreorderStack(root Node, stack []Node, f func(n Node, stack []Node) bool) {
+	before := len(stack)
+	Inspect(root, func(n 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")
+	}
+}

src/go/ast/walk_test.go

@@ -8,10 +8,13 @@ import (
 	"go/ast"
 	"go/parser"
 	"go/token"
+	"reflect"
+	"slices"
+	"strings"
 	"testing"
 )
 
-func TestPreorderBreak(t *testing.T) {
+func TestPreorder_Break(t *testing.T) {
 	// This test checks that Preorder correctly handles a break statement while
 	// in the middle of walking a node. Previously, incorrect handling of the
 	// boolean returned by the yield function resulted in the iterator calling
@@ -31,3 +34,54 @@ func TestPreorderBreak(t *testing.T) {
 		}
 	}
 }
+
+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
+	ast.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.
+	wantEvents := []string{
+		"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 !slices.Equal(events, wantEvents) {
+		t.Errorf("PreorderStack events:\ngot:  %s\nwant: %s", events, wantEvents)
+	}
+
+	// Check captured stack.
+	wantStack := []string{"File", "FuncDecl", "BlockStmt", "ExprStmt", "CallExpr"}
+	if !slices.Equal(gotStack, wantStack) {
+		t.Errorf("PreorderStack stack:\ngot:  %s\nwant: %s", gotStack, wantStack)
+	}
+}