| 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. |
| 4 | |
| 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. |
| 8 | // |
| 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. |
| 13 | // |
| 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. |
| 19 | package inspector |
| 20 | |
| 21 | // There are four orthogonal features in a traversal: |
| 22 | // 1 type filtering |
| 23 | // 2 pruning |
| 24 | // 3 postorder calls to f |
| 25 | // 4 stack |
| 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 |
| 32 | // is not justified. |
| 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%). |
| 36 | |
| 37 | import ( |
| 38 | "go/ast" |
| 39 | ) |
| 40 | |
| 41 | // An Inspector provides methods for inspecting |
| 42 | // (traversing) the syntax trees of a package. |
| 43 | type Inspector struct { |
| 44 | events []event |
| 45 | } |
| 46 | |
| 47 | // New returns an Inspector for the specified syntax trees. |
| 48 | func New(files []*ast.File) *Inspector { |
| 49 | return &Inspector{traverse(files)} |
| 50 | } |
| 51 | |
| 52 | // An event represents a push or a pop |
| 53 | // of an ast.Node during a traversal. |
| 54 | type event struct { |
| 55 | node ast.Node |
| 56 | typ uint64 // typeOf(node) |
| 57 | index int // 1 + index of corresponding pop event, or 0 if this is a pop |
| 58 | } |
| 59 | |
| 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 |
| 62 | // n's children. |
| 63 | // |
| 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). |
| 71 | |
| 72 | mask := maskOf(types) |
| 73 | for i := 0; i < len(in.events); { |
| 74 | ev := in.events[i] |
| 75 | if ev.typ&mask != 0 { |
| 76 | if ev.index > 0 { |
| 77 | f(ev.node) |
| 78 | } |
| 79 | } |
| 80 | i++ |
| 81 | } |
| 82 | } |
| 83 | |
| 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 |
| 88 | // f(n, false). |
| 89 | // |
| 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)) { |
| 94 | mask := maskOf(types) |
| 95 | for i := 0; i < len(in.events); { |
| 96 | ev := in.events[i] |
| 97 | if ev.typ&mask != 0 { |
| 98 | if ev.index > 0 { |
| 99 | // push |
| 100 | if !f(ev.node, true) { |
| 101 | i = ev.index // jump to corresponding pop + 1 |
| 102 | continue |
| 103 | } |
| 104 | } else { |
| 105 | // pop |
| 106 | f(ev.node, false) |
| 107 | } |
| 108 | } |
| 109 | i++ |
| 110 | } |
| 111 | } |
| 112 | |
| 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) |
| 119 | var stack []ast.Node |
| 120 | for i := 0; i < len(in.events); { |
| 121 | ev := in.events[i] |
| 122 | if ev.index > 0 { |
| 123 | // push |
| 124 | stack = append(stack, ev.node) |
| 125 | if ev.typ&mask != 0 { |
| 126 | if !f(ev.node, true, stack) { |
| 127 | i = ev.index |
| 128 | stack = stack[:len(stack)-1] |
| 129 | continue |
| 130 | } |
| 131 | } |
| 132 | } else { |
| 133 | // pop |
| 134 | if ev.typ&mask != 0 { |
| 135 | f(ev.node, false, stack) |
| 136 | } |
| 137 | stack = stack[:len(stack)-1] |
| 138 | } |
| 139 | i++ |
| 140 | } |
| 141 | } |
| 142 | |
| 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 (!). |
| 148 | var extent int |
| 149 | for _, f := range files { |
| 150 | extent += int(f.End() - f.Pos()) |
| 151 | } |
| 152 | // This estimate is based on the net/http package. |
| 153 | capacity := extent * 33 / 100 |
| 154 | if capacity > 1e6 { |
| 155 | capacity = 1e6 // impose some reasonable maximum |
| 156 | } |
| 157 | events := make([]event, 0, capacity) |
| 158 | |
| 159 | var stack []event |
| 160 | for _, f := range files { |
| 161 | ast.Inspect(f, func(n ast.Node) bool { |
| 162 | if n != nil { |
| 163 | // push |
| 164 | ev := event{ |
| 165 | node: n, |
| 166 | typ: typeOf(n), |
| 167 | index: len(events), // push event temporarily holds own index |
| 168 | } |
| 169 | stack = append(stack, ev) |
| 170 | events = append(events, ev) |
| 171 | } else { |
| 172 | // pop |
| 173 | ev := stack[len(stack)-1] |
| 174 | stack = stack[:len(stack)-1] |
| 175 | |
| 176 | events[ev.index].index = len(events) + 1 // make push refer to pop |
| 177 | |
| 178 | ev.index = 0 // turn ev into a pop event |
| 179 | events = append(events, ev) |
| 180 | } |
| 181 | return true |
| 182 | }) |
| 183 | } |
| 184 | |
| 185 | return events |
| 186 | } |
| 187 |
Members