Source File
	inspector.go

Belonging Package
	golang.org/x/tools/go/ast/inspector

// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package inspector provides helper functions for traversal over the
// syntax trees of a package, including node filtering by type, and
// materialization of the traversal stack.
//
// During construction, the inspector does a complete traversal and
// builds a list of push/pop events and their node type. Subsequent
// method calls that request a traversal scan this list, rather than walk
// the AST, and perform type filtering using efficient bit sets.
// This representation is sometimes called a "balanced parenthesis tree."
//
// Experiments suggest the inspector's traversals are about 2.5x faster
// than [ast.Inspect], but it may take around 5 traversals for this
// benefit to amortize the inspector's construction cost.
// If efficiency is the primary concern, do not use Inspector for
// one-off traversals.
//
// The [Cursor] type provides a more flexible API for efficient
// navigation of syntax trees in all four "cardinal directions". For
// example, traversals may be nested, so you can find each node of
// type A and then search within it for nodes of type B. Or you can
// traverse from a node to its immediate neighbors: its parent, its
// previous and next sibling, or its first and last child. We
// recommend using methods of Cursor in preference to Inspector where
// possible.
package inspector

// There are four orthogonal features in a traversal:
//  1 type filtering
//  2 pruning
//  3 postorder calls to f
//  4 stack
// Rather than offer all of them in the API,
// only a few combinations are exposed:
// - Preorder is the fastest and has fewest features,
//   but is the most commonly needed traversal.
// - Nodes and WithStack both provide pruning and postorder calls,
//   even though few clients need it, because supporting two versions
//   is not justified.
// More combinations could be supported by expressing them as
// wrappers around a more generic traversal, but this was measured
// and found to degrade performance significantly (30%).

import (
	"go/ast"

	"golang.org/x/tools/go/ast/edge"
)

// An Inspector provides methods for inspecting
// (traversing) the syntax trees of a package.
type Inspector struct {
	events []event
}

func packEdgeKindAndIndex(ek edge.Kind, index int) int32 {
	return int32(uint32(index+1)<<7 | uint32(ek))
}

// unpackEdgeKindAndIndex unpacks the edge kind and edge index (within
// an []ast.Node slice) from the parent field of a pop event.
func unpackEdgeKindAndIndex(x int32) (edge.Kind, int) {
	// The "parent" field of a pop node holds the
	// edge Kind in the lower 7 bits and the index+1
	// in the upper 25.
	return edge.Kind(x & 0x7f), int(x>>7) - 1
}

// New returns an Inspector for the specified syntax trees.
func New(files []*ast.File) *Inspector {
	return &Inspector{traverse(files)}
}

// An event represents a push or a pop
// of an ast.Node during a traversal.
type event struct {
	node   ast.Node
	typ    uint64 // typeOf(node) on push event, or union of typ strictly between push and pop events on pop events
	index  int32  // index of corresponding push or pop event
	parent int32  // index of parent's push node (push nodes only), or packed edge kind/index (pop nodes only)
}

// TODO: Experiment with storing only the second word of event.node (unsafe.Pointer).
// Type can be recovered from the sole bit in typ.
// [Tried this, wasn't faster. --adonovan]

// Preorder visits all the nodes of the files supplied to [New] in
// depth-first order. It calls f(n) for each node n before it visits
// n's children.
//
// The complete traversal sequence is determined by [ast.Inspect].
// The types argument, if non-empty, enables type-based filtering of
// events. The function f is called only for nodes whose type
// matches an element of the types slice.
//
// The [Cursor.Preorder] method provides a richer alternative interface.
// Example:
//
//	for c := range in.Root().Preorder(types) { ... }
func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node)) {
	// Because it avoids postorder calls to f, and the pruning
	// check, Preorder is almost twice as fast as Nodes. The two
	// features seem to contribute similar slowdowns (~1.4x each).

	// This function is equivalent to the PreorderSeq call below,
	// but to avoid the additional dynamic call (which adds 13-35%
	// to the benchmarks), we expand it out.
	//
	// in.PreorderSeq(types...)(func(n ast.Node) bool {
	// 	f(n)
	// 	return true
	// })

	mask := maskOf(types)
	for i := int32(0); i < int32(len(in.events)); {
		ev := in.events[i]
		if ev.index > i {
			// push
			if ev.typ&mask != 0 {
				f(ev.node)
			}
			pop := ev.index
			if in.events[pop].typ&mask == 0 {
				// Subtrees do not contain types: skip them and pop.
				i = pop + 1
				continue
			}
		}
		i++
	}
}

// Nodes visits the nodes of the files supplied to [New] in depth-first
// order. It calls f(n, true) for each node n before it visits n's
// children. If f returns true, Nodes invokes f recursively for each
// of the non-nil children of the node, followed by a call of
// f(n, false).
//
// The complete traversal sequence is determined by [ast.Inspect].
// The types argument, if non-empty, enables type-based filtering of
// events. The function f if is called only for nodes whose type
// matches an element of the types slice.
//
// The [Cursor.Inspect] method provides a richer alternative interface.
// Example:
//
//	in.Root().Inspect(types, func(c Cursor) bool {
//		...
//		return true
//	}
func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proceed bool)) {
	mask := maskOf(types)
	for i := int32(0); i < int32(len(in.events)); {
		ev := in.events[i]
		if ev.index > i {
			// push
			pop := ev.index
			if ev.typ&mask != 0 {
				if !f(ev.node, true) {
					i = pop + 1 // jump to corresponding pop + 1
					continue
				}
			}
			if in.events[pop].typ&mask == 0 {
				// Subtrees do not contain types: skip them.
				i = pop
				continue
			}
		} else {
			// pop
			push := ev.index
			if in.events[push].typ&mask != 0 {
				f(ev.node, false)
			}
		}
		i++
	}
}

// WithStack visits nodes in a similar manner to Nodes, but it
// supplies each call to f an additional argument, the current
// traversal stack. The stack's first element is the outermost node,
// an *ast.File; its last is the innermost, n.
//
// The [Cursor.Inspect] method provides a richer alternative interface.
// Example:
//
//	in.Root().Inspect(types, func(c Cursor) bool {
//		stack := slices.Collect(c.Enclosing())
//		...
//		return true
//	})
func (in *Inspector) WithStack(types []ast.Node, f func(n ast.Node, push bool, stack []ast.Node) (proceed bool)) {
	mask := maskOf(types)
	var stack []ast.Node
	for i := int32(0); i < int32(len(in.events)); {
		ev := in.events[i]
		if ev.index > i {
			// push
			pop := ev.index
			stack = append(stack, ev.node)
			if ev.typ&mask != 0 {
				if !f(ev.node, true, stack) {
					i = pop + 1
					stack = stack[:len(stack)-1]
					continue
				}
			}
			if in.events[pop].typ&mask == 0 {
				// Subtrees does not contain types: skip them.
				i = pop
				continue
			}
		} else {
			// pop
			push := ev.index
			if in.events[push].typ&mask != 0 {
				f(ev.node, false, stack)
			}
			stack = stack[:len(stack)-1]
		}
		i++
	}
}

// traverse builds the table of events representing a traversal.
func traverse(files []*ast.File) []event {
	// Preallocate approximate number of events
	// based on source file extent of the declarations.
	// (We use End-Pos not FileStart-FileEnd to neglect
	// the effect of long doc comments.)
	// This makes traverse faster by 4x (!).
	var extent int
	for _, f := range files {
		extent += int(f.End() - f.Pos())
	}
	// This estimate is based on the net/http package.
	capacity := min(extent*33/100, 1e6) // impose some reasonable maximum (1M)

	v := &visitor{
		events: make([]event, 0, capacity),
		stack:  []item{{index: -1}}, // include an extra event so file nodes have a parent
	}
	for _, file := range files {
		walk(v, edge.Invalid, -1, file)
	}
	return v.events
}

type visitor struct {
	events []event
	stack  []item
}

type item struct {
	index            int32  // index of current node's push event
	parentIndex      int32  // index of parent node's push event
	typAccum         uint64 // accumulated type bits of current node's descendants
	edgeKindAndIndex int32  // edge.Kind and index, bit packed
}

func (v *visitor) push(ek edge.Kind, eindex int, node ast.Node) {
	var (
		index       = int32(len(v.events))
		parentIndex = v.stack[len(v.stack)-1].index
	)
	v.events = append(v.events, event{
		node:   node,
		parent: parentIndex,
		typ:    typeOf(node),
		index:  0, // (pop index is set later by visitor.pop)
	})
	v.stack = append(v.stack, item{
		index:            index,
		parentIndex:      parentIndex,
		edgeKindAndIndex: packEdgeKindAndIndex(ek, eindex),
	})

	// 2B nodes ought to be enough for anyone!
	if int32(len(v.events)) < 0 {
		panic("event index exceeded int32")
	}

	// 32M elements in an []ast.Node ought to be enough for anyone!
	if ek2, eindex2 := unpackEdgeKindAndIndex(packEdgeKindAndIndex(ek, eindex)); ek2 != ek || eindex2 != eindex {
		panic("Node slice index exceeded uint25")
	}
}

func (v *visitor) pop(node ast.Node) {
	top := len(v.stack) - 1
	current := v.stack[top]

	push := &v.events[current.index]
	parent := &v.stack[top-1]

	push.index = int32(len(v.events))              // make push event refer to pop
	parent.typAccum |= current.typAccum | push.typ // accumulate type bits into parent

	v.stack = v.stack[:top]

	v.events = append(v.events, event{
		node:   node,
		typ:    current.typAccum,
		index:  current.index,
		parent: current.edgeKindAndIndex, // see [unpackEdgeKindAndIndex]
	})
}

The pages are generated with Golds v0.8.4. (GOOS=linux GOARCH=amd64)