inspector.go

     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  	// This function is equivalent to the PreorderSeq call below,
    77  	// but to avoid the additional dynamic call (which adds 13-35%
    78  	// to the benchmarks), we expand it out.
    79  	//
    80  	// in.PreorderSeq(types...)(func(n ast.Node) bool {
    81  	// 	f(n)
    82  	// 	return true
    83  	// })
    84  
    85  	mask := maskOf(types)
    86  	for i := 0; i < len(in.events); {
    87  		ev := in.events[i]
    88  		if ev.index > i {
    89  			// push
    90  			if ev.typ&mask != 0 {
    91  				f(ev.node)
    92  			}
    93  			pop := ev.index
    94  			if in.events[pop].typ&mask == 0 {
    95  				// Subtrees do not contain types: skip them and pop.
    96  				i = pop + 1
    97  				continue
    98  			}
    99  		}
   100  		i++
   101  	}
   102  }
   103  
   104  // Nodes visits the nodes of the files supplied to New in depth-first
   105  // order. It calls f(n, true) for each node n before it visits n's
   106  // children. If f returns true, Nodes invokes f recursively for each
   107  // of the non-nil children of the node, followed by a call of
   108  // f(n, false).
   109  //
   110  // The complete traversal sequence is determined by ast.Inspect.
   111  // The types argument, if non-empty, enables type-based filtering of
   112  // events. The function f if is called only for nodes whose type
   113  // matches an element of the types slice.
   114  func (in *Inspector) Nodes(types []ast.Node, f func(n ast.Node, push bool) (proceed bool)) {
   115  	mask := maskOf(types)
   116  	for i := 0; i < len(in.events); {
   117  		ev := in.events[i]
   118  		if ev.index > i {
   119  			// push
   120  			pop := ev.index
   121  			if ev.typ&mask != 0 {
   122  				if !f(ev.node, true) {
   123  					i = pop + 1 // jump to corresponding pop + 1
   124  					continue
   125  				}
   126  			}
   127  			if in.events[pop].typ&mask == 0 {
   128  				// Subtrees do not contain types: skip them.
   129  				i = pop
   130  				continue
   131  			}
   132  		} else {
   133  			// pop
   134  			push := ev.index
   135  			if in.events[push].typ&mask != 0 {
   136  				f(ev.node, false)
   137  			}
   138  		}
   139  		i++
   140  	}
   141  }
   142  
   143  // WithStack visits nodes in a similar manner to Nodes, but it
   144  // supplies each call to f an additional argument, the current
   145  // traversal stack. The stack's first element is the outermost node,
   146  // an *ast.File; its last is the innermost, n.
   147  func (in *Inspector) WithStack(types []ast.Node, f func(n ast.Node, push bool, stack []ast.Node) (proceed bool)) {
   148  	mask := maskOf(types)
   149  	var stack []ast.Node
   150  	for i := 0; i < len(in.events); {
   151  		ev := in.events[i]
   152  		if ev.index > i {
   153  			// push
   154  			pop := ev.index
   155  			stack = append(stack, ev.node)
   156  			if ev.typ&mask != 0 {
   157  				if !f(ev.node, true, stack) {
   158  					i = pop + 1
   159  					stack = stack[:len(stack)-1]
   160  					continue
   161  				}
   162  			}
   163  			if in.events[pop].typ&mask == 0 {
   164  				// Subtrees does not contain types: skip them.
   165  				i = pop
   166  				continue
   167  			}
   168  		} else {
   169  			// pop
   170  			push := ev.index
   171  			if in.events[push].typ&mask != 0 {
   172  				f(ev.node, false, stack)
   173  			}
   174  			stack = stack[:len(stack)-1]
   175  		}
   176  		i++
   177  	}
   178  }
   179  
   180  // traverse builds the table of events representing a traversal.
   181  func traverse(files []*ast.File) []event {
   182  	// Preallocate approximate number of events
   183  	// based on source file extent of the declarations.
   184  	// (We use End-Pos not FileStart-FileEnd to neglect
   185  	// the effect of long doc comments.)
   186  	// This makes traverse faster by 4x (!).
   187  	var extent int
   188  	for _, f := range files {
   189  		extent += int(f.End() - f.Pos())
   190  	}
   191  	// This estimate is based on the net/http package.
   192  	capacity := extent * 33 / 100
   193  	if capacity > 1e6 {
   194  		capacity = 1e6 // impose some reasonable maximum
   195  	}
   196  	events := make([]event, 0, capacity)
   197  
   198  	var stack []event
   199  	stack = append(stack, event{}) // include an extra event so file nodes have a parent
   200  	for _, f := range files {
   201  		ast.Inspect(f, func(n ast.Node) bool {
   202  			if n != nil {
   203  				// push
   204  				ev := event{
   205  					node:  n,
   206  					typ:   0,           // temporarily used to accumulate type bits of subtree
   207  					index: len(events), // push event temporarily holds own index
   208  				}
   209  				stack = append(stack, ev)
   210  				events = append(events, ev)
   211  			} else {
   212  				// pop
   213  				top := len(stack) - 1
   214  				ev := stack[top]
   215  				typ := typeOf(ev.node)
   216  				push := ev.index
   217  				parent := top - 1
   218  
   219  				events[push].typ = typ            // set type of push
   220  				stack[parent].typ |= typ | ev.typ // parent's typ contains push and pop's typs.
   221  				events[push].index = len(events)  // make push refer to pop
   222  
   223  				stack = stack[:top]
   224  				events = append(events, ev)
   225  			}
   226  			return true
   227  		})
   228  	}
   229  
   230  	return events
   231  }
   232
View as plain text