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  	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
View as plain text