1 // Copyright 2017 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.
14 // An ApplyFunc is invoked by Apply for each node n, even if n is nil,
15 // before and/or after the node's children, using a Cursor describing
16 // the current node and providing operations on it.
18 // The return value of ApplyFunc controls the syntax tree traversal.
19 // See Apply for details.
20 type ApplyFunc func(*Cursor) bool
22 // Apply traverses a syntax tree recursively, starting with root,
23 // and calling pre and post for each node as described below.
24 // Apply returns the syntax tree, possibly modified.
26 // If pre is not nil, it is called for each node before the node's
27 // children are traversed (pre-order). If pre returns false, no
28 // children are traversed, and post is not called for that node.
30 // If post is not nil, and a prior call of pre didn't return false,
31 // post is called for each node after its children are traversed
32 // (post-order). If post returns false, traversal is terminated and
33 // Apply returns immediately.
35 // Only fields that refer to AST nodes are considered children;
36 // i.e., token.Pos, Scopes, Objects, and fields of basic types
37 // (strings, etc.) are ignored.
39 // Children are traversed in the order in which they appear in the
40 // respective node's struct definition. A package's files are
41 // traversed in the filenames' alphabetical order.
43 func Apply(root ast.Node, pre, post ApplyFunc) (result ast.Node) {
44 parent := &struct{ ast.Node }{root}
46 if r := recover(); r != nil && r != abort {
51 a := &application{pre: pre, post: post}
52 a.apply(parent, "Node", nil, root)
56 var abort = new(int) // singleton, to signal termination of Apply
58 // A Cursor describes a node encountered during Apply.
59 // Information about the node and its parent is available
60 // from the Node, Parent, Name, and Index methods.
62 // If p is a variable of type and value of the current parent node
63 // c.Parent(), and f is the field identifier with name c.Name(),
64 // the following invariants hold:
66 // p.f == c.Node() if c.Index() < 0
67 // p.f[c.Index()] == c.Node() if c.Index() >= 0
69 // The methods Replace, Delete, InsertBefore, and InsertAfter
70 // can be used to change the AST without disrupting Apply.
74 iter *iterator // valid if non-nil
78 // Node returns the current Node.
79 func (c *Cursor) Node() ast.Node { return c.node }
81 // Parent returns the parent of the current Node.
82 func (c *Cursor) Parent() ast.Node { return c.parent }
84 // Name returns the name of the parent Node field that contains the current Node.
85 // If the parent is a *ast.Package and the current Node is a *ast.File, Name returns
86 // the filename for the current Node.
87 func (c *Cursor) Name() string { return c.name }
89 // Index reports the index >= 0 of the current Node in the slice of Nodes that
90 // contains it, or a value < 0 if the current Node is not part of a slice.
91 // The index of the current node changes if InsertBefore is called while
92 // processing the current node.
93 func (c *Cursor) Index() int {
100 // field returns the current node's parent field value.
101 func (c *Cursor) field() reflect.Value {
102 return reflect.Indirect(reflect.ValueOf(c.parent)).FieldByName(c.name)
105 // Replace replaces the current Node with n.
106 // The replacement node is not walked by Apply.
107 func (c *Cursor) Replace(n ast.Node) {
108 if _, ok := c.node.(*ast.File); ok {
109 file, ok := n.(*ast.File)
111 panic("attempt to replace *ast.File with non-*ast.File")
113 c.parent.(*ast.Package).Files[c.name] = file
118 if i := c.Index(); i >= 0 {
121 v.Set(reflect.ValueOf(n))
124 // Delete deletes the current Node from its containing slice.
125 // If the current Node is not part of a slice, Delete panics.
126 // As a special case, if the current node is a package file,
127 // Delete removes it from the package's Files map.
128 func (c *Cursor) Delete() {
129 if _, ok := c.node.(*ast.File); ok {
130 delete(c.parent.(*ast.Package).Files, c.name)
136 panic("Delete node not contained in slice")
140 reflect.Copy(v.Slice(i, l), v.Slice(i+1, l))
141 v.Index(l - 1).Set(reflect.Zero(v.Type().Elem()))
146 // InsertAfter inserts n after the current Node in its containing slice.
147 // If the current Node is not part of a slice, InsertAfter panics.
148 // Apply does not walk n.
149 func (c *Cursor) InsertAfter(n ast.Node) {
152 panic("InsertAfter node not contained in slice")
155 v.Set(reflect.Append(v, reflect.Zero(v.Type().Elem())))
157 reflect.Copy(v.Slice(i+2, l), v.Slice(i+1, l))
158 v.Index(i + 1).Set(reflect.ValueOf(n))
162 // InsertBefore inserts n before the current Node in its containing slice.
163 // If the current Node is not part of a slice, InsertBefore panics.
164 // Apply will not walk n.
165 func (c *Cursor) InsertBefore(n ast.Node) {
168 panic("InsertBefore node not contained in slice")
171 v.Set(reflect.Append(v, reflect.Zero(v.Type().Elem())))
173 reflect.Copy(v.Slice(i+1, l), v.Slice(i, l))
174 v.Index(i).Set(reflect.ValueOf(n))
178 // application carries all the shared data so we can pass it around cheaply.
179 type application struct {
185 func (a *application) apply(parent ast.Node, name string, iter *iterator, n ast.Node) {
186 // convert typed nil into untyped nil
187 if v := reflect.ValueOf(n); v.Kind() == reflect.Ptr && v.IsNil() {
191 // avoid heap-allocating a new cursor for each apply call; reuse a.cursor instead
193 a.cursor.parent = parent
198 if a.pre != nil && !a.pre(&a.cursor) {
204 // (the order of the cases matches the order of the corresponding node types in go/ast)
205 switch n := n.(type) {
209 // Comments and fields
213 case *ast.CommentGroup:
215 a.applyList(n, "List")
219 a.apply(n, "Doc", nil, n.Doc)
220 a.applyList(n, "Names")
221 a.apply(n, "Type", nil, n.Type)
222 a.apply(n, "Tag", nil, n.Tag)
223 a.apply(n, "Comment", nil, n.Comment)
226 a.applyList(n, "List")
229 case *ast.BadExpr, *ast.Ident, *ast.BasicLit:
233 a.apply(n, "Elt", nil, n.Elt)
236 a.apply(n, "Type", nil, n.Type)
237 a.apply(n, "Body", nil, n.Body)
239 case *ast.CompositeLit:
240 a.apply(n, "Type", nil, n.Type)
241 a.applyList(n, "Elts")
244 a.apply(n, "X", nil, n.X)
246 case *ast.SelectorExpr:
247 a.apply(n, "X", nil, n.X)
248 a.apply(n, "Sel", nil, n.Sel)
251 a.apply(n, "X", nil, n.X)
252 a.apply(n, "Index", nil, n.Index)
255 a.apply(n, "X", nil, n.X)
256 a.apply(n, "Low", nil, n.Low)
257 a.apply(n, "High", nil, n.High)
258 a.apply(n, "Max", nil, n.Max)
260 case *ast.TypeAssertExpr:
261 a.apply(n, "X", nil, n.X)
262 a.apply(n, "Type", nil, n.Type)
265 a.apply(n, "Fun", nil, n.Fun)
266 a.applyList(n, "Args")
269 a.apply(n, "X", nil, n.X)
272 a.apply(n, "X", nil, n.X)
274 case *ast.BinaryExpr:
275 a.apply(n, "X", nil, n.X)
276 a.apply(n, "Y", nil, n.Y)
278 case *ast.KeyValueExpr:
279 a.apply(n, "Key", nil, n.Key)
280 a.apply(n, "Value", nil, n.Value)
284 a.apply(n, "Len", nil, n.Len)
285 a.apply(n, "Elt", nil, n.Elt)
287 case *ast.StructType:
288 a.apply(n, "Fields", nil, n.Fields)
291 a.apply(n, "Params", nil, n.Params)
292 a.apply(n, "Results", nil, n.Results)
294 case *ast.InterfaceType:
295 a.apply(n, "Methods", nil, n.Methods)
298 a.apply(n, "Key", nil, n.Key)
299 a.apply(n, "Value", nil, n.Value)
302 a.apply(n, "Value", nil, n.Value)
309 a.apply(n, "Decl", nil, n.Decl)
314 case *ast.LabeledStmt:
315 a.apply(n, "Label", nil, n.Label)
316 a.apply(n, "Stmt", nil, n.Stmt)
319 a.apply(n, "X", nil, n.X)
322 a.apply(n, "Chan", nil, n.Chan)
323 a.apply(n, "Value", nil, n.Value)
325 case *ast.IncDecStmt:
326 a.apply(n, "X", nil, n.X)
328 case *ast.AssignStmt:
329 a.applyList(n, "Lhs")
330 a.applyList(n, "Rhs")
333 a.apply(n, "Call", nil, n.Call)
336 a.apply(n, "Call", nil, n.Call)
338 case *ast.ReturnStmt:
339 a.applyList(n, "Results")
341 case *ast.BranchStmt:
342 a.apply(n, "Label", nil, n.Label)
345 a.applyList(n, "List")
348 a.apply(n, "Init", nil, n.Init)
349 a.apply(n, "Cond", nil, n.Cond)
350 a.apply(n, "Body", nil, n.Body)
351 a.apply(n, "Else", nil, n.Else)
353 case *ast.CaseClause:
354 a.applyList(n, "List")
355 a.applyList(n, "Body")
357 case *ast.SwitchStmt:
358 a.apply(n, "Init", nil, n.Init)
359 a.apply(n, "Tag", nil, n.Tag)
360 a.apply(n, "Body", nil, n.Body)
362 case *ast.TypeSwitchStmt:
363 a.apply(n, "Init", nil, n.Init)
364 a.apply(n, "Assign", nil, n.Assign)
365 a.apply(n, "Body", nil, n.Body)
367 case *ast.CommClause:
368 a.apply(n, "Comm", nil, n.Comm)
369 a.applyList(n, "Body")
371 case *ast.SelectStmt:
372 a.apply(n, "Body", nil, n.Body)
375 a.apply(n, "Init", nil, n.Init)
376 a.apply(n, "Cond", nil, n.Cond)
377 a.apply(n, "Post", nil, n.Post)
378 a.apply(n, "Body", nil, n.Body)
381 a.apply(n, "Key", nil, n.Key)
382 a.apply(n, "Value", nil, n.Value)
383 a.apply(n, "X", nil, n.X)
384 a.apply(n, "Body", nil, n.Body)
387 case *ast.ImportSpec:
388 a.apply(n, "Doc", nil, n.Doc)
389 a.apply(n, "Name", nil, n.Name)
390 a.apply(n, "Path", nil, n.Path)
391 a.apply(n, "Comment", nil, n.Comment)
394 a.apply(n, "Doc", nil, n.Doc)
395 a.applyList(n, "Names")
396 a.apply(n, "Type", nil, n.Type)
397 a.applyList(n, "Values")
398 a.apply(n, "Comment", nil, n.Comment)
401 a.apply(n, "Doc", nil, n.Doc)
402 a.apply(n, "Name", nil, n.Name)
403 a.apply(n, "Type", nil, n.Type)
404 a.apply(n, "Comment", nil, n.Comment)
410 a.apply(n, "Doc", nil, n.Doc)
411 a.applyList(n, "Specs")
414 a.apply(n, "Doc", nil, n.Doc)
415 a.apply(n, "Recv", nil, n.Recv)
416 a.apply(n, "Name", nil, n.Name)
417 a.apply(n, "Type", nil, n.Type)
418 a.apply(n, "Body", nil, n.Body)
420 // Files and packages
422 a.apply(n, "Doc", nil, n.Doc)
423 a.apply(n, "Name", nil, n.Name)
424 a.applyList(n, "Decls")
425 // Don't walk n.Comments; they have either been walked already if
426 // they are Doc comments, or they can be easily walked explicitly.
429 // collect and sort names for reproducible behavior
431 for name := range n.Files {
432 names = append(names, name)
435 for _, name := range names {
436 a.apply(n, name, nil, n.Files[name])
440 panic(fmt.Sprintf("Apply: unexpected node type %T", n))
443 if a.post != nil && !a.post(&a.cursor) {
450 // An iterator controls iteration over a slice of nodes.
451 type iterator struct {
455 func (a *application) applyList(parent ast.Node, name string) {
456 // avoid heap-allocating a new iterator for each applyList call; reuse a.iter instead
460 // must reload parent.name each time, since cursor modifications might change it
461 v := reflect.Indirect(reflect.ValueOf(parent)).FieldByName(name)
462 if a.iter.index >= v.Len() {
466 // element x may be nil in a bad AST - be cautious
468 if e := v.Index(a.iter.index); e.IsValid() {
469 x = e.Interface().(ast.Node)
473 a.apply(parent, name, &a.iter, x)
474 a.iter.index += a.iter.step