proposal: x/tools/go/ast/inspector: add Cursor, to enable partial and multi-level traversals

[adonovan]

Background: The inspector package enables efficient repeated traversal over the list of ast.Files that make up a package. However, unlike ast.Inspect, it provides no way to inspect just some subtree, which is a common and recurring need. In particular, it is common to need a two-level traversal, in which one first inspects to find each node of type A, then conditionally visits some subtrees of it looking for nodes of type B.

This need could be addressed if the inspector could iterate over a sequence of abstract cursor positions from which the actual ast.Node could be queried. A cursor can then be used as the basis for a new traversal of the subtree rooted at that node, with a different type filter.

As a bonus, the WithStack operation would not be needed because the stack can be derived relatively efficiently from the cursor position by walking back from the current point, skipping from pop to push nodes.

Unfortunately the inspector as defined is stateless, so we can't simply add a Current() Cursor method to it, nor sneak a Cursor among the arguments to the existing callback functions. Cursor would need all the same visitation methods as Inspector, and Inspector's methods would conceptually delegate to the "root" Cursor. However, we can do better than simply duplicating the old API.

Proposal: We propose to add the Cursor type and related methods:

package inspect

type Inspector

// Conceptually the first three of these existing methods could
// become wrappers around in.Cursor().Preorder,
// but for efficiency they should remain independent.

func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proceed bool)) 
func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node))
func (in *Inspector) PreorderSeq(types ...ast.Node) iter.Seq[ast.Node]
func (in *Inspector) WithStack(types []ast.Node, f func(n ast.Node, push bool, stack []ast.Node) (proceed bool))

// New declarations:

// A Cursor represents an ast.Node. It is immutable.
type Cursor struct {
    in *Inspector
    index int
}    

// Root returns a cursor for the virtual root node,
// whose children are the ast.Files provided to NewInspector.
//
// Its Node and Stack methods returns nil.
func (in *Inspector) Root() Cursor

// Node returns the node at the current cursor position.
func (Cursor) Node() ast.Node

// Parent returns the parent of the current node.
// It may return the root node (whose Cursor.Node methods returns nil).
func (Cursor) Parent() Cursor

// Stack returns the stack of enclosing nodes,
// from the ast.File down to the current cursor's node.
func (Cursor) Stack() []Cursor {
   // The implementation would construct the stack on demand
   // by skip-walking up the event tree. (This is probably
   // more efficient than the current implementation!)
}

// Preorder returns a sequence of cursor positions as the nodes
// in the tree are visited in preorder.
func (Cursor) Preorder(types []ast.Node) iter.Seq[Cursor]

// We could add a pre+post order variant like WithStack, but in practice the `push bool` is rarely needed.

Example usage:

var (
   filterA = []ast.Node{(*ast.File)(nil)}
   filterB = []ast.Node{(*ast.BinaryExpr)(nil)}
)


var in *inspect.Inspector = ...

for curFile:= range in.Cursor().Preorder(filterA) {
    file := curFile.Node().(*ast.File)
    ...
    for curBinary := range c.Preorder(filterB) {
        binary := curBinary.Node().(*ast.BinaryExpr)
	...c.Stack()...
    }
}

See https://go.dev/cl/636156 for the implementation.

@findleyr @timothy-king @griesemer

[gabyhelp]

Related Issues

• x/tools/go/ast/inspector: add Inspector.PreorderSeq and top-level func All #67795 (closed)
• go/ast: add Preorder iterator #66339 (closed)

_{(Emoji vote if this was helpful or unhelpful; more detailed feedback welcome in this discussion.)}

[gopherbot]

Change https://go.dev/cl/636156 mentions this issue: go/ast/inspector: Cursor: partial and multi-level traversals.

[gopherbot]

Change https://go.dev/cl/636656 mentions this issue: gopls/internal/analysis/inspect: hack for Cursor access

[findleyr]

Summarizing discussion from the CL: we're a little concerned about Cursor methods with suboptimal performance, since the Inspector itself was created as an optimization. Specifically:

• The fact that Cursor.Preorder does not allow for pruning branches of the traversal may unsatisfactory. The advantage of a Cursor is that it allows for nested traversal, and in such cases you almost always want to prune the parent traversal.
• Using Stack is not as fast as a WithStack traversal, at least in naive benchmarks. In practice it may still be better, if most iterations of the body don't access the stack.

We are going to use the Cursor internally to see which API is the most useful.