go/ast/inspector/inspector.go - tools.git - Git at Google

 // 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]
 	})
 }
	// 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]
	})
	}