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) on push event, or union of typ strictly between push and pop events on pop events 57 index int // index of corresponding push or pop event 58 } 59 60 // TODO: Experiment with storing only the second word of event.node (unsafe.Pointer). 61 // Type can be recovered from the sole bit in typ. 62 63 // Preorder visits all the nodes of the files supplied to New in 64 // depth-first order. It calls f(n) for each node n before it visits 65 // n's children. 66 // 67 // The complete traversal sequence is determined by ast.Inspect. 68 // The types argument, if non-empty, enables type-based filtering of 69 // events. The function f is called only for nodes whose type 70 // matches an element of the types slice. 71 func (in *Inspector) Preorder(types []ast.Node, f func(ast.Node)) { 72 // Because it avoids postorder calls to f, and the pruning 73 // check, Preorder is almost twice as fast as Nodes. The two 74 // features seem to contribute similar slowdowns (~1.4x each). 75 76 mask := maskOf(types) 77 for i := 0; i < len(in.events); { 78 ev := in.events[i] 79 if ev.index > i { 80 // push 81 if ev.typ&mask != 0 { 82 f(ev.node) 83 } 84 pop := ev.index 85 if in.events[pop].typ&mask == 0 { 86 // Subtrees do not contain types: skip them and pop. 87 i = pop + 1 88 continue 89 } 90 } 91 i++ 92 } 93 } 94 95 // Nodes visits the nodes of the files supplied to New in depth-first 96 // order. It calls f(n, true) for each node n before it visits n's 97 // children. If f returns true, Nodes invokes f recursively for each 98 // of the non-nil children of the node, followed by a call of 99 // f(n, false). 100 // 101 // The complete traversal sequence is determined by ast.Inspect. 102 // The types argument, if non-empty, enables type-based filtering of 103 // events. The function f if is called only for nodes whose type 104 // matches an element of the types slice. 105 func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proceed bool)) { 106 mask := maskOf(types) 107 for i := 0; i < len(in.events); { 108 ev := in.events[i] 109 if ev.index > i { 110 // push 111 pop := ev.index 112 if ev.typ&mask != 0 { 113 if !f(ev.node, true) { 114 i = pop + 1 // jump to corresponding pop + 1 115 continue 116 } 117 } 118 if in.events[pop].typ&mask == 0 { 119 // Subtrees do not contain types: skip them. 120 i = pop 121 continue 122 } 123 } else { 124 // pop 125 push := ev.index 126 if in.events[push].typ&mask != 0 { 127 f(ev.node, false) 128 } 129 } 130 i++ 131 } 132 } 133 134 // WithStack visits nodes in a similar manner to Nodes, but it 135 // supplies each call to f an additional argument, the current 136 // traversal stack. The stack's first element is the outermost node, 137 // an *ast.File; its last is the innermost, n. 138 func (in *Inspector) WithStack(types []ast.Node, f func(n ast.Node, push bool, stack []ast.Node) (proceed bool)) { 139 mask := maskOf(types) 140 var stack []ast.Node 141 for i := 0; i < len(in.events); { 142 ev := in.events[i] 143 if ev.index > i { 144 // push 145 pop := ev.index 146 stack = append(stack, ev.node) 147 if ev.typ&mask != 0 { 148 if !f(ev.node, true, stack) { 149 i = pop + 1 150 stack = stack[:len(stack)-1] 151 continue 152 } 153 } 154 if in.events[pop].typ&mask == 0 { 155 // Subtrees does not contain types: skip them. 156 i = pop 157 continue 158 } 159 } else { 160 // pop 161 push := ev.index 162 if in.events[push].typ&mask != 0 { 163 f(ev.node, false, stack) 164 } 165 stack = stack[:len(stack)-1] 166 } 167 i++ 168 } 169 } 170 171 // traverse builds the table of events representing a traversal. 172 func traverse(files []*ast.File) []event { 173 // Preallocate approximate number of events 174 // based on source file extent. 175 // This makes traverse faster by 4x (!). 176 var extent int 177 for _, f := range files { 178 extent += int(f.End() - f.Pos()) 179 } 180 // This estimate is based on the net/http package. 181 capacity := extent * 33 / 100 182 if capacity > 1e6 { 183 capacity = 1e6 // impose some reasonable maximum 184 } 185 events := make([]event, 0, capacity) 186 187 var stack []event 188 stack = append(stack, event{}) // include an extra event so file nodes have a parent 189 for _, f := range files { 190 ast.Inspect(f, func(n ast.Node) bool { 191 if n != nil { 192 // push 193 ev := event{ 194 node: n, 195 typ: 0, // temporarily used to accumulate type bits of subtree 196 index: len(events), // push event temporarily holds own index 197 } 198 stack = append(stack, ev) 199 events = append(events, ev) 200 } else { 201 // pop 202 top := len(stack) - 1 203 ev := stack[top] 204 typ := typeOf(ev.node) 205 push := ev.index 206 parent := top - 1 207 208 events[push].typ = typ // set type of push 209 stack[parent].typ |= typ | ev.typ // parent's typ contains push and pop's typs. 210 events[push].index = len(events) // make push refer to pop 211 212 stack = stack[:top] 213 events = append(events, ev) 214 } 215 return true 216 }) 217 } 218 219 return events 220 } 221