From: Alan Donovan Date: Wed, 14 May 2025 20:15:40 +0000 (-0400) Subject: go/ast: add PreorderStack, a variant of Inspect that builds a stack X-Git-Tag: go1.25rc1~236 X-Git-Url: http://www.git.cypherpunks.su/?a=commitdiff_plain;h=5afada035ced1f89267d3177a9fb75fab4df81ff;p=gostls13.git 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 LUCI-TryBot-Result: Go LUCI Auto-Submit: Alan Donovan --- diff --git a/api/next/73319.txt b/api/next/73319.txt new file mode 100644 index 0000000000..39a3ece5f8 --- /dev/null +++ b/api/next/73319.txt @@ -0,0 +1 @@ +pkg go/ast, func PreorderStack(Node, []Node, func(Node, []Node) bool) #73319 diff --git a/doc/next/6-stdlib/99-minor/go/ast/73319.md b/doc/next/6-stdlib/99-minor/go/ast/73319.md new file mode 100644 index 0000000000..b99e20e316 --- /dev/null +++ b/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. diff --git a/src/go/ast/walk.go b/src/go/ast/walk.go index ec9a8901c4..24cdc60d73 100644 --- a/src/go/ast/walk.go +++ b/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") + } +} diff --git a/src/go/ast/walk_test.go b/src/go/ast/walk_test.go index b8b4f958ec..172b2e3f5d 100644 --- a/src/go/ast/walk_test.go +++ b/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) + } +}