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