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

// -- existing API -- 

type Inspector

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

// -- new API --

// 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 (*Inspector) Root() Cursor

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

// String returns a description of the current node.
func (Cursor) String() string

// -- traversals --

// All the usual traversals over the inspector can
// be offered as traversals within a subtree rooted at a given
// node, allowing multi-level traversals with independent
// control over type-based node filtering and independent control
// over pruning descent into subtrees. We can also provide fast
// means to obtain the cursor for a particular node or position.

func (Cursor) Preorder(types ...ast.Node) iter.Seq[Cursor]
func (Cursor) Inspect(types []ast.Node, f func(c Cursor, push bool) (descend bool))

// -- navigation in the four cardinal directions --

// 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.
// It appends to the provided slice, which must be empty,
// but may have spare capacity.
func (Cursor) Stack([]Cursor) []Cursor

func (Cursor) FirstChild() (Cursor, bool)
func (Cursor) LastChild() (Cursor, bool)
func (Cursor) Children() iter.Seq[Cursor]

func (Cursor) NextSibling() (Cursor, bool)
func (Cursor) PrevSibling() (Cursor, bool)

// -- search --

func (c Cursor) FindNode(n ast.Node) (Cursor, bool)
func (c Cursor) FindPos(start, end token.Pos) (Cursor, bool)

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)
	fmt.Println(c.Stack(nil))
    }
}

See https://go.dev/cl/636656 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.