1 // Copyright 2018 The Go Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style
3 // license that can be found in the LICENSE file.
5 // Package inspector provides helper functions for traversal over the
6 // syntax trees of a package, including node filtering by type, and
7 // materialization of the traversal stack.
9 // During construction, the inspector does a complete traversal and
10 // builds a list of push/pop events and their node type. Subsequent
11 // method calls that request a traversal scan this list, rather than walk
12 // the AST, and perform type filtering using efficient bit sets.
14 // Experiments suggest the inspector's traversals are about 2.5x faster
15 // than ast.Inspect, but it may take around 5 traversals for this
16 // benefit to amortize the inspector's construction cost.
17 // If efficiency is the primary concern, do not use Inspector for
18 // one-off traversals.
21 // There are four orthogonal features in a traversal:
24 // 3 postorder calls to f
26 // Rather than offer all of them in the API,
27 // only a few combinations are exposed:
28 // - Preorder is the fastest and has fewest features,
29 // but is the most commonly needed traversal.
30 // - Nodes and WithStack both provide pruning and postorder calls,
31 // even though few clients need it, because supporting two versions
33 // More combinations could be supported by expressing them as
34 // wrappers around a more generic traversal, but this was measured
35 // and found to degrade performance significantly (30%).
41 // An Inspector provides methods for inspecting
42 // (traversing) the syntax trees of a package.
43 type Inspector struct {
47 // New returns an Inspector for the specified syntax trees.
48 func New(files []*ast.File) *Inspector {
49 return &Inspector{traverse(files)}
52 // An event represents a push or a pop
53 // of an ast.Node during a traversal.
56 typ uint64 // typeOf(node)
57 index int // 1 + index of corresponding pop event, or 0 if this is a pop
60 // Preorder visits all the nodes of the files supplied to New in
61 // depth-first order. It calls f(n) for each node n before it visits
64 // The types argument, if non-empty, enables type-based filtering of
65 // events. The function f if is called only for nodes whose type
66 // matches an element of the types slice.
67 func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node)) {
68 // Because it avoids postorder calls to f, and the pruning
69 // check, Preorder is almost twice as fast as Nodes. The two
70 // features seem to contribute similar slowdowns (~1.4x each).
73 for i := 0; i < len(in.events); {
84 // Nodes visits the nodes of the files supplied to New in depth-first
85 // order. It calls f(n, true) for each node n before it visits n's
86 // children. If f returns true, Nodes invokes f recursively for each
87 // of the non-nil children of the node, followed by a call of
90 // The types argument, if non-empty, enables type-based filtering of
91 // events. The function f if is called only for nodes whose type
92 // matches an element of the types slice.
93 func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proceed bool)) {
95 for i := 0; i < len(in.events); {
100 if !f(ev.node, true) {
101 i = ev.index // jump to corresponding pop + 1
113 // WithStack visits nodes in a similar manner to Nodes, but it
114 // supplies each call to f an additional argument, the current
115 // traversal stack. The stack's first element is the outermost node,
116 // an *ast.File; its last is the innermost, n.
117 func (in *Inspector) WithStack(types []ast.Node, f func(n ast.Node, push bool, stack []ast.Node) (proceed bool)) {
118 mask := maskOf(types)
120 for i := 0; i < len(in.events); {
124 stack = append(stack, ev.node)
125 if ev.typ&mask != 0 {
126 if !f(ev.node, true, stack) {
128 stack = stack[:len(stack)-1]
134 if ev.typ&mask != 0 {
135 f(ev.node, false, stack)
137 stack = stack[:len(stack)-1]
143 // traverse builds the table of events representing a traversal.
144 func traverse(files []*ast.File) []event {
145 // Preallocate approximate number of events
146 // based on source file extent.
147 // This makes traverse faster by 4x (!).
149 for _, f := range files {
150 extent += int(f.End() - f.Pos())
152 // This estimate is based on the net/http package.
153 capacity := extent * 33 / 100
155 capacity = 1e6 // impose some reasonable maximum
157 events := make([]event, 0, capacity)
160 for _, f := range files {
161 ast.Inspect(f, func(n ast.Node) bool {
167 index: len(events), // push event temporarily holds own index
169 stack = append(stack, ev)
170 events = append(events, ev)
173 ev := stack[len(stack)-1]
174 stack = stack[:len(stack)-1]
176 events[ev.index].index = len(events) + 1 // make push refer to pop
178 ev.index = 0 // turn ev into a pop event
179 events = append(events, ev)