--- /dev/null
+pkg go/ast, func PreorderStack(Node, []Node, func(Node, []Node) bool) #73319
--- /dev/null
+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.
// 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)
}
// 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
})
}
}
+
+// 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")
+ }
+}
"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
}
}
}
+
+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)
+ }
+}